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

TYPE 
  Tri = RECORD u, v, w: CARDINAL END;
  BOOLS = ARRAY OF BOOLEAN;
  LONGS = ARRAY OF LONGREAL;

REVEAL
  T = Public BRANDED OBJECT
      top: Topology;
      NT: CARDINAL;          (* Number of triangles on "triang" *)
      K: LONGREAL;           (* Energy normalization factor *)      
      vVar: REF BOOLS;       (* Which vertices are variable *)
      tri: REF ARRAY OF Tri; (* The existing triangles *)
      tVar: REF BOOLS;       (* Which triangles in "tri" are variable *)
      bar: REF Coords;       (* (Work) Barycenters of triangles in "tri" *)
      eDbar: REF Coords;     (* (Work) Gradient of "e" rel "bar" *)
      rsq: REF LONGS;        (* (Work) Nominal triangle radius, squared *)
      eDrsq: REF LONGS;      (* (Work) Gradient of "e" rel "rsq" *)
    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
    erg.top := top;
    (* Work areas: *)
    erg.tri := CollectExistingTriangles(top);
    erg.NT := NUMBER(erg.tri^);
    erg.K := 2.0d0/FLOAT(erg.NT*erg.NT, LONGREAL);
    erg.vVar := NEW(REF BOOLS, top.NV);
    erg.tVar := NEW(REF BOOLS, erg.NT);
    erg.bar := NEW(REF Coords, erg.NT);
    erg.eDbar := NEW(REF Coords, erg.NT);
    erg.rsq := NEW(REF ARRAY OF LONGREAL, erg.NT);
    erg.eDrsq := NEW(REF ARRAY OF LONGREAL, erg.NT);
    (* In case the client forgets to call "defVar": *)
    FOR i := 0 TO erg.NT-1 DO erg.tVar[i] := TRUE END;
  END DefTop;
  
PROCEDURE CollectExistingTriangles(READONLY top: Topology): REF ARRAY OF Tri =
  VAR NT: CARDINAL := 0;
  BEGIN
    WITH
      NF = top.NF,
      t = NEW(REF ARRAY OF Tri, NF)^
    DO
      FOR i := 0 TO NF-1 DO 
        WITH 
          f = top.face[i], a = top.side[i]
        DO
          IF f.exists AND Degree(a) = 3 THEN 
            WITH
              u = OrgV(Rot(a)), 
              v = OrgV(Rot(Onext(a))), 
              w = OrgV(Rot(Onext(Onext(a))))
            DO
              <* ASSERT u.exists AND v.exists AND w.exists *>
              t[NT] := Tri{u.num, v.num, w.num}
            END;
            INC(NT)
          END
        END
      END;
      WITH r = NEW(REF ARRAY OF Tri, NT) DO
        r^ := SUBARRAY(t, 0, NT);
        RETURN r
      END
    END
  END CollectExistingTriangles;
  
PROCEDURE DefVar(erg: T; READONLY variable: ARRAY OF BOOLEAN) =
  BEGIN
    (*
      Separate the existing triangles now in "triang"
      into the variable list "tvar" and the fixed list 
      "tfix".  (A triangle is variable if any of its vertices
      is variable.)
    *)
    WITH
      NT = erg.NT,
      NV = erg.top.NV,
      vVar = erg.vVar^,
      tri = erg.tri^,
      tVar = erg.tVar^
    DO
      <* ASSERT NUMBER(variable) = NV *>
      vVar := variable;

      FOR i := 0 TO NT-1 DO
        WITH ti = tri[i] DO
          tVar[i] := vVar[ti.u] OR vVar[ti.v] OR vVar[ti.w]
        END
      END;
    END
  END DefVar;
  
PROCEDURE Eval(
    erg: T; 
    READONLY c: Coords; 
    VAR e: LONGREAL; 
    grad: BOOLEAN;
    VAR eDc: Gradient;
  ) =
  BEGIN
    WITH
      NT = erg.NT,
      NV = erg.top.NV,
      vVar = erg.vVar^,
      tri = erg.tri^,
      tVar = erg.tVar^,
      bar = erg.bar^,
      eDbar = erg.eDbar^,
      rsq = erg.rsq^,
      eDrsq = erg.eDrsq^,

      fuzz = FLOAT(erg.fuzz, LONGREAL),
      fuzz2 = fuzz*fuzz,
      K = erg.K
    DO
    
      PROCEDURE Compute_rsq(READONLY t: Tri): LONGREAL =
        (*
          Computes the 'radius squared' of "t", actually
          the mean squared vertex-barycenter distance
        *)
        VAR r: LONGREAL := 0.0d0;
        BEGIN
          WITH
            u = c[t.u], 
            v = c[t.v], 
            w = c[t.w]
          DO
            WITH uv0 = u[0] - v[0], vw0 = v[0] - w[0], wu0 = w[0] - u[0] DO
              r := r + uv0*uv0 + vw0*vw0 + wu0*wu0
            END;
            WITH uv1 = u[1] - v[1], vw1 = v[1] - w[1], wu1 = w[1] - u[1] DO
              r := r + uv1*uv1 + vw1*vw1 + wu1*wu1
            END;
            WITH uv2 = u[2] - v[2], vw2 = v[2] - w[2], wu2 = w[2] - u[2] DO
              r := r + uv2*uv2 + vw2*vw2 + wu2*wu2
            END;
            RETURN r/9.0d0
          END;
        END Compute_rsq;
        
      PROCEDURE Distribute_eDrsq(READONLY t: Tri; READONLY eDr2: LONGREAL) =
        (*
          Adds to the gradient "eDc" the component due to the influence
          "eDr2" of the radius squared of "tri" on the energy.
        *)
        BEGIN
          WITH
            eDs = 2.0d0/9.0d0*eDr2,
            u = c[t.u], v = c[t.v], w = c[t.w]
          DO
            IF vVar [t.u] THEN
              WITH eDu = eDc[t.u] DO
                eDu[0] := eDu[0] + eDs * (u[0] - v[0] + u[0] - w[0]);
                eDu[1] := eDu[1] + eDs * (u[1] - v[1] + u[1] - w[1]);
                eDu[2] := eDu[2] + eDs * (u[2] - v[2] + u[2] - w[2]);
              END
            END;

            IF vVar [t.v] THEN
              WITH eDv = eDc[t.v] DO
                eDv[0] := eDv[0] + eDs * (v[0] - w[0] + v[0] - u[0]);
                eDv[1] := eDv[1] + eDs * (v[1] - w[1] + v[1] - u[1]);
                eDv[2] := eDv[2] + eDs * (v[2] - w[2] + v[2] - u[2]);
              END
            END;

            IF vVar [t.w] THEN
              WITH eDw = eDc[t.w] DO
                eDw[0] := eDw[0] + eDs * (w[0] - u[0] + w[0] - v[0]);
                eDw[1] := eDw[1] + eDs * (w[1] - u[1] + w[1] - v[1]);
                eDw[2] := eDw[2] + eDs * (w[2] - u[2] + w[2] - v[2]);
              END
            END
          END;
        END Distribute_eDrsq;
          
      PROCEDURE Compute_bar(READONLY t: Tri): LR3.T =
        VAR b: LR3.T;
        BEGIN
          WITH
            u = c[t.u], 
            v = c[t.v], 
            w = c[t.w]
          DO
            b[0]:= (u[0] + v[0] + w[0])/3.0d0;
            b[1]:= (u[1] + v[1] + w[1])/3.0d0;
            b[2]:= (u[2] + v[2] + w[2])/3.0d0;
            RETURN b
          END;
        END Compute_bar;
      
      PROCEDURE Distribute_eDbar(READONLY t: Tri; READONLY eDb: LR3.T) =
        (*
          Adds to the gradient "eDc" the component due to the influence
          "eDb" of "tri"'s barycenter on the energy.
        *)
        BEGIN
          WITH
            eDs0 = eDb[0]/3.0d0,
            eDs1 = eDb[1]/3.0d0,
            eDs2 = eDb[2]/3.0d0
          DO
            IF vVar [t.u] THEN
              WITH eDu = eDc[t.u] DO
                eDu[0] := eDu[0] + eDs0;
                eDu[1] := eDu[1] + eDs1;
                eDu[2] := eDu[2] + eDs2;
              END
            END;

            IF vVar [t.v] THEN
              WITH eDv = eDc[t.v] DO
                eDv[0] := eDv[0] + eDs0;
                eDv[1] := eDv[1] + eDs1;
                eDv[2] := eDv[2] + eDs2;
              END
            END;

            IF vVar [t.w] THEN
              WITH eDw = eDc[t.w] DO
                eDw[0] := eDw[0] + eDs0;
                eDw[1] := eDw[1] + eDs1;
                eDw[2] := eDw[2] + eDs2;
              END
            END;
          END;
        END Distribute_eDbar;

      PROCEDURE AddTrianglePairEnergy(i, j: CARDINAL) =
        (*
          Adds to "e" the electric potential energy between the 
          triangles "tri[i]" and "tri[j]", from "bar[i],
          "bar[j]", "rsq[i]", and "rsq[j]". Assumes "i # j".

          Also adds to "eDrsq[i]", "eDrsq[j]", "eDbar[i]"
          and "eDbar[j]" the corresponding derivatives,
          if the triangles are variable.
        *)
        BEGIN
          WITH
            bi = bar[i],      bj = bar[j],
            eDbi = eDbar[i],  eDbj = eDbar[j],
            si = rsq[i],      sj = rsq[j],
            eDsi = eDrsq[i],  eDsj = eDrsq[j],
           
            vij = LR3.Sub(bj, bi),
            s = si + sj,
            r2 = LR3.NormSqr(vij) + fuzz2 * s,
            r = Math.sqrt(r2),
            a = K/r
          DO
            e := e + a;
            IF grad THEN
              WITH
                aDr = - a/r,
                aDr2 = aDr * 0.5d0 / r,
                aDs = aDr2 * fuzz2,
                aDvij = LR3.Scale(2.0d0 * aDr2, vij)
              DO
                IF tVar[i] THEN
                  eDsi := eDsi + aDs;
                  eDbi := LR3.Sub(eDbi, aDvij)
                END;
                IF tVar[j] THEN
                  eDsj := eDsj + aDs;
                  eDbj := LR3.Add(eDbj, aDvij)
                END
              END
            END
          END;
        END AddTrianglePairEnergy;

      BEGIN
        (* Compute baricenters and face sizes, clear derivatives: *)
        FOR i := 0 TO NT-1 DO 
          bar[i] := Compute_bar(tri[i]);
          rsq[i] := Compute_rsq(tri[i]);
          IF grad THEN
            eDbar[i] := LR3.T{0.0d0, 0.0d0, 0.0d0};
            eDrsq[i] := 0.0d0;
          END
        END;
        
        (* Compute energies and the internal derivatives "eDrsq", "eDbar": *)
        e := 0.0d0;
        FOR i := 0 TO NT-1 DO
          IF tVar[i] THEN
            FOR j := 0 TO NT-1 DO
              IF i # j AND (NOT tVar[j] OR i < j) THEN
                AddTrianglePairEnergy(i, j)
              END
            END
          END
        END;

        (* Distribute "eDrsq", "eDbar" onto "eDc": *)
        FOR i := 0 TO NV-1 DO eDc[i] := LR3.T{0.0d0, 0.0d0, 0.0d0} END;
        IF grad THEN
          FOR i := 0 TO NT-1 DO 
            IF tVar[i] THEN
              Distribute_eDbar(tri[i], eDbar[i]);
              Distribute_eDrsq(tri[i], eDrsq[i]);
            END
          END
        END
      END
    END
  END Eval;

PROCEDURE Name(erg: T): TEXT =
  BEGIN
    RETURN "Elect(fuzz := " & Fmt.Real(erg.fuzz, prec := 3) & ")"
  END Name;

BEGIN
END ElectEnergy.