MODULE TraceSphere EXPORTS Main;

IMPORT R3, VGM, Wr, Thread, Math, Fmt, Rd, FileStream;
IMPORT ParseParams, Scan, RTMisc;
FROM Stdio IMPORT stderr;

TYPE 
  Options = RECORD
      texture: TEXT;          (* Name of ".vgm" file, without the extension. *)
      imageSize: CARDINAL;    (* Image width and height. *)
      sphereRadius: LONGREAL; (* Sphere radius in voxels *)
      noShadows: BOOLEAN;
      filling: Filling;       (* Filling material for "0" texture value *)
    END;
    
  Filling = {Air, Glass, Plastic};
  
  Scene = RECORD
      sphereRadius: LONGREAL;
      texture: VGM.T;
      
      txtTrans: ARRAY VGM.BYTE OF R3.T; (* Texture transparency table *)
      txtColor: ARRAY VGM.BYTE OF R3.T; (* Texture color table *)
      
      (* The lighting: *)
      noShadows: BOOLEAN;      (* TRUE to skip computation of proper shadows *)
      lampDir: R3.T;           (* Lamp direction from origin *)
      lampColor: R3.T;         (* Lamp light color *)
      ambientLightColor: R3.T; (* Ambient light color. *)
      
      (* The background: *)
      wallCorner: R3.T;      (* Position of wall/floor corner *)
      wallColor: R3.T;       (* Color of walls *)
    END;
    
PROCEDURE Main() =
  <* FATAL Thread.Alerted, Rd.Failure, Wr.Failure, Rd.EndOfFile *>
  VAR 
    ppmName, vgmName: TEXT;
    txt: VGM.T;
  BEGIN
    WITH
      o = GetOptions()
    DO
      vgmName := o.texture & ".vgm";
      WITH
        vf = FileStream.OpenRead(vgmName)
      DO
        txt := VGM.Read(vf); 
        VGM.PrintParms(stderr, txt);
        Rd.Close(vf)
      END;
      
      ppmName := o.texture;
      CASE o.filling OF
      | Filling.Air     => ppmName := ppmName & "-air";
      | Filling.Glass   => ppmName := ppmName & "-gls";
      | Filling.Plastic => ppmName := ppmName & "-pls";
      END;
      ppmName := ppmName & "-sphere.ppm";
      WITH
        pf = FileStream.OpenWrite(ppmName),
        scene = MakeScene(
          sphereRadius := o.sphereRadius, 
          texture := txt,
          noShadows := o.noShadows,
          filling := o.filling
        )^
      DO
        RayTrace(pf, 
          width := o.imageSize, height := o.imageSize,
          obs := R3.Scale(o.sphereRadius, R3.T{1.00, 4.00, 1.50}), 
          focus := R3.T{0.0, 0.0, 0.0},
          pixSize := 3.0d0 * o.sphereRadius / FLOAT(o.imageSize, LONGREAL),
          scene := scene
        );
        Wr.Close(pf);
      END
    END
  END Main;
  
PROCEDURE MakeScene(
    sphereRadius: LONGREAL; 
    READONLY texture: VGM.T; 
    noShadows: BOOLEAN;
    filling: Filling;
  ): REF Scene =
  (* Assumes observer in first quadrant. *)
  BEGIN
    WITH
      rs = NEW(REF Scene),
      s = rs^
    DO
      (* The sphere: *)
      s.sphereRadius := sphereRadius;
      s.texture := texture;
      
      (* Build texture optical tables: *)
      VAR
        t0, c0: R3.T;    (* Transmissivity and color where texture is 0 *)
        t1, c1: R3.T;    (* Transmissivity and color where texture is 0. *)
      BEGIN
        t1 := R3.Zero;
        c1 := R3.T{1.0, 0.90, 0.30};
        CASE filling OF
        | Filling.Plastic =>
          t0 := R3.Zero;
          c0 := R3.T{0.40, 0.60, 0.80};
        | Filling.Glass =>
          t0 := R3.T{0.997, 0.999, 0.999};
          c0 := R3.Zero;
        | Filling.Air =>
          t0 := R3.Ones;
          c0 := R3.Zero;
        END;
        MakeTextureTables(t0, c0,  t1, c1, s.texture.maxval, s.txtTrans, s.txtColor);
      END;
      
      (* The lighting: *)
      s.noShadows := noShadows;
      s.lampDir := R3.Dir(R3.T{1.0, 2.3, 4.7});
      s.lampColor := R3.Ones;
      s.ambientLightColor := R3.Sub(R3.Ones, s.lampColor);
      
      (* The background: *)
      s.wallCorner := R3.Scale(-1.5d0 * sphereRadius, R3.Ones);
      s.wallColor := R3.T{1.0, 1.0, 1.0};
      RETURN rs
    END;
  END MakeScene;
  
PROCEDURE MakeTextureTables(
    READONLY t0, c0: R3.T;
    READONLY t1, c1: R3.T;
    maxval: [0..255];
    VAR t, c: ARRAY [0..255] OF R3.T;
  ) =
  BEGIN
    FOR b := 0 TO 255 DO
      WITH
        v = MIN(1.0, FLOAT(b)/FLOAT(maxval)),
        u = 1.0 - v
      DO
        FOR i := 0 TO 2 DO
          MixMaterials(u, t0[i], c0[i], v, t1[i], c1[i], t[b][i], c[b][i])
        END
      END;
    END;
    PrintTextureTables(maxval, t, c);
  END MakeTextureTables;

PROCEDURE PrintTextureTables(
    maxval: [0..255];
    VAR t, c: ARRAY [0..255] OF R3.T;
  ) =
  <* FATAL Thread.Alerted, Wr.Failure *>
  BEGIN
    FOR b := 0 TO 255 DO
      Wr.PutText(stderr, Fmt.Pad(Fmt.Int(b), 3));
      Wr.PutChar(stderr, ' ');
      Wr.PutChar(stderr, ' ');
      R3.Print(stderr, t[b]);
      Wr.PutChar(stderr, ' ');
      Wr.PutChar(stderr, ' ');
      R3.Print(stderr, c[b]);
      Wr.PutChar(stderr, '\n');
    END;
  END PrintTextureTables;

PROCEDURE MixMaterials(
    u, tu, cu: REAL;  (* Fraction, unit transparency, and particle color of material "0" *)
    v, tv, cv: REAL;  (* Fraction, unit transparency, and particle color of material "1" *)
    VAR t, c: REAL
  ) =
  (* 
    Computes the transparency and color for a mixture 
    of "u" of material "0" and "v" of material "1".
    The mixture is non-linear. *)
    
  BEGIN
    IF u <= 0.0 THEN
      t := tv; c := cv 
    ELSIF v <= 0.0 THEN
      t := tu; c := cu
    ELSIF tu = 0.0 AND tv = 0.0 THEN
      t := 0.0; c := u * cu + v * cv
    ELSE
      WITH 
        uu = FLOAT(u, LONGREAL),
        vv = FLOAT(v, LONGREAL),
        
        ttu = FLOAT(tu, LONGREAL), 
        ttv = FLOAT(tv, LONGREAL),
        
        utu = 1.0d0 - uu * (1.0d0 - ttu),
        vtv = 1.0d0 - vv * (1.0d0 - ttv),
        
        ccu = FLOAT(cu, LONGREAL), 
        ccv = FLOAT(cv, LONGREAL),
        
        ddu = - Math.log(utu),    (* Incremental density of particles of matter "0" *)
        ddv = - Math.log(vtv),    (* Incremental density of particles of matter "1" *)
        dd = ddu + ddv            (* Incremental density of particles *)
      DO
        t := FLOAT(Math.exp(-dd));
        c := FLOAT((ddu * ccu + ddv * ccv)/dd)
      END;
    END;
  END MixMaterials;

PROCEDURE RayTrace(
    wr: Wr.T; 
    width, height: CARDINAL; 
    obs, focus: R3.T; 
    pixSize: LONGREAL;
    READONLY scene: Scene;
  ) =

  VAR 
    gammaTable: ARRAY [0..255] OF CHAR;

  PROCEDURE PutHeader() =
    <* FATAL Thread.Alerted, Wr.Failure *>
    BEGIN
      Wr.PutText(wr, "P6\n"& Fmt.Int(width)& " " & Fmt.Int(height)& "\n255\n");
    END PutHeader;

  PROCEDURE PutPixel(READONLY color: R3.T) =
    <* FATAL Thread.Alerted, Wr.Failure *>
    BEGIN
      Wr.PutChar(wr, gammaTable[ROUND(color[0] * 255.0)]);
      Wr.PutChar(wr, gammaTable[ROUND(color[1] * 255.0)]);
      Wr.PutChar(wr, gammaTable[ROUND(color[2] * 255.0)]);
    END PutPixel;
  
  PROCEDURE InitGammaTable() =
    CONST Gamma = 1.8d0;
    BEGIN
      gammaTable[0] := VAL(0, CHAR);
      FOR i := 1 TO 254 DO 
        WITH
          y = FLOAT(i, LONGREAL)/255.0d0,
          v = Math.exp(Math.log(y)/Gamma)
        DO
          gammaTable[i] := VAL(ROUND(v * 255.0d0), CHAR)
        END;
      END;
      gammaTable[255] := VAL(255, CHAR)
    END InitGammaTable;

  BEGIN
    PutHeader();
    InitGammaTable();
    WITH
      vt = R3.Dir(R3.Sub(obs, focus)),
      vr = R3.Dir(R3.Cross(vt, R3.Axis[2])),
      vs = R3.Cross(vt,vr)
    DO
      FOR u := 1 TO height DO 
        FOR v := 1 TO width DO 
          WITH
            pixel = R3.Add(
              focus, 
              R3.Add(
                R3.Scale(FLOAT (u - height DIV 2, LONGREAL) * pixSize, vs),
                R3.Scale(FLOAT (v - width DIV 2, LONGREAL) * pixSize, vr)
              )
            ),
            color = ComputePixelColor(obs, pixel, scene)
          DO
            PutPixel(color);
          END        
        END;
      END;
    END;
  END RayTrace;

PROCEDURE ComputePixelColor(READONLY obs, pixel: R3.T; READONLY scene: Scene): R3.T =
 
  PROCEDURE TextureTransmitAndColor (READONLY q: R3.T; VAR txtTrans, txtColor: R3.T) =
    (* 
      Returns the optical properties of "txt" matter in the neighborhood 
      of point "q". 
      
      Specifically, for each wavlength band "i", returns 
        
        "txtTrans[i]": the probability of a photon not being blocked 
          (scattered or absorbed) when traversing a unit length of the 
          "txt" material; 
          
        "txtColor[i]": the probability of a *blocked* photon being 
          scattered instead of absorbed.
    *)
    BEGIN
      WITH
        b = VGM.Eval(scene.texture, q[0], q[1], q[2])
      DO
        txtTrans := scene.txtTrans[b];
        txtColor := scene.txtColor[b]
      END
    END TextureTransmitAndColor;
    
  PROCEDURE LayerTransmitAndScatter (
      ds: LONGREAL; 
      txtTrans, txtColor: REAL; 
      VAR layerTrans, layerColor: REAL
    ) = 
    (*
      Computes (for one color band) the transparency and color of a layer
      of material with thickness "ds", given the material's
      transmission probability per unit length "txtTrans" and the color
      "txtColor" of the plastic particles.
    *)
    BEGIN
      IF txtTrans <= 0.0 THEN 
        layerTrans := 0.0; layerColor := txtColor
      ELSIF txtTrans >= 1.0 THEN
        layerTrans := 1.0; layerColor := 0.0
      ELSE
        layerTrans := FLOAT(Math.exp(ds * Math.log(FLOAT(txtTrans, LONGREAL))));
        layerColor := (1.0 - layerTrans) * txtColor
      END;
    END LayerTransmitAndScatter;
    
  VAR
    pixColor, pixTrans: R3.T;

  BEGIN
    pixTrans := R3.Ones;
    pixColor := R3.Zero;
    WITH
      rayDir = R3.Dir(R3.Sub(pixel, obs)),

      (* Object's equation along ray "q(s) = obs + s * rayDir" *)
      radius = scene.sphereRadius,
      a = 1.0d0,
      b = 2.0d0 * R3.Dot(obs, rayDir),
      c = R3.Dot(obs, obs) - radius*radius,

      delta = b * b - 4.0d0 * a * c
    DO
      IF delta > 0.0d0 THEN 
        WITH
          srdelta = Math.sqrt(delta),
          s1 = (- b - srdelta)/(2.0d0 * a),
          s2 = (- b + srdelta)/(2.0d0 * a)
        DO
          VAR s: LONGREAL := MAX(0.0d0, s1);
              ds: LONGREAL := 0.25d0;
              q: R3.T := R3.Add(obs, R3.Scale(s, rayDir));
              dq: R3.T := R3.Scale(ds, rayDir);
              norm: R3.T := R3.Dir(q);
              txtTransmit, txtScatter: R3.T;
              tt, tc: R3.T;
          BEGIN
            WHILE 
              (s < s2) AND 
              (pixTrans[0]>0.01 OR pixTrans[1]>0.01 OR pixTrans[2]>0.01)
            DO 
              (* Compute material properties at point "q": *)
              TextureTransmitAndColor(q, txtTransmit, txtScatter);

              (* Compute color and transparency of layer between "s" and "s+ds": *)
              LayerTransmitAndScatter(ds, txtTransmit[0], txtScatter[0], tt[0], tc[0]);
              LayerTransmitAndScatter(ds, txtTransmit[1], txtScatter[1], tt[1], tc[1]);
              LayerTransmitAndScatter(ds, txtTransmit[2], txtScatter[2], tt[2], tc[2]);

              (* Attenuate color of layer by intervening layers: *)
              tc[0] := pixTrans[0] * tc[0];
              tc[1] := pixTrans[1] * tc[1];
              tc[2] := pixTrans[2] * tc[2];

              (* Adjust color of layer for illumination: *)
              IF tc[0] > 0.001 OR tc[1] > 0.001 OR tc[2] > 0.001 THEN
                tc := ComputeIllumination(q, norm, tc, scene);
              END;

              (* Add color of layer to pixel color: *)
              pixColor[0] := pixColor[0] + tc[0];
              pixColor[1] := pixColor[1] + tc[1];
              pixColor[2] := pixColor[2] + tc[2];

              (* Combine transp of layer with that of previous layers: *)
              pixTrans[0] := pixTrans[0] * tt[0];
              pixTrans[1] := pixTrans[1] * tt[1];
              pixTrans[2] := pixTrans[2] * tt[2];

              s := s + ds;
              q := R3.Add(q, dq);
              norm := R3.Zero;
              IF ds < 0.5d0 THEN 
                (* Increase integration step: *)
                ds := 2.0d0 * ds; 
                dq := R3.Scale(ds, rayDir)
              END;
            END;
          END;
        END
      END;

      (* Add background color: *)
      IF (pixTrans[0]>0.01 OR pixTrans[1]>0.01 OR pixTrans[2]>0.01) THEN
        
        (* Compute background color: *)
        VAR q, n: R3.T;
        BEGIN
          ComputeWallHit(obs, rayDir, scene.wallCorner, q, n);
          WITH
            bc = ComputeIllumination(q, n, scene.wallColor, scene)
          DO
            pixColor[0] := pixColor[0] + pixTrans[0] * bc[0];
            pixColor[1] := pixColor[1] + pixTrans[1] * bc[1];
            pixColor[2] := pixColor[2] + pixTrans[2] * bc[2];
          END;
        END;
      END;
    END;
    RETURN pixColor
  END ComputePixelColor;
  
PROCEDURE ComputeIllumination (
    READONLY p, normal, color: R3.T;
    READONLY scene: Scene
  ): R3.T = 
  (* 
    Returns "color" multiplied componentwise by the total light color
    incident on "p".  
    
    If the "normal" is (0,0,0), the light is assumed to be
    perpendicular to the surface.
  *)

  PROCEDURE TextureTransmit (READONLY q: R3.T; VAR txtTrans: R3.T) =
    (* 
      Same as "TextureTransmitAndColor", omitting the particle color data. 
    *)
    BEGIN
      WITH
        b = VGM.Eval(scene.texture, q[0], q[1], q[2])
     DO
        txtTrans := scene.txtTrans[b]
      END
    END TextureTransmit;
    
  PROCEDURE LayerTransmit (ds: LONGREAL; txtTrans: REAL; VAR layerTrans: REAL) = 
    (*
      Computes (for one color band) the transparency of a layer
      of material with thickness "ds", given the material's
      transmission probability per unit length "txtTrans".
    *)
    BEGIN
      IF txtTrans <= 0.0 THEN 
        layerTrans := 0.0;
      ELSIF txtTrans >= 1.0 THEN
        layerTrans := 1.0;
      ELSE
        layerTrans := FLOAT(Math.exp(ds * Math.log(FLOAT(txtTrans, LONGREAL))))
      END;
    END LayerTransmit;
    
  VAR lc: R3.T := Weigh(scene.lampColor, color);
      tt: R3.T;
  BEGIN
    IF normal # R3.Zero THEN
      lc := R3.Scale(MAX(0.0d0, R3.Dot(normal, scene.lampDir)), lc)
    END;
    IF (NOT scene.noShadows)
    AND (lc[0]>0.01 OR lc[1]>0.01 OR lc[2]>0.01) THEN
     WITH
        rayDir = scene.lampDir,
        
        (* Object's equation along ray "q(s) = p + s * " *)
        radius = scene.sphereRadius,
        a = 1.0d0,
        b = 2.0d0 * R3.Dot(p, rayDir),
        c = R3.Dot(p, p) - radius*radius,

        delta = b * b - 4.0d0 * a * c
      DO
        IF delta > 0.0d0 AND (c < 0.0d0 OR b < 0.0d0) THEN 
          (* Ray crosses sphere; attenuate light by matter in it. *)
          WITH
            srdelta = Math.sqrt(delta),
            s1 = (- b - srdelta)/(2.0d0 * a),
            s2 = (- b + srdelta)/(2.0d0 * a)
          DO
            VAR s: LONGREAL := MAX(0.0d0, s1) + 0.1d0;
                ds: LONGREAL := 0.25d0;
                q: R3.T := R3.Add(p, R3.Scale(s, rayDir));
                dq: R3.T := R3.Scale(ds, rayDir);
                txtTransmit: R3.T;
            BEGIN
              WHILE 
                (s < s2) AND 
                (lc[0]>0.01 OR lc[1]>0.01 OR lc[2]>0.01)
              DO 
                (* Compute material properties at point "q": *)
                TextureTransmit(q, txtTransmit);

                (* Compute transp of layer between "s" and "s+ds": *)
                LayerTransmit(ds, txtTransmit[0], tt[0]);
                LayerTransmit(ds, txtTransmit[1], tt[1]);
                LayerTransmit(ds, txtTransmit[2], tt[2]);

                (* Combine transp of layer with that of previous layers: *)
                lc[0] := lc[0] * tt[0];
                lc[1] := lc[1] * tt[1];
                lc[2] := lc[2] * tt[2];

                s := s + ds;
                q := R3.Add(q, dq);
                IF ds < 0.5d0 THEN 
                  (* Increase integration step: *)
                  ds := 2.0d0 * ds; 
                  dq := R3.Scale(ds, rayDir)
                END;
              END;
            END;
          END
        END
      END
    END;

    (* Add ambient light: *)
    lc := R3.Add(lc, Weigh(color, scene.ambientLightColor));
    RETURN lc
  END ComputeIllumination;

PROCEDURE ComputeWallHit(READONLY obs, rayDir, corner: R3.T; VAR hit, normal: R3.T) =
  (* Assumes "obs" is in the "+++" octant, and "corner" in the "---" *)
  VAR s0, s1, s2, s: REAL;
  BEGIN
    IF rayDir[0] < 0.0 THEN s0 := (corner[0] - obs[0])/rayDir[0] ELSE s0 := 1.0e20 END;
    IF rayDir[1] < 0.0 THEN s1 := (corner[1] - obs[1])/rayDir[1] ELSE s1 := 1.0e20 END;
    IF rayDir[2] < 0.0 THEN s2 := (corner[2] - obs[2])/rayDir[2] ELSE s2 := 1.0e20 END;
    s := MIN(s0, MIN(s1, s2));
    hit := R3.Add(obs, R3.Scale(FLOAT(s, LONGREAL), rayDir));
    IF s = s0 THEN
      normal := R3.T{1.0, 0.0, 0.0}
    ELSIF s = s1 THEN
      normal := R3.T{0.0, 1.0, 0.0}
    ELSE
      normal := R3.T{0.0, 0.0, 1.0}
    END;
  END ComputeWallHit;
  
PROCEDURE Weigh(READONLY x, y: R3.T): R3.T =
  BEGIN
    RETURN R3.T{x[0]*y[0], x[1]*y[1], x[2]*y[2]}
  END Weigh;

PROCEDURE GetOptions (): Options =
  <* FATAL Thread.Alerted, Wr.Failure *>
  VAR o: Options;
  BEGIN
    TRY
      ParseParams.BeginParsing(stderr);                         
      
      ParseParams.GetKeyword("-texture");                               
      o.texture := ParseParams.GetNext();                         

      IF ParseParams.KeywordPresent("-imageSize") THEN                               
        o.imageSize := ParseParams.GetNextInt(1,4096)
      ELSE
        o.imageSize := 400
      END;
      
      IF ParseParams.KeywordPresent("-sphereRadius") THEN                               
        o.sphereRadius := FLOAT(ParseParams.GetNextReal(1.0, 1000.0), LONGREAL)
      ELSE
        o.sphereRadius := 160.0d0
      END;
      
      o.noShadows := ParseParams.KeywordPresent("-noShadows");

      IF ParseParams.KeywordPresent("-plastic") THEN
        o.filling := Filling.Plastic
      ELSIF ParseParams.KeywordPresent("-air") THEN
        o.filling := Filling.Air
      ELSIF ParseParams.KeywordPresent("-glass") THEN
        o.filling := Filling.Glass
      ELSE
        o.filling := Filling.Glass
      END;

      ParseParams.EndParsing();                                       
    EXCEPT                                                            
    | Scan.BadFormat =>                                              
        Wr.PutText(stderr, "Usage: tarcesphere\\\n");
        Wr.PutText(stderr, "  -texture <name>\\\n");        
        Wr.PutText(stderr, "  [ -sphereRadius <num> ]\\\n");        
        Wr.PutText(stderr, "  [ -imageSize <int> ]\\\n"); 
        Wr.PutText(stderr, "  [ -noShadows ]\\\n"); 
        Wr.PutText(stderr, "  [ -plastic | -hollow | -glass ]\n"); 
        RTMisc.Exit (1);                                              
    END;
    RETURN o
  END GetOptions;
  
BEGIN
  Main();
END TraceSphere.