MODULE PotenEnergy;
IMPORT LR3; 
FROM Triang IMPORT Topology;
FROM Energy IMPORT Coords, Gradient;

REVEAL
  T = Public BRANDED OBJECT
      top: Topology;
      K: LONGREAL;                     (* Energy normalization constant *)
      (* Defined by "defVar: *)
      termVar: REF ARRAY OF BOOLEAN;   (* Tells which terms are variable *)
    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;
    erg.K := 1.0d0/FLOAT(top.NV, LONGREAL);
    erg.termVar := NEW(REF ARRAY OF BOOLEAN, top.NV);
    (* In case the client forgets to call "defVar": *)
    FOR i := 0 TO LAST(erg.termVar^) DO erg.termVar[i] := FALSE END;
  END DefTop;
  
PROCEDURE DefVar(erg: T; READONLY variable: ARRAY OF BOOLEAN) =
  BEGIN
    (*
      A vertex contributes to the energy only if it exists
      and is variable.
    *)
    WITH
      NV = erg.top.NV,
      vertex = erg.top.vertex^,
      termVar = erg.termVar^
    DO
      <* ASSERT NUMBER(variable) = NV *>
      FOR v := 0 TO NV-1 DO 
        termVar[v] := variable[v] AND vertex[v].exists
      END;    
    END
  END DefVar;
  
PROCEDURE Eval(
    erg: T; 
    READONLY c: Coords; 
    VAR e: LONGREAL; 
    grad: BOOLEAN;
    VAR eDc: Gradient;
  ) =
  BEGIN
    WITH
      NV = erg.top.NV,
      termVar = erg.termVar^,
      K = erg.K
    DO

      PROCEDURE AccumTerm (
          READONLY cv: LR3.T; 
          VAR evDcv: LR3.T;
        ) =
        (*
          Adds to "e" the energy term corresponding to a vertex at "cv".
          Returns also the gradient "evDcv" of that term relative to "cv".
        *)
        BEGIN
          e := e + K * LR3.NormSqr(cv);
          IF grad THEN
            evDcv := LR3.Scale(2.0d0 * K, cv)
          ELSE
            evDcv := LR3.T{0.0d0, 0.0d0, 0.0d0}
          END;
        END AccumTerm;

      BEGIN
        e := 0.0d0;
        FOR v := 0 TO NV-1 DO
          IF termVar[v] THEN
            AccumTerm(c[v], eDc[v]);
          ELSE
            eDc[v] := LR3.T{0.0d0, 0.0d0, 0.0d0};
          END;
        END;
      END;
    END
  END Eval;

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

BEGIN
END PotenEnergy.