MODULE PAreaEnergy;
IMPORT LR3, LR3Extras, Fmt, Triang, Math; 
FROM Oct IMPORT Onext, Rot, Degree;
FROM Triang IMPORT OrgV, Topology;
FROM Energy IMPORT Coords, Gradient;

TYPE
  BOOLS = ARRAY OF BOOLEAN;
  LONGS = ARRAY OF LONGREAL;

REVEAL
  T = Public BRANDED OBJECT
      top: Topology;            (* The topology *)
      refArea: LONGREAL;        (* Ideal area of a patch, for sphere of unit radius *)
      NP: CARDINAL;             (* Max patch number + 1 *)
      vVar: REF BOOLS;          (* TRUE if vertex is variable *)
      faceRelevant: REF BOOLS;  (* TRUE if face is relevant *)
      patchRelevant: REF BOOLS; (* TRUE if logical patch is relevant *)
      fn: REF ARRAY OF LR3.T;   (* (Work) Cross product of face sides *)
      pa: REF LONGS;            (* (Work) Total area of each patch *)
      eDpa: REF LONGS;          (* (Work) Derivative of energy rel. to patch areas *)
    OVERRIDES
      init := Init;
      defTop := DefTop;
      defVar := DefVar;
      eval := Eval;
      name := Name;
    END;

PROCEDURE Init(erg: T): T =
  BEGIN
    RETURN erg
  END Init;

PROCEDURE DefTop(erg: T; READONLY top: Topology) =
  BEGIN
    WITH
      NP = erg.NP,
      NF = top.NF
    DO
      erg.top := top;
      erg.vVar := NEW(REF BOOLS, top.NV);
      erg.faceRelevant := NEW(REF BOOLS, top.NF);
      erg.fn := NEW(REF ARRAY OF LR3.T, top.NF);

      (* Find highest patch number: *)
      NP := 0;
      FOR i := 0 TO NF-1 DO
        WITH fi = top.face[i] DO
          IF fi.exists AND Degree(top.side[i]) = 3 THEN 
            NP := MAX(NP, fi.patch+1);
          END
        END
      END;
      
      (* Allocate patch tables: *)
      erg.patchRelevant := NEW(REF BOOLS, NP);
      erg.pa := NEW(REF LONGS, NP);
      erg.eDpa := NEW(REF LONGS, NP);

      (* Count distinct patches and compute ideal patch area: *)
      VAR nPatches: CARDINAL := 0;
      BEGIN
        (* Borrow the "patchRelevant" vector for a moment... *)
        FOR i := 0 TO erg.NP-1 DO erg.patchRelevant[i] := TRUE END;
        FOR i := 0 TO NF-1 DO
          WITH fi = top.face[i] DO
            IF fi.exists AND Degree(top.side[i]) = 3 THEN 
              IF erg.patchRelevant[fi.patch] THEN 
                INC(nPatches); erg.patchRelevant[fi.patch] := FALSE
              END;
            END
          END
        END;
        erg.refArea := 4.0d0*FLOAT(Math.Pi, LONGREAL)/FLOAT(nPatches, LONGREAL)
      END;

      (* Just in case the client forgets to call "defVar": *)
      FOR i := 0 TO top.NV-1 DO erg.vVar[i] := FALSE END;
      FOR i := 0 TO top.NF-1 DO erg.faceRelevant[i] := FALSE END;
      FOR i := 0 TO erg.NP-1 DO erg.patchRelevant[i] := FALSE END;
    END
  END DefTop;
 
PROCEDURE DefVar(erg: T; READONLY variable: BOOLS) =
  BEGIN
    (*
      Decide which faces and patches are relevant.

      A face is relevant iff it exists, has exactly three sides, and
      belongs to a relevant patch.  A patch is relevant iff it
      includes at least one existing triangular face having at least
      one variable corner.
    *)
    WITH
      NP = erg.NP,
      NV = erg.top.NV,
      NF = erg.top.NF,
      vVar = erg.vVar^,
      face = erg.top.face^,
      side = erg.top.side^,
      faceRelevant = erg.faceRelevant^,
      patchRelevant = erg.patchRelevant^
    DO
      <* ASSERT NUMBER(variable) = NV *>
      vVar := variable;
      
      (* Find the relevant patches: *)
      FOR p := 0 TO NP-1 DO patchRelevant[p] := FALSE END;
      FOR i := 0 TO NF-1 DO
        WITH 
          f = face[i], 
          a = side[i],
          u = OrgV(Rot(a)),
          v = OrgV(Rot(Onext(a))),
          w = OrgV(Rot(Onext(Onext(a)))),
          fvar = vVar[u.num] OR vVar[v.num] OR vVar[w.num]
        DO
          IF f.exists AND Degree(a) = 3 AND fvar THEN
            <* ASSERT u.exists AND v.exists AND w.exists *>
            WITH p = f.patch DO patchRelevant[p] := TRUE END
          END
        END
      END;
      
      (* Find relevant faces: *)
      FOR i := 0 TO NF-1 DO
        WITH 
          f = face[i], 
          a = side[i]
        DO
          IF f.exists AND Degree(a) = 3 THEN
            faceRelevant[i] := patchRelevant[f.patch]
          ELSE
            faceRelevant[i] := FALSE
          END
        END
      END
    END
  END DefVar;
                
PROCEDURE Eval(
    erg: T; 
    READONLY c: Coords; 
    VAR e: LONGREAL; 
    grad: BOOLEAN;
    VAR eDc: Gradient;
  ) =
  BEGIN
    WITH
      NP = erg.NP,
      NV = erg.top.NV,
      NF = erg.top.NF,
      face = erg.top.face^,
      side = erg.top.side^,
      vVar = erg.vVar^,
      faceRelevant = erg.faceRelevant^,
      fn = erg.fn^,
      patchRelevant = erg.patchRelevant^,
      pa = erg.pa^,
      eDpa = erg.eDpa^,
      refArea = erg.refArea,
      aCoef = 1.0d0/(FLOAT(erg.area, LONGREAL) * refArea)
    DO

      PROCEDURE Accum_a(READONLY u, v, w: LR3.T; VAR a: LONGREAL; VAR n: LR3.T) =
        (*
          Adds to "a" the area of the triangle "u v w".  Also stores in
          "n" the cross product "(v - u)X(w - v)".
        *)
        BEGIN
          WITH
            uv = LR3.Sub(v, u),
            vw = LR3.Sub(w, v)
          DO
            n := LR3Extras.Cross(uv, vw);
            a := a + 0.5d0 * LR3.Norm(n);
          END
        END Accum_a;

      PROCEDURE Accum_e_from_a(a: LONGREAL; VAR e: LONGREAL; VAR etDa: LONGREAL) =
        (*
          Adds to "e" the energy term corresponting to a patch 
          with area "a".  Also, if "grad" is true, stores in "eDa" the
          derivative of that term.
        *)
        BEGIN
          IF a <= 0.0d0 THEN 
            etDa := 0.0d0
          ELSE
            WITH
              r = a * aCoef,
              s = 1.0d0/r,
              h = r + s - 2.0d0
            DO
              e := e + h;
              IF grad THEN
                WITH
                  eDr = 1.0d0 - s * s
                DO
                  etDa := eDr * aCoef
                END
              ELSE
                etDa := 0.0d0;
              END
            END
          END
        END Accum_e_from_a;

      PROCEDURE Distribute_eDa(
          iu, iv, iw: CARDINAL; 
          READONLY n: LR3.T; 
          eDa: LONGREAL;
        )=
        (*
          Accumulates in "eDc" the gradient of "e" relative
          to the corners of the triangle "iu iv iw", given 
          the cross-product "n = (v-u)X(w-v)" and the derivative 
          "eDa" of "e" relative to the triangle's area "a".
        *)
        BEGIN
          WITH
            u = c[iu], v = c[iv], w = c[iw],
            uv = LR3.Sub(v, u),
            vw = LR3.Sub(w, v),
            s = LR3.Norm(n)
          DO
            IF s # 0.0d0 THEN
              WITH
                eDs = 0.5d0 * eDa,
                eDn = LR3.Scale(eDs/s, n),
                eDuv = LR3Extras.Cross(vw, eDn),
                eDvw = LR3Extras.Cross(eDn, uv)
              DO
                IF vVar[iu] THEN
                  eDc[iu] := LR3.Sub(eDc[iu], eDuv)
                END;
                IF vVar[iv] THEN
                  eDc[iv] := LR3.Add(eDc[iv], LR3.Sub(eDuv, eDvw))
                END;
                IF vVar[iw] THEN
                  eDc[iw] := LR3.Add(eDc[iw], eDvw)
                END
              END
            END
          END
        END Distribute_eDa;

      BEGIN
        (* Clear area accumulators: *)
        FOR p := 0 TO NP-1 DO pa[p] := 0.0d0 END;

        (* Enumerate triangles and accumulate patch areas: *)
        FOR j := 0 TO NF-1 DO
          IF faceRelevant[j] THEN
            WITH
              t = side[j],
              un = OrgV(Rot(t)).num,
              vn = OrgV(Rot(Onext(t))).num,
              wn = OrgV(Rot(Onext(Onext(t)))).num,
              p = face[j].patch
            DO
              Accum_a(c[un], c[vn], c[wn], pa[p], fn[j])
            END
          END            
        END;
        
        (* Compute energy "e" from patch areas, and the gradient "eDpa": *)
        (* Note that patches with 0 triangles will have patchRelevant[k] = FALSE. *)
        e := 0.0d0;
        FOR p := 0 TO NP-1 DO 
          IF patchRelevant[p] THEN
            Accum_e_from_a(pa[p], e, eDpa[p])
          ELSE
            eDpa[p] := 0.0d0
          END
        END;

        (* Now distribute "eDpa" over "eDc": *)
        FOR i := 0 TO NV-1 DO eDc[i] := LR3.T{0.0d0, ..} END;
        IF grad THEN
          FOR j := 0 TO NF-1 DO 
            IF faceRelevant[j] THEN
              WITH
                t = side[j],
                un = OrgV(Rot(t)).num,
                vn = OrgV(Rot(Onext(t))).num,
                wn = OrgV(Rot(Onext(Onext(t)))).num,
                p = face[j].patch
              DO
                Distribute_eDa(un, vn, wn, fn[j], eDpa[p])
              END
            END
          END
        END;
      END
    END
  END Eval;
  
PROCEDURE Name(erg: T): TEXT =
  BEGIN
    RETURN 
      "PArea(ideal := " & Fmt.Real(erg.area, prec := 3) & ")"
  END Name;

BEGIN
END PAreaEnergy.