MODULE CurvEnergy;
IMPORT LR3, LR3Extras, Triang, Math; 
FROM Oct IMPORT Onext, Rot, Tor, Sym, Degree;
FROM Triang IMPORT OrgV, LeftF, Topology;
FROM Energy IMPORT Coords, Gradient;

REVEAL
  T = Public BRANDED OBJECT
      top: Topology;
      NS: CARDINAL; 
      K: LONGREAL;
      vVar: REF ARRAY OF BOOLEAN;      (* Which vertices are variable *)
      (* Number of relevant springs *)
      su, sv: REF ARRAY OF CARDINAL;   (* The endpoints of those springs *)
      sp, sq: REF ARRAY OF CARDINAL;   (* The wingtips of those springs *)
    OVERRIDES
      init := Init;
      defTop := DefTop;
      defVar := DefVar;
      eval := Eval;
      name := Name;
    END;
    
CONST
  Zero = LR3.T{0.0d0, 0.0d0, 0.0d0};
    
PROCEDURE Init(erg: T): T =
  BEGIN
    RETURN erg
  END Init;

PROCEDURE DefTop(erg: T; READONLY top: Topology) =
  BEGIN
    erg.top := top;
    erg.vVar := NEW(REF ARRAY OF BOOLEAN, top.NV);
    erg.su := NEW(REF ARRAY OF CARDINAL, top.NE);
    erg.sv := NEW(REF ARRAY OF CARDINAL, top.NE);
    erg.sp := NEW(REF ARRAY OF CARDINAL, top.NE);
    erg.sq := NEW(REF ARRAY OF CARDINAL, top.NE);
    erg.K := Math.sqrt(FLOAT(top.NF, LONGREAL))/64.0d0;
    (* Just in case the client forgets to call "defVar": *)
    erg.NS := 0;
    FOR i := 0 TO top.NV-1 DO 
      erg.vVar[i] := FALSE;
    END;
  END DefTop;

PROCEDURE DefVar(erg: T; READONLY variable: ARRAY OF BOOLEAN) =
  BEGIN
    (*
      Decide which edges and faces are relevant.

      An edge is relevant if it exists, has the "spring" bit set,
      the two adjacent faces exist and are triangles, 
      and at least one of the four vertices in
      those faces is variable.
      
      A face is relevant if it is adjacent to a relevant edge.
    *)
    WITH
      NS = erg.NS,
      NV = erg.top.NV,
      NE = erg.top.NE,
      edge = erg.top.edge^,
      vVar = erg.vVar^,
      su = erg.su^,
      sv = erg.sv^,
      sp = erg.sp^,
      sq = erg.sq^
    DO
      <* ASSERT NUMBER(variable) = NV *>
      vVar := variable;
      
      NS := 0;
      FOR i := 0 TO NE-1 DO
        WITH 
          a = edge[i], 
          e = NARROW(a.edge, Triang.Edge),
          f = LeftF(a),
          g = LeftF(Sym(a)),
          u = OrgV(a),                   uvar = vVar[u.num],
          v = OrgV(Sym(a)),              vvar = vVar[v.num],
          p = OrgV(Sym(Onext(a))),       pvar = vVar[p.num],
          q = OrgV(Sym(Onext(Sym(a)))),  qvar = vVar[q.num]
        DO
          IF e.spring
          AND f.exists AND g.exists   
          AND Degree(Rot(a)) = 3 AND Degree(Tor(a)) = 3 
          AND (uvar OR vvar OR pvar OR qvar)
          THEN
            <* ASSERT e.exists *>
            <* ASSERT u.exists AND v.exists AND p.exists AND q.exists *>
            su[NS] := u.num; sv[NS] := v.num;
            sp[NS] := p.num; sq[NS] := q.num;
            INC(NS)
          END
        END
      END;
    END
  END DefVar;

PROCEDURE Eval(
    erg: T; 
    READONLY c: Coords; 
    VAR e: LONGREAL; 
    grad: BOOLEAN;
    VAR eDc: Gradient;
  ) =
  BEGIN
    WITH
      NS = erg.NS,
      NV = erg.top.NV,
      vVar = erg.vVar^,
      su = erg.su^,
      sv = erg.sv^,
      sp = erg.sp^,
      sq = erg.sq^,
      K = erg.K
    DO
      
      PROCEDURE AccumulateTerms() =
        BEGIN
          e := 0.0d0;
          FOR v := 0 TO NV-1 DO eDc[v] := Zero END;
          FOR i := 0 TO NS-1 DO
            AccumulateTerm(su[i], sv[i], sp[i], sq[i]);
          END;
        END AccumulateTerms;

      PROCEDURE AccumulateTerm(iu, iv, ip, iq: CARDINAL) =
        (*
          Adds to "e" the energy for one hinge with endpoints
          "c[iu]" and "c[iv]", and wingtips "c[ip]" and "c[iq]".  
          Also adds to "eDc" the corresponding gradients.
        *)
        BEGIN
          WITH
            u = c[iu], v = c[iv],  
            p = c[ip], q = c[iq]
          DO
            IF u[0] # v[0] OR u[1] # v[1] OR u[2] # v[2] THEN
              WITH
                h = LR3.Sub(v, u),  (* Hinge vector *)
                a = LR3.Sub(q, u),  (* One wing vector *)
                b = LR3.Sub(v, p)   (* Another wing vector *)
              DO
                VAR eTerm: LONGREAL;
                    eDa, eDb, eDh: LR3.T;
                BEGIN
                  ComputeTermFromHingeVectors(a, b, h, eTerm, eDa, eDb, eDh);
                  e := e + eTerm;
                  IF grad THEN
                    IF vVar[ip] THEN
                      eDc[ip] := LR3.Sub(eDc[ip], eDb)
                    END;
                    IF vVar[iq] THEN
                      eDc[iq] := LR3.Add(eDc[iq], eDa)
                    END;
                    IF vVar[iu] THEN
                      eDc[iu] := LR3.Sub(eDc[iu], LR3.Add(eDh, eDa))
                    END;
                    IF vVar[iv] THEN
                      eDc[iv] := LR3.Add(eDc[iv], LR3.Add(eDb, eDh))
                    END
                  END
                END
              END
            END
          END
        END AccumulateTerm;
              
      PROCEDURE ComputeTermFromHingeVectors(
          READONLY a, b, h: LR3.T;
          VAR eTerm: LONGREAL;
          VAR eDa, eDb, eDh: LR3.T;
        ) =
        (*
          Computes the hinge energy "eTerm" given the hinge vector "h"
          and two vectors "a" and "b" that define the exterior dihedral
          angle.  Also sets the term's derivatives "eDa", "eDb", "eDh".
          
          Specifically, the energy is the length of the hinge times
          some concave function of the angle between the projections
          of "a" and "b" on a plane perpendicular to "h".
          The energy is minimum when the angle is zero.
        *)
        BEGIN
          WITH
            m = LR3.Norm(h),
            r = LR3Extras.Cross(a, h),
            s = LR3Extras.Cross(b, h)
          DO
            IF m = 0.0d0 THEN
              eTerm := 0.0d0;
              eDa := Zero; eDb := Zero; eDh := Zero;
            ELSE
              VAR z: LONGREAL; (* Energy density per unit length *)
                  zDr, zDs: LR3.T;
              BEGIN
                ComputeEnergyDensityFromNormalVectors(r, s, z, zDr, zDs);
                eTerm := m * z;
                IF grad THEN
                  WITH
                    eDr = LR3.Scale(m, zDr),
                    eDs = LR3.Scale(m, zDs)
                  DO
                    eDa := LR3Extras.Cross(h, eDr);
                    eDb := LR3Extras.Cross(h, eDs);
                    eDh := LR3.Add(
                      LR3.Scale(z/m, h),
                      LR3.Add(
                        LR3Extras.Cross(eDr, a),
                        LR3Extras.Cross(eDs, b)
                      )
                    )
                  END
                ELSE
                  eDa := Zero; eDb := Zero; eDh := Zero;
                END
              END
            END
          END
        END ComputeTermFromHingeVectors;

      PROCEDURE ComputeEnergyDensityFromNormalVectors(
          READONLY r, s: LR3.T;
          VAR z: LONGREAL;
          VAR zDr, zDs: LR3.T;
        ) =
        (*
          Computes the linear energy density "z" given two vectors "r" and "s",
          orthogonal to the hinge, that define the external dihedral angle. *)
        BEGIN
          WITH
            m = LR3.Norm(r),
            n = LR3.Norm(s),
            o = LR3.Dot(r, s),
            d = m*n,
            q = o/d
          DO
            IF d # 0.0d0 THEN
              z := 2.0d0 * K * (1.0d0 - q);
              IF grad THEN 
                WITH
                  zDq = - 2.0d0 * K,
                  zDo = zDq /d,
                  zDd = - zDq * q / d,
                  zDm = zDd * n,
                  zDn = zDd * m
                DO
                  zDr := LR3.Mix(zDo, s, zDm/m, r);
                  zDs := LR3.Mix(zDo, r, zDn/n, s);
                END
              END
            END
          END
        END ComputeEnergyDensityFromNormalVectors;

      BEGIN
        AccumulateTerms();
      END;
    END
  END Eval;

PROCEDURE Name(<*UNUSED*> erg: T): TEXT =
  BEGIN
    RETURN "Curv()"
  END Name;

BEGIN
END CurvEnergy.