MODULE PGMNormSignal EXPORTS Main;
(*
  This program removes the low-frequency multiplicative component
  of an image and normalizes the local contrast OF what is left.
  
  The input is assumed to be a a multiplicative image (e.g. a photograph)
  with linear energy encoding.  The output is a signed additive 
  image with "squeeze" encoding.
  
  Let "X" be log(intensity) for a given pixel "p", and "M" and "V" be
  the weighted mean and variance of log(intensity) in the neighborhood
  of a pixel "p".  The value "exp(M)" is a low-frequency
  multiplicative component of the image, so "exp(X-M)" would be the
  high-frequency component.
  
  The output is "(X-M)/sqrt(V)", squeezed to "[0_1]".  Note that
  this formula roughly normalizes the local variance to 1.
*)  

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

CONST
  Tiny = LongReal.MinPos;
  Huge = LongReal.MaxFinite;
  MaxWidth = 1024;

TYPE
  Distr = OBJECT METHODS eval(x, y, r2: LONGREAL): LONGREAL END;

TYPE
  Options = RECORD
      input: Rd.T;           (* The input image *)
      distr: Distr;          (* Weight distribution type *)
      width: CARDINAL;       (* Window width (an odd INTEGER) *)
      printWeights: BOOLEAN; (* TRUE to print pixel weight table *)
    END;
    
PROCEDURE Main() =
  <* FATAL Wr.Failure, Thread.Alerted, Rd.Failure *>
  VAR NX, NY, maxVal, c: CARDINAL;
  BEGIN
    WITH 
      o = GetOptions(),
      NW = o.width,
      HW = NW DIV 2,
      rd = NEW(PGM.Reader).init(o.input, NX, NY, maxVal),
      wr = NEW(PGM.Writer).init(stdout, NX, NY, maxVal),
      (* input table: "inStretch[c] = Stretch(c/maxVal)" *)
      inStretch = ComputeInStretchTable(maxVal)^,
      (* input quantization variance table: *)
      inStretchVar = ComputeInStretchVarTable(maxVal)^,
      outScale = FLOAT(maxVal, LONGREAL),
      stretchBrite = Stretch(o.brightness),
      w = ComputeWeightTable(NW, o.distr)^,
      line = NEW(REF ARRAY OF ARRAY OF CARDINAL, NW, NX)^
    DO
      IF o.printWeights THEN PrintWeightTable(stderr, w) END;
      
      (* Read first "width/2" scanlines *)
      FOR yp := 0 TO MIN(HW, NY)-1 DO
        FOR xp := 0 TO NX-1 DO
          rd.get(c); 
          line[yp,xp] := c
        END
      END;
      
      (* Perform filtering *)
      FOR yp := 0 TO NY-1 DO
        (* Invariant: scanline "k" OF input image is stored in "line[k MOD NW]", *)
        (*   for "k" in "MAX(0,yp-HW) .. MIN(NY,yp+HW)-1". *)
        
        WITH ypm = yp MOD NW 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
                rd.get(c); 
                line[yqm,xq] := c
              END
            END
          END;

          (* Process pixels: *)
          FOR xp := 0 TO NX-1 DO
            VAR sumW, sumWX, sumWX2: LONGREAL := 0.0d0;
            BEGIN
              (* Compute weighted count, sum, and  sum of squares of pixels in window: *)
              FOR yw := -MIN(yp, HW) TO MIN(HW, NY-1-yp) DO
                WITH yqm = (yp + yw) MOD NW DO
                  FOR xw := -MIN(xp, HW) TO MIN(HW, NX-1-xp) DO
                    WITH 
                      xq = xp + xw, 
                      q = line[yqm, xq], 
                      sxq = inStretch[q],
                      svq = inStretchVar[q],
                      wq = w[yw+HW, xw+HW]
                    DO
                      sumW := sumW + wq;
                      sumWX := sumWX + wq * sxq;
                      sumWX2 := sumWX2 + wq * (sxq * sxq + svq);
                    END
                  END
                END
              END;

              (* Compute local mean "mu" and deviation "sd" around "p"  *)
              (* Assumes the input pixels are always contaminated with *)
              (*   a quantization error of "quantVar = (1/maxVal)^2/12". *)
              WITH
                mu = sumWX/sumW,
                var = sumWX2/sumW - mu*mu,
                sd = Math.sqrt(var),
                P = inStretch[line[ypm, xp]],
                R = Adjust(P, 
                  M := mu, D := sd, 
                  c := o.contrast, r := o.reference, B := stretchBrite
                ),
                r = Squeeze(R)
              DO
                wr.put(ROUND(outScale * r))
                (* wr.put(ROUND(outScale * mu)) *)
              END
            END
          END
        END
      END
    END
  END Main;

TYPE
  WeightTable = ARRAY OF ARRAY OF LONGREAL;

PROCEDURE ComputeWeightTable(
    NW: CARDINAL; 
    distr: Distr; 
  ): REF WeightTable =
  VAR e: LONGREAL;
  BEGIN
    WITH
      HW = NW DIV 2,
      r2max = FLOAT(HW*HW, LONGREAL),
      wRef = NEW(REF WeightTable, NW, NW),
      w = wRef^
    DO
      <* ASSERT NW MOD 2 = 1 *>
      FOR yw := 0 TO HW DO
        FOR xw := yw TO HW DO
          WITH 
            x = FLOAT(xw, LONGREAL),
            y = FLOAT(yw, LONGREAL),
            r2 = x*x + y*y
          DO
            IF r2 > r2max THEN
              e := 0.0d0
            ELSE
              e := distr.eval(x, y, r2)
            END;
            w[HW+yw, HW+xw] := e;
            w[HW-yw, HW+xw] := e;
            w[HW+yw, HW-xw] := e;
            w[HW-yw, HW-xw] := e;
            
            w[HW+xw, HW+yw] := e;
            w[HW+xw, HW-yw] := e;
            w[HW-xw, HW+yw] := e;
            w[HW-xw, HW-yw] := e;
          END
        END
      END;
      w[HW, HW] := 0.0d0;
      (* Normalize weights: *)
      VAR sum: LONGREAL := 0.0d0;
      BEGIN
        FOR yw := 0 TO NW-1 DO
          FOR xw := 0 TO NW-1 DO
            WITH e = w[yw,xw] DO sum := sum + e END
          END
        END;
        FOR yw := 0 TO NW-1 DO
          FOR xw := 0 TO NW-1 DO
            WITH e = w[yw,xw] DO e := e/sum END
          END
        END;
      END;
      RETURN wRef
    END
  END ComputeWeightTable;

PROCEDURE PrintWeightTable(wr: Wr.T; READONLY w: WeightTable) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    WITH NW = NUMBER(w) DO
      FOR yw := 0 TO NW-1 DO
        FOR xw := 0 TO NW-1 DO
          WITH e = w[yw,xw] DO 
            Wr.PutText(wr, Fmt.Pad(Fmt.LongReal(e, Fmt.Style.Fix, 5), 8))
          END
        END;
        Wr.PutText(wr, "\n")
      END
    END
  END PrintWeightTable;
  
PROCEDURE ComputeInStretchTable(maxVal: CARDINAL): REF ARRAY OF LONGREAL =
  BEGIN
    WITH
      rt = NEW(REF ARRAY OF LONGREAL, maxVal+1),
      t = rt^,
      mv = FLOAT(maxVal+1, LONGREAL)
    DO
      t[0] := Stretch(0.25d0/mv);
      FOR v := 1 TO maxVal-1 DO 
        t[v] := Stretch((FLOAT(v, LONGREAL) + 0.5d0)/mv)
      END;
      t[maxVal] := Stretch(1.0d0 - 0.25d0/mv);
      RETURN rt
    END
  END ComputeInStretchTable;
  
PROCEDURE ComputeInStretchVarTable(maxVal: CARDINAL): REF ARRAY OF LONGREAL =
  BEGIN
    WITH
      rt = NEW(REF ARRAY OF LONGREAL, maxVal+1),
      t = rt^,
      mv = FLOAT(maxVal+1, LONGREAL)
    DO
      WITH
        x = 0.25d0/mv,
        rx = 0.25d0/mv,
        vx = rx*rx/3.0d0,
        dstretch = DStretch(x),
        vstretch = dstretch*dstretch * vx
      DO
        <* ASSERT vstretch > 0.0d0 *>
        t[0] := vstretch;
        t[maxVal] := vstretch
      END;
      FOR v := 1 TO maxVal-1 DO
        WITH
          x = FLOAT(v, LONGREAL)/mv,
          rx = 0.50d0/mv,
          vx = rx*rx/3.0d0,
          dstretch = DStretch(x),
          vstretch = dstretch*dstretch * vx
        DO
          <* ASSERT vstretch > 0.0d0 *>
          t[v] := vstretch
        END;
      END;
      RETURN rt
    END
  END ComputeInStretchVarTable;

TYPE GaussDistr = (* Gaussian distribution: "w(r) = exp(-r^2*r2scale)" *)
  Distr OBJECT r2Scale: LONGREAL OVERRIDES eval := GaussEval END;
  
PROCEDURE GaussEval(d: GaussDistr; <*UNUSED*> x, y: LONGREAL; r2: LONGREAL): LONGREAL =
  BEGIN
    RETURN Math.exp(- r2 * d.r2Scale)
  END GaussEval;
  
TYPE PowerDistr = (* Power distribution: "w(r) = 1/(r^2)^alpha" *)
  Distr OBJECT alpha: LONGREAL OVERRIDES eval := PowerEval END;
  
PROCEDURE PowerEval(d: PowerDistr; <*UNUSED*> x, y: LONGREAL; r2: LONGREAL): LONGREAL =
  BEGIN
    IF r2 = 0.0d0 THEN 
      RETURN 0.0d0
    ELSE
      RETURN Math.exp(-d.alpha*Math.log(r2))
    END
  END PowerEval;

TYPE FlatDistr = (* Flat distribution: "w(r) = 1" *)
  Distr OBJECT OVERRIDES eval := FlatEval END;
  
PROCEDURE FlatEval(<*UNUSED*> d: FlatDistr; <*UNUSED*> x, y, r2: LONGREAL): LONGREAL =
  BEGIN
    RETURN 1.0d0
  END FlatEval;
  
PROCEDURE Stretch(x: LONGREAL): LONGREAL = 
  (* Maps [0 _ 1] to [-oo _ +oo] *)
  BEGIN
    RETURN 0.25d0 * (Math.log(x) - Math.log(1.0d0 - x))
  END Stretch;
  
PROCEDURE DStretch(x: LONGREAL): LONGREAL = 
  (* Derivative of "Stretch" *)
  BEGIN
    RETURN 0.25d0/(x * (1.0d0 - x))
  END DStretch;
  
PROCEDURE Squeeze(xh: LONGREAL): LONGREAL =
  (* Functional inverse of "Stretch" *)
  BEGIN
    WITH e = Math.exp(4.0d0 * xh) DO 
      RETURN e/(e + 1.0d0)
    END
  END Squeeze;
  
PROCEDURE Adjust(
    X: LONGREAL;          (* Stretched pixel value *)
    M, D: LONGREAL;       (* Local mean and stddev of stretched pixels *)
    c: LONGREAL;          (* Output contrast *)
    r: LONGREAL;          (* Relative stretched reference brightness on input *)
    B: LONGREAL;          (* Stretched reference brightness on output *)
  ): LONGREAL = 
  (*
    Adjusts the stretched pixel value "X" as described above.
    Returns the stretched output pixel value.
  *)
  CONST D0 = 0.453449841058555d0; (* Deviation of stretched white noise *)
  BEGIN
    <* ASSERT D > 0.0d0 *>
    WITH
      (* Compute stretch factor "alpha": *)
      alpha = c * D0 / D,
      (* Compute the shift "beta": *)
      beta = B - alpha*(M + r * D),
      (* Compute output value and unstretch: *)
      Y = alpha * X + beta
    DO
      RETURN Y
    END
  END Adjust;
  
PROCEDURE GetOptions (): Options =
  <* FATAL Thread.Alerted, Wr.Failure *>
  VAR o: Options;
  BEGIN
    WITH
      pp = NEW(ParseParams.T).init(stderr)
    DO

      TRY
        pp.getKeyword("-width");
        o.width := pp.getNextInt(3, MaxWidth);
        IF o.width MOD 2 # 1 THEN pp.error("\"-width\" must be odd") END;

        IF pp.keywordPresent("-flat") THEN
          o.distr := NEW(FlatDistr)
        ELSIF pp.keywordPresent("-power") THEN
          WITH expt = pp.getNextLongReal(0.0d0, 10.0d0) DO
            o.distr := NEW(PowerDistr, alpha := expt/2.0d0)
          END
        ELSIF pp.keywordPresent("-gauss") THEN
          WITH radius = pp.getNextLongReal(Tiny, Huge) DO
            o.distr := NEW(GaussDistr, r2Scale := 1.0d0/(radius*radius))
          END
        ELSE
          pp.error("distribution not specified")
        END;
        
        o.printWeights := pp.keywordPresent("-printWeights");
        
        IF pp.keywordPresent("-contrast") THEN
          o.contrast := pp.getNextLongReal()
        ELSE
          o.contrast := 1.0d0
        END;
          
        IF pp.keywordPresent("-reference") THEN
          o.reference := pp.getNextLongReal()
        ELSE
          o.reference := 0.0d0
        END;
        
        IF pp.keywordPresent("-brightness") THEN
          o.brightness := pp.getNextLongReal(0.00001d0, 0.99999d0)
        ELSE
          o.brightness := 0.5d0
        END;
        
        pp.skipParsed();
        IF pp.next < NUMBER(pp.arg^) THEN 
          WITH name = pp.getNext() DO
            IF Text.Equal(name, "-") THEN
              o.input := stdin
            ELSE
              TRY
                o.input := FileRd.Open(name)
              EXCEPT
              | OSError.E => pp.error("can't open file \"" & name & "\"")
              END
            END
          END
        ELSE
          o.input := stdin
        END;

        pp.finish();                                       
      EXCEPT                                                            
      | ParseParams.Error =>                                              
          Wr.PutText(stderr, "Usage: PGMNormSignal \\\n");
          Wr.PutText(stderr, "  -width <wd> \\\n");
          Wr.PutText(stderr, "  { -flat | -gauss <radius> | -power <expt> } \\\n");
          Wr.PutText(stderr, "  [ -printWeights ] \\\n");
          Wr.PutText(stderr, "  [ -brightness <num> ] \\\n");
          Wr.PutText(stderr, "  [ -contrast <num> ] \\\n");
          Wr.PutText(stderr, "  [ <pgmfile> ]\n");
          Process.Exit(1);
      END;
    END;
    RETURN o
  END GetOptions;
  
BEGIN 
  Main();
END PGMNormSignal.