MODULE PGMSignal EXPORTS Main;
(*
  This program removes the low-frequency multiplicative component
  of a grayscale image, and outputs the resulting signal,
  optionally normalized with respect to (a) the quantization noise,
  (b) the local contrast, or (c) a specified reference image.
  
  The input image is interpreted as "exp(M + S)" where "exp(S)" is the
  desired high-frequency signal, and "exp(M)" is an undesired
  low-frequency modulation.
  
  The signal "S" is therefore obtained by taking the logarithm "L" of
  the input image, and subtracting the estimated "M".  The component
  "M" is either the logarithm of the reference image, or is obtained
  internally by convolving "L" with a user-specified weight matrix "W".
  
  The program ignores any pixels whose value is below a user-specified
  threshold, or where the reference image (if given) is zero.
  
  The resulting signal "S" is optionally normalized, in two
  possible ways:
  
    If "-normalizeNoise" is requested, the signal "S" is scaled by the
    factor "(1/maxVal)*(1/sqrt(US))", where "US" is the variance of the
    quantization-related uncertainty in "S".  Thus, in the normalized
    signal, the variance of the quantization-related noise is approximately
    "1/(maxVal)" everywhere.

    If "-normalizeContrast" is requested, the signal "S" is scaled by
    "1/sqrt(8*VS)", where "VS" is the total variance of "L" (including
    the quantization error) in the local neighborhood.  Thus the
    scaled signal will have variance (roughly) 1/8.
    
  In any case, the resulting signal "S" gets multiplied by a
  user-specified "sharpening" factor "F", then compressed to the range
  [0_1] by the "squeeze" function, and finally converted to an integer
  "c" by the formula "c := 1 + ROUND(y * (maxVal-1))".  Thus, pixels
  with "S = 0" (meaning the input image is neither darker nor lighter
  than the background) are mapped to approximately "maxVal/2".
  
  Also, if the signal "S" is not normalized, pixels that are
  "1+eps" times as bright as the background (for small eps), 
  will be mapped to approximately "(1+F*eps)*(maxval/2)".

  Pixels where the signal could not be computed are mapped to 0.
  
  TO DO: Should use the error function "erf" for squeezing/stretching.
  
*)

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

TYPE
  Weight = PixelWeightTable.T;
  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 *)
      weight: REF Weight;      (* Weight distribution type *)
      printWeights: BOOL;      (* TRUE to print pixel weight table *)
      printPixelValues: BOOL;  (* TRUE to print pixel avg/var tables *)
      maskRelative: BOOL;      (* TRUE to use the mask as the reference level *)
      normalizeNoise: BOOL;    (* Normalize so that noise = "1/maxVal" *)
      normalizeContrast: BOOL; (* Normalize so that local contrast = 1. *) 
      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 *)
      inC, mkC: NAT;
  BEGIN
    WITH 
      o = GetOptions(),
      w = o.weight^,
      NW = NUMBER(w),
      HW = NW DIV 2,
      inPGM = NEW(PGM.Reader).init(o.inRd, NX, NY, inMax),
      otPGM = NEW(PGM.Writer).init(stdout, NX, NY, inMax),
      quantRange = FLOAT(inMax-1, LONG),
      quantStep = 1.0d0/quantRange,
      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
      <* ASSERT NW MOD 2 = 1 *>
      <* ASSERT NUMBER(w[0]) = NW *>
      w[HW,HW] := 0.0d0;
      IF o.printWeights THEN 
        PixelWeightTable.Print(stderr, w)
      END;

      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:: *)
        inAvgLog = ComputePixelAvgLogTable(inMax)^,
        inVarLog = ComputePixelVarLogTable(inMax)^,

        (* de-quantization tables for mask image: *)
        mkAvgLog = ComputePixelAvgLogTable(mkMax)^,
        mkVarLog = ComputePixelVarLogTable(mkMax)^
      DO

        IF o.printPixelValues THEN 
          PrintPixelLogTables(stderr, "input image",    inAvgLog, inVarLog);
          PrintPixelLogTables(stderr, "mask/ref image", mkAvgLog, mkVarLog);
        END;

        PROCEDURE ComputeLocalSums(
            xp, yp: NAT;       (* Pixel coordinates *)
            VAR sumW: LONG;    (* "sum(w[r])" *)                      
            VAR sumWL: LONG;   (* "sum(w[r]*L[p-r])" *)               
            VAR sumWL2: LONG;  (* "sum(w[r]*(L[p-r]^2 + UL[p-r]))" *) 
            VAR sumW2UL: LONG; (* "sum(w[r]^2*UL[p-r])" *)            
          ) =
          (*
            Computes various sums over the unmasked pixels 
            in the neighborhood of pixel "[yp,xp]"
          *)
          BEGIN
            sumW := 0.0d0;     
            sumWL := 0.0d0;    
            sumWL2 := 0.0d0;   
            sumW2UL := 0.0d0;  
            FOR yw := -MIN(yp, HW) TO MIN(HW, NY-1-yp) DO
              WITH 
                yqm = (yp + yw) MOD NW, 
                lyq = imag[yqm], 
                myq = mask[yqm],
                wyw = w[yw+HW]
              DO
                FOR xw := -MIN(xp, HW) TO MIN(HW, NX-1-xp) DO
                  WITH xq = xp + xw DO
                    IF myq[xq] > 0 THEN 
                      WITH
                        q = lyq[xq],
                        L = inAvgLog[q],
                        UL = inVarLog[q],
                        W = wyw[xw+HW]
                      DO
                        sumW := sumW + W;
                        sumWL := sumWL + W * L;
                        sumWL2 := sumWL2 + W * (L * L + UL);
                        sumW2UL := sumW2UL + W * W * UL
                      END
                    END
                  END
                END
              END
            END
          END ComputeLocalSums;

        PROCEDURE ComputeOutputPixel(
            xp, yp: NAT;  (* Pixel coordinates *)
            p: NAT;       (* Pixel value *)
            m: NAT;       (* Mask value *)
          ): NAT =
          BEGIN
            IF m = 0 THEN
              RETURN 0
            ELSE
              VAR sumW, sumWL, sumWL2, sumW2UL: LONG;
              VAR M: LONG;  (* Log of the unwanted modulation *)
                  UM: LONG; (* Variance of quantization noise in "M" *)
                  R: LONG;  (* Normalized signal *)
              BEGIN
                ComputeLocalSums(xp, yp, sumW, sumWL, sumWL2, sumW2UL);
                IF sumW = 0.0d0 THEN
                  RETURN 0
                ELSE
                  IF o.maskRelative THEN
                    M := mkAvgLog[m]; 
                    UM := mkVarLog[m];
                  ELSE 
                    M := sumWL/sumW; 
                    UM := sumW2UL/(sumW * sumW)
                  END;
                  WITH
                    L = inAvgLog[p],
                    S = L - M,
                    US = UM + inVarLog[p]
                  DO
                    IF o.normalizeNoise THEN
                      R := S*(quantStep/Math.sqrt(US))
                    ELSIF o.normalizeContrast THEN
                      WITH
                        ML = (sumWL + L)/(sumW + 1.0d0),
                        VL = (sumWL2 + L*L)/(sumW + 1.0d0),
                        VS = MAX(1.0d-10, VL  - M*(2.0d0*ML - M))
                         (* 
                           "VS" is Avg((log(image)-M)^2) in neighborhood, 
                           including the current pixel *)
                      DO
                        R := S/Math.sqrt(8.0d0*VS)
                      END
                    ELSE
                      R := S
                    END;
                  END
                END;
                RETURN 1 + ROUND(quantRange * Squeeze(o.sharpen * R))
              END
            END
          END ComputeOutputPixel;

        BEGIN
          (* Read first "HW" scanlines *)
          FOR yp := 0 TO MIN(HW, NY)-1 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[yp,xp] := inC; mask[yp,xp] := mkC
            END
          END;

          (* Process all scanlines: *)
          FOR yp := 0 TO NY-1 DO

            IF (yp + HW) < NY THEN
              (* Read line "yp+HW" into its proper place: *)
              WITH yqm = (yp + HW) MOD NW DO
                FOR xq := 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[yqm,xq] := inC; mask[yqm,xq] := mkC
                END
              END
            END;

            FOR xp := 0 TO NX-1 DO
              WITH ypm = yp MOD NW DO
                otPGM.put(ComputeOutputPixel(xp, yp, imag[ypm,xp], mask[ypm,xp]))
              END
            END
          END;
        END;
      END;
      otPGM.finish()
    END
  END Main;

(*
  The following pixel translation tables assume that the original
  pixel intensities were uniformly distributed in "[0_1]", and 
  were quantized to "[0..inMax]" by the formula "c = ROUND(p*inMax)".
*)

PROCEDURE ComputePixelAvgLogTable(inMax: NAT): REF ARRAY OF LONG =
  (*
    The returned table maps each pixel code to the corresponding
    mean log(intensity).
    See the algebra in "~stolfi/maple/PixelLogIntensityTables.ms". 
  *)
  VAR a, b, ya, yb: LONG;
  BEGIN
    WITH
      rt = NEW(REF ARRAY OF LONG, inMax+1),
      t = rt^,
      mv = FLOAT(inMax, LONG)
    DO
      b := 0.5d0/mv;
      t[0] := Math.log(b) - 1.0d0;
      yb := b*Math.log(b);
      FOR v := 1 TO inMax-1 DO 
        a := b;
        ya := yb;
        b := (FLOAT(v, LONG) + 0.5d0)/mv;
        yb := b*Math.log(b);
        t[v] := mv * (yb - ya) - 1.0d0;
      END;
      a := b;
      ya := yb;
      t[inMax] := - 2.0d0*mv * ya - 1.0d0;
      RETURN rt
    END
  END ComputePixelAvgLogTable;
  
PROCEDURE ComputePixelVarLogTable(inMax: NAT): REF ARRAY OF LONG =
  (*
    The returned table maps each pixel code to the corresponding
    variance of the quantization error in log(intensity).
    See the algebra in "~stolfi/maple/PixelLogIntensityTables.ms". 
  *)
  VAR a, b, lna, lnb: LONG;
  BEGIN
    WITH
      rt = NEW(REF ARRAY OF LONG, inMax+1),
      t = rt^,
      mv = FLOAT(inMax, LONG)
    DO
      b := 0.5d0/mv;
      lnb := Math.log(b);
      t[0] := 1.0d0;
      FOR v := 1 TO inMax-1 DO 
        a := b;
        lna := lnb;
        b := (FLOAT(v, LONG) + 0.5d0)/mv;
        lnb := Math.log(b);
        t[v] := 1.0d0 - mv*mv*a*b*(lnb - lna)*(lnb - lna);
      END;
      a := b;
      lna := lnb;
      t[inMax] := 1.0d0 - 4.0d0*mv*mv*a*lna*lna;
      RETURN rt
    END
  END ComputePixelVarLogTable;

PROCEDURE PrintPixelLogTables(wr: Wr.T; tit: TEXT; READONLY ta, tv: ARRAY OF LONG) =
  <* FATAL Wr.Failure, Thread.Alerted *>
  BEGIN
    WITH
      inMax = LAST(ta),
      ndigs = Text.Length(Fmt.Int(inMax))
    DO
      Wr.PutText(wr, "pixel tables for ");
      Wr.PutText(wr, tit);
      Wr.PutText(wr, "\n");
      Wr.PutText(wr, Fmt.Pad("x", ndigs));
      Wr.PutChar(wr, ' ');
      Wr.PutText(wr, Fmt.Pad("Avg(log(x))", 8));
      Wr.PutChar(wr, ' ');
      Wr.PutText(wr, Fmt.Pad("Var(log(x))", 12));
      Wr.PutText(wr, "\n");
      Wr.PutText(wr, Fmt.Pad("", ndigs, '-'));
      Wr.PutChar(wr, ' ');
      Wr.PutText(wr, Fmt.Pad("", 8, '-'));
      Wr.PutChar(wr, ' ');
      Wr.PutText(wr, Fmt.Pad("", 12, '-'));
      Wr.PutText(wr, "\n");
      
      FOR v := 0 TO LAST(ta) DO 
        Wr.PutText(wr, Fmt.Pad(Fmt.Int(v), ndigs, '0'));
        Wr.PutChar(wr, ' ');
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(ta[v], Fmt.Style.Fix, 4), 8));
        Wr.PutChar(wr, ' ');
        Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(tv[v], Fmt.Style.Fix, 8), 12));
        Wr.PutText(wr, "\n")
      END;
    END
  END PrintPixelLogTables;

<*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
        pp.getKeyword("-weight");
        o.weight := PixelWeightTable.FromParams(pp);
        
        o.printWeights := pp.keywordPresent("-printWeights");
        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.normalizeNoise := pp.keywordPresent("-normalizeNoise");
        o.normalizeContrast := pp.keywordPresent("-normalizeContrast");
        IF o.normalizeNoise AND o.normalizeContrast THEN
          pp.error("only one of \"-normalizeNoise\" or \"normalizeContrast\" allowed")
        END;
        
        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: PGMSignal \\\n");
          Wr.PutText(stderr, "  -weight \\\n");
          Wr.PutText(stderr, PixelWeightTable.ParamsHelp & " \\\n");
          Wr.PutText(stderr, "  [ -printWeights ] [ -printPixelValues ] \\\n");
          Wr.PutText(stderr, "  [ -minVal <N> ] \\\n");
          Wr.PutText(stderr, "  [ -mask <pgmfile> | -reference <pgmfile> ] \\\n");
          Wr.PutText(stderr, "  [ -normalizeNoise | -normalizeContrast ] \\\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 PGMSignal.