MODULE TextureInvariants2D EXPORTS Main;
(*
  This program computes three "shape index" images
  from a given grayscale image "f".  These shape indices are
  non-linear differential operators defined as follows:

    lambda = L*M/M^2;   # Bumpiness

    chi = R^2/M^2;      # Saddlism squared

    rho = L*R/M^2;      # Ridgicity

  where

    M = sqrt(L^2+S^2+T^2)
    R = sqrt(S^2+T^2)

    L = fxx + fyy;  # Laplacian
    T = 2*fxy;      # Twist
    S = fxx - fyy;  # Sault

  Let "f*" be the deviation from "f" to its first-order taylor
  approximation. The parameter "lambda", "chi" and "rho" describe the
  shape of the level curves of "f*" around the point "(x,y)".  The
  values are as follows:
  
       lambda chi    rho
       ------ ------ ------
       +1     0      0      circular curves around a minimum,
       +H     +1/2   +1/2   parallel lines along a valley
       0      +1     0      saddle-shaped curves (or "f*" is null)
       -H     +1/2   -1/2   parallel lines along a ridge
       -1     0      0      circular curves around a maximum
       
  The program optionally reads a mask or reference image "mk".  Pixels
  where "mk" is zero are assumed to be invalid.  In addition, when
  "mk" is given, the input image may be converted to a relative signal
  "s", by the formula "s = log(f) - log(mk)", before applying the 
  operators.

*)

IMPORT Texture2DGray AS Tx;
IMPORT PGM, Math, PixelWeightTable;
IMPORT Text, Rd, Fmt, FileRd, Wr, Thread, OSError, Process, ParseParams;
FROM Stdio IMPORT stdin, stdout, stderr;

TYPE
  LONG = LONGREAL;
  NAT = CARDINAL;
  BOOL = BOOLEAN;

TYPE
  Options = RECORD
      inRd: Rd.T;              (* The input image file *)
      mkRd: Rd.T;              (* The pixel mask/reference image file, or NIL *)
      minVal: NAT;             (* Ignore pixels lower than this *)
      maskRelative: BOOL;      (* TRUE to use the mask as the reference level *)
      logarithmic: BOOL;       (* TRUE to take logs of input and mask *)
      printPixelValues: BOOL;  (* TRUE to print log conversion tables *)
      sharpen: LONG;           (* Contrast-sharpening factor "F" *)
    END;
    
PROCEDURE Main() =
  <* FATAL Wr.Failure, Thread.Alerted, Rd.Failure *>
  VAR NX, NY: NAT;         (* Image size *)
      mkPGM: PGM.Reader;   (* Mask pixel reader *)
      inMax, mkMax: NAT;   (* Max pixel values of image and mask *)
  BEGIN
    WITH 
      NW = 3,  (* Number of scanlines needed for processing (odd) *)
      HW = 1,  (* NW DIV 2 *)

      o = GetOptions(),
      inPGM = NEW(PGM.Reader).init(o.inRd, NX, NY, inMax),
      otPGM = NEW(PGM.Writer).init(stdout, NX, NY, inMax),
      otRange = FLOAT(inMax - 1, LONG),
      otStep = 1.0d0/otRange,
      
      imag = NEW(REF ARRAY OF ARRAY OF NAT, NW, NX)^,
      mask = NEW(REF ARRAY OF ARRAY OF NAT, NW, NX)^
        (* 
          Invariant: while processing scanline "yp", line "k" of the
          input image is stored in "imag[k MOD NW]", for "k" in 
          "MAX(0,yp-HW) .. MIN(NY-1,yp+HW)"; and similarly for the mask. *)
    DO
      IF o.mkRd # NIL THEN 
        VAR  mkNX, mkNY: NAT;
        BEGIN
          mkPGM := NEW(PGM.Reader).init(o.mkRd, mkNX, mkNY, mkMax);
          <* ASSERT mkNX = NX AND mkNY = NY *>
        END
      ELSE
        mkPGM := NIL ; mkMax := 1
      END;
      
      WITH 
        (* de-quantization tables for input image: *)
        inPixVal = PixValTable(o.minVal, inMax, o.logarithmic, allZeros := FALSE)^,
        inLo = inPixVal[o.minVal],
        inHi = inPixVal[inMax],

        (* de-quantization tables for mask image: *)
        mkPixVal = PixValTable(1, mkMax, o.logarithmic, allZeros := NOT o.maskRelative)^,
        mkLo = mkPixVal[1],
        mkHi = mkPixVal[mkMax],

        (* Relative signal range: *)
        sgLo = imLo - mkHi,
        sgHi = imHi - mkLo
      DO
      
        IF o.printPixelValues THEN 
          PrintPixelLogTables(stderr, "input image",    inPixVal);
          PrintPixelLogTables(stderr, "mask/ref image", mkPixVal);
        END;

        PROCEDURE ComputeOutputPixel(
            xp, yp: NAT;          (* Pixel coordinates *)
            pmm, pzm, ppm: NAT;   (* Image pixel values in neighborhood *)
            pmz, pzz, ppz: NAT;   (* " " *)
            pmp, pzp, ppp: NAT;   (* " " *)
            mmm, mzm, mpm: NAT;   (* Mask pixel values in neighborhood *)
            mmz, mzz, mpz: NAT;   (* " " *)
            mmp, mzp, mpp: NAT;   (* " " *)
          ): NAT =
          BEGIN
            IF mmm = 0 OR mzm = 0 OR mpm = 0
            OR mmz = 0 OR mzz = 0 OR mpz = 0
            OR mmp = 0 OR mzp = 0 OR mpp = 0 
            THEN
              RETURN 0
            ELSE
              ???
              WITH
                L = inPixVal[p], (* Input image value *)
                M = mkPixVal[m], (* Reference signal *)
                R = L - M        (* Relative signal, in [sgLo _ sgHi] *)
              DO
                VAR S: LONG;  (* Output signal in [0 _ 1]: *)
                BEGIN
                  IF o.logarithmic THEN 
                    S := Squeeze(o.sharpen * R) 
                  ELSE
                    S := (R - sgLo)/(sgHi - sgLo);
                    IF o.sharpen # 1.0d0 THEN 
                      S := Squeeze(o.sharpen * Stretch(S));
                    END;
                  END;
                  RETURN 1 + ROUND(otRange * S)
                END
              END
            END
          END ComputeOutputPixel;

        PROCEDURE ReadScanLine(yp: CARDINAL) =
          VAR inC, mkC: NAT;
          BEGIN
            WITH ypm = yp MOD NW DO 
              FOR xp := 0 TO NX-1 DO
                inPGM.get(inC);
                IF mkPGM # NIL THEN mkPGM.get(mkC) ELSE mkC := mkMax END;
                IF inC < o.minVal THEN mkC := 0 END;
                imag[ypm,xp] := inC; mask[ypm,xp] := mkC
              END
            END;
          END ReadScanLine;

        BEGIN
          (* Read first HW scanlines: *)
          FOR yp := 0 TO MIN(HW-1, NY-1) DO ReadScanLine(yp) END;
          
          (* Process all scanlines, except for a frame of width "HW" *)
          FOR yp := 0 TO NY-1 DO

            IF yp + HW < NY THEN ReadScanLine(yp + HW) END;

            (* Process scanline: *)
            FOR xp := 0 TO NX-1 DO
              IF yp < HW OR yp > NY-1-HW OR xp < HW OR xp > NX-1-HW THEN
                otPGM.put(0)
              ELSE
                WITH
                  ym = (yp - 1) MOD 3, inm = imag[ym], mkm = mask[ym],
                  yz = yp MOD 3,       inz = imag[yz], mkm = mask[yz],
                  yp = (yp + 1) MOD 3, inp = imag[yp], mkm = mask[yp]
                DO
                  otPGM.put(ComputeOutputPixel(xp, yp, 
                    inm[xp-1], inm[xp], inm[xp+1],
                    inz[xp-1], inz[xp], inz[xp+1],
                    inp[xp-1], inp[xp], inp[xp+1],
                    
                    mkm[xp-1], mkm[xp], mkm[xp+1],
                    mkz[xp-1], mkz[xp], mkz[xp+1],
                    mkp[xp-1], mkp[xp], mkp[xp+1],
                  ))
                END
              END
            END
          END;
        END;
      END;
      otPGM.finish()
    END
  END Main;

(*
  The following pixel translation tables assume that 
*)

PROCEDURE PixValTable(
    minCode, maxCode: NAT; 
    logarithmic: BOOL; 
    allZeros: BOOL;
  ): REF ARRAY OF LONG =
  (*
    The returned table gives the expected intensity that was
    quantized to each pixel code, assuming that the original
    pixel intensities were uniformly distributed in "[0_1]", and 
    were quantized to "[minCode..maxCode]" by the formula 
    "c = minCode + ROUND(p*(maxCode-minCode))".
    
    If "logarithmic" is true, the returned table gives the expected
    log (natural) of the signal, rather than the expected signal.  The
    distribution is still assumed to be uniform in the signal itself.
    
    If "allZeros" is true, all entries are set to zero.  This option
    is useful when building the reference image table, when the latter
    does not exist or is to be used only as a boolean mask.
  *)
  VAR a, b, ya, yb: LONG;
  BEGIN
    WITH
      rt = NEW(REF ARRAY OF LONG, maxCode+1),
      t = rt^,
      quantRange = FLOAT(maxCode-minCode, LONG)
    DO
      <* ASSERT minCode <= maxCode *>
      FOR v := 0 TO minCode-1 DO t[v] := 0.0d0 END
      IF allZeros OR minCode = maxCode THEN
        FOR v := minCode TO maxCode DO t[v] := 0.0d0 END
      ELSIF logarithmic THEN
        (* See the algebra in "~stolfi/maple/PixelLogIntensityTables.ms". *)
        b := 0.5d0/quantRange;
        t[minCode] := Math.log(b) - 1.0d0;
        yb := b*Math.log(b);
        FOR v := minCode+1 TO maxCode-1 DO 
          a := b;
          ya := yb;
          b := (FLOAT(v, LONG) + 0.5d0)/quantRange;
          yb := b*Math.log(b);
          t[v] := quantRange * (yb - ya) - 1.0d0;
        END;
        a := b;
        ya := yb;
        t[maxCode] := - 2.0d0 * quantRange * ya - 1.0d0;
      ELSE
        t[minCode] := 0.25d0/quantRange;
        FOR v := 1 TO maxCode-1 DO 
          t[v] := FLOAT(v, LONG)/quantRange;
        END;
        t[maxCode] := 1.0d0 - 0.25d0/quantRange;
      END;
      RETURN rt
    END
  END PixValTable;
  
<*UNUSED*>
PROCEDURE Stretch(x: LONG): LONG = 
  (* Maps [0 _ 1] to [-oo _ +oo] *)
  BEGIN
    RETURN 0.25d0 * (Math.log(x) - Math.log(1.0d0 - x))
  END Stretch;
  
<*UNUSED*>
PROCEDURE DStretch(x: LONG): LONG = 
  (* Derivative of "Stretch" *)
  BEGIN
    RETURN 0.25d0/(x * (1.0d0 - x))
  END DStretch;
  
PROCEDURE Squeeze(xh: LONG): LONG =
  (* Functional inverse of "Stretch" *)
  BEGIN
    WITH e = Math.exp(4.0d0 * xh) DO 
      RETURN e/(e + 1.0d0)
    END
  END Squeeze;
  
PROCEDURE GetOptions (): Options =
  <* FATAL Thread.Alerted, Wr.Failure *>
  CONST
    MaxMaxVal: NAT = 256*256-1;
  VAR o: Options;
  BEGIN
    WITH
      pp = NEW(ParseParams.T).init(stderr)
    DO

      TRY
        o.printPixelValues := pp.keywordPresent("-printPixelValues");
        
        IF pp.keywordPresent("-minVal") THEN
          o.minVal := pp.getNextInt(0, MaxMaxVal);
        ELSE
          o.minVal := 0;
        END;
        
        IF pp.keywordPresent("-mask") THEN
          o.mkRd := GetFileOption(pp);
          o.maskRelative := FALSE
        ELSIF pp.keywordPresent("-reference") THEN
          o.mkRd := GetFileOption(pp);
          o.maskRelative := TRUE
        ELSE
          o.mkRd := NIL;
          o.maskRelative := FALSE
        END;
        
        o.logarithmic := pp.keywordPresent("-logarithmic");
        
        IF pp.keywordPresent("-sharpen") THEN
          o.sharpen := pp.getNextLongReal(0.001d0, 100.0d0);
        ELSE
          o.sharpen := 1.0d0
        END;
        
        pp.skipParsed();
        IF pp.next < NUMBER(pp.arg^) THEN 
          o.inRd := GetFileOption(pp)
        ELSE
          o.inRd := stdin
        END;

        pp.finish();                                       
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: TextureInvariants2D \\\n");
          Wr.PutText(stderr, PixelWeightTable.ParamsHelp & " \\\n");
          Wr.PutText(stderr, "  [ -printPixelValues ] \\\n");
          Wr.PutText(stderr, "  [ -minVal <N> ] \\\n");
          Wr.PutText(stderr, "  [ -mask <pgmfile> | -reference <pgmfile> ] \\\n");
          Wr.PutText(stderr, "  [ -logarithmic ] \\\n");
          Wr.PutText(stderr, "  [ -sharpen <F> ] \\\n");
          Wr.PutText(stderr, "  [ <pgmfile> ]\n");
          Process.Exit(1);
      END;
    END;
    RETURN o
  END GetOptions;
  
PROCEDURE GetFileOption(pp: ParseParams.T): Rd.T RAISES {ParseParams.Error} =
  BEGIN
    WITH name = pp.getNext() DO
      IF Text.Equal(name, "-") THEN
        RETURN stdin
      ELSE
        TRY
          RETURN FileRd.Open(name)
        EXCEPT
        | OSError.E => 
            pp.error("can't open file \"" & name & "\"");
            RETURN NIL
        END
      END
    END
  END GetFileOption;

BEGIN 
  Main();
END TextureInvariants2D.