MODULE EquaAngleEnergy;

IMPORT Triangulation, Math, LR4, Octf;

FROM Triangulation IMPORT Topology, Ppos, Pneg, OrgV, Pair;
FROM Energy IMPORT Coords, Gradient;
FROM Octf IMPORT Fnext, Enext, Enext_1, Clock;


TYPE
  BOOLS = ARRAY OF BOOLEAN;
  LONGS = ARRAY OF LONGREAL;
  DRF = REF ARRAY OF CARDINAL;
  Face = Triangulation.Face;
  Edge = Triangulation.Edge;
  FACES = ARRAY OF Face;
  Cosines = REF LONGS;
  Faces = REF FACES;

REVEAL
  T = Public BRANDED OBJECT
      K: LONGREAL;                        (* The energy normalization factor *)
      top: Topology;                      (* The topology *)
      vVar: REF BOOLS;                    (* TRUE if vertex is variable *)
      edgeRelevant: REF BOOLS;            (* TRUE if edge is relevant *)
      cosine: REF ARRAY OF Cosines;       (* The diedral angles *)
      drf: DRF;                           (* The Facets Ring Degree for edge *)
      faces: REF ARRAY OF Faces;          (* The faces-ring for edge *)
      eDcosine: REF ARRAY OF Cosines;     (* (Work) Gradient of "e" rel. to
                                             the dihedral average angle "a".
					  *)
    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.vVar := NEW(REF BOOLS, top.NV);
    erg.edgeRelevant := NEW(REF BOOLS, top.NE);
    erg.faces := CollectFacesRing(top); 
    erg.drf := ComputeFaceRingDegree(erg.faces,top);
    erg.K := 1.0d0;    
    erg.cosine := NEW(REF ARRAY OF Cosines, top.NE);
    erg.eDcosine := NEW(REF ARRAY OF Cosines, top.NE);

    FOR i := 0 TO top.NE-1 DO
      erg.cosine[i] := NEW(REF LONGS, erg.drf^[i])
    END;

    FOR i := 0 TO top.NE-1 DO
      erg.eDcosine[i] := NEW(REF LONGS, erg.drf^[i])
    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.NE-1 DO erg.edgeRelevant[i] := FALSE END;
  END DefTop;

  PROCEDURE CollectFaces(
     READONLY e: Edge;
     READONLY top: Topology;
     ):
     REF FACES =
     VAR
       NT: CARDINAL := 0;
       ct : CARDINAL;
     BEGIN
       WITH
         NF = top.NF,
         t = NEW(REF ARRAY OF Face, NF)^
       DO
         FOR i := 0 TO NF-1 DO
           ct := 0;
           WITH
             f = top.face[i],
             fun = f.vertex[0].num,
             fvn = f.vertex[1].num,
             fwn = f.vertex[2].num,
             eun = e.vertex[0].num,
             evn = e.vertex[1].num
           DO
             IF (fun = eun OR fun = evn) THEN INC(ct) END;
             IF (fvn = eun OR fvn = evn) THEN INC(ct) END;
             IF (fwn = eun OR fwn = evn) THEN INC(ct) END;
             IF ct = 2 THEN
               t[NT] := f;
               INC(NT);
             END;
           END; 
         END;
         WITH r = NEW(REF ARRAY OF Face, NT) DO
           r^ := SUBARRAY(t,0,NT);
           RETURN r;
         END;
       END;
     END CollectFaces;

  PROCEDURE ComputeFaceRingDegree(READONLY faces: REF ARRAY OF Faces;
                                  READONLY top: Topology) : DRF =
    BEGIN
      WITH
        drf = NEW(DRF, top.NE)
      DO
        FOR l := 0 TO top.NE-1 DO
          WITH
            fie = faces[l]
          DO
            drf^[l] :=  NUMBER(fie^);
          END
        END;
        RETURN drf;
      END
    END ComputeFaceRingDegree;


  PROCEDURE CollectFacesRing(READONLY top: Topology) : REF ARRAY OF Faces =
    BEGIN
      WITH
        faces = NEW(REF ARRAY OF Faces, top.NE)
      DO
        FOR l := 0 TO top.NE-1 DO
          WITH
            e = top.edge[l],
            fie = CollectFaces(e,top)
          DO
            faces[l] :=  fie;
          END
        END;
        RETURN faces;
      END
    END CollectFacesRing;

PROCEDURE DefVar(erg: T; READONLY variable: BOOLS) =
  VAR n : CARDINAL := 0;
  BEGIN
    (*
      This energy is computed only for edges that exist, have DegreeRingFacets
      three or more, and have only existing faces and vertices incident to
      them.
    *)
    WITH
      NV = erg.top.NV,
      NE = erg.top.NE,
      vVar = erg.vVar^,
      edge = erg.top.edge^,
      edgeRelevant = erg.edgeRelevant^,
      drf = erg.drf^,
      faces = erg.faces^
    DO
      <* ASSERT NUMBER(variable) = NV *>
      vVar := variable;
      
      (* Find the relevant edges: *)
      FOR i := 0 TO NE-1 DO edgeRelevant[i] := FALSE; END;
      
      VAR
        fexists,pexists,interface,interpoly: BOOLEAN;
      BEGIN
        FOR i := 0 TO NE-1 DO  
          WITH
            e = edge[i],
            u = e.vertex[0],
            v = e.vertex[1],
            vvar = vVar[u.num] OR vVar[v.num],
            de = drf[e.num],
            cf = faces[e.num]
          DO
            FOR j := 0 TO de-1 DO
              WITH
                f1 = cf^[j],
	        a = f1.pa^,
	        aPpos = Ppos(a),
	        aPneg = Pneg(a)
              DO
	        IF( aPpos #  NIL) AND (aPneg # NIL)
	          THEN
	            interpoly := TRUE;
	          ELSE  
	            interpoly := FALSE;
		END;
                IF f1.exists
	          THEN
                    interface := TRUE;
		  ELSE
		    interface := FALSE;
		END;    
		fexists := f1.exists AND pexists;
	      END;
	      fexists := TRUE AND interface;
	      pexists := TRUE AND interpoly;
	    END; 	         
	                  
            IF e.exists AND fexists AND pexists AND vvar AND ( de >= 3 ) THEN
              WITH
                u  = NARROW(u, Triangulation.Vertex),
	        v = NARROW(v, Triangulation.Vertex)
              DO
                <* ASSERT u.exists AND v.exists *>
              END;
              edgeRelevant[i] := TRUE;
	      INC(n);
            ELSE
              edgeRelevant[i] := FALSE;
            END
          END 
        END
      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,
      NE = erg.top.NE,
      edge = erg.top.edge^,
      edgeRelevant = erg.edgeRelevant^,
      cosine = erg.cosine^,
      K = erg.K,
      drf = erg.drf^,
      faces = erg.faces^,
      eDcosine = erg.eDcosine^,
      vVar = erg.vVar
    DO

      PROCEDURE Compute_cosine(
         READONLY f1,f2: Pair;
         VAR cos: LONGREAL;
	 ) =
        (*
          Compute the diedral angle between the faces f1 and f2. 
        *)
        BEGIN
          WITH
            a = f1,
            b = f2,
	    ao = OrgV(a).num,
	    ad = OrgV(Enext(a)).num,
	    ae = OrgV(Enext_1(a)).num,
	    bo = OrgV(b).num,
	    bd = OrgV(Enext(b)).num,
	    be = OrgV(Enext_1(b)).num,
	    v1a = LR4.Sub(c[ad],c[ao]), 
	    v2a = LR4.Sub(c[ae],c[ao]),
	    v1b = LR4.Sub(c[bd],c[bo]),
	    v2b = LR4.Sub(c[be],c[bo]),
            p1 = FindOrthogonal(v1a,v2a),
	    p2 = FindOrthogonal(v1b,v2b),
            m = LR4.Norm(p1),
            n = LR4.Norm(p2),
            o = LR4.Dot(p1,p2),
            d = m * n
	  DO
            <* ASSERT (ao = bo) AND (ad = bd) *>
            cos := o/d;
	    (*Wr.PutText(stderr, Fmt.LongReal(cos) & "\n");*)
	  END  
        END Compute_cosine;

      PROCEDURE FindOrthogonal(READONLY v1,v2: LR4.T) : LR4.T =
       (* compute a orthogonal vector to v1, by the Gram-Schmidt
          decomposition.
       *)
       BEGIN
	 WITH
	   u1 = v1,
	   v2_u1 = LR4.Project(v2,u1),
           u2 = LR4.Sub(v2,v2_u1)
	 DO
           RETURN u2;
	 END  
       END FindOrthogonal;
	
      PROCEDURE Accum_e_from_cosine(
        cos: LONGREAL;
        d: CARDINAL;
        VAR eDcos: LONGREAL;
        ) =
        (*
          Adds to "e" the energy term corresponding to dihedral average
          angle "a".
        *)
        BEGIN
          WITH
	    thetan = 2.0d0 * FLOAT(Math.Pi,LONGREAL),
	    theta  = thetan/FLOAT(d, LONGREAL),
            cosideal = Math.cos(theta),
	    sinideal = Math.sin(theta),
            cos2 = cos * cos,
	    sqr = Math.sqrt(1.0d0-cos2),
	    first = cosideal * cos,
	    secon = sinideal * sqr,
	    nund = sinideal * cos,
	    firstd = nund/sqr 
          DO
            e := e + K * ( 1.0d0 - 2.0d0 * (first+secon) );
            IF grad THEN
              eDcos:= 2.0d0 * (firstd - cosideal);
            ELSE
              eDcos := 0.0d0;
            END
          END
        END Accum_e_from_cosine;

      PROCEDURE Distribute_eDcosine(
        READONLY f1,f2: Pair;
        READONLY eDcos: LONGREAL;
        ) =
        VAR
          eDv1a,eDv2a,eDv2b,eDp1,eDp2: LR4.T;
        BEGIN
          WITH
            a = f1,
            b = f2,
	    ao = OrgV(a).num,
	    ad = OrgV(Enext(a)).num,
	    ae = OrgV(Enext_1(a)).num,
	    be = OrgV(Enext_1(b)).num,
            v1a = LR4.Sub(c[ad],c[ao]), 
	    v2a = LR4.Sub(c[ae],c[ao]),
	    v2b = LR4.Sub(c[be],c[ao]),
            p1 = FindOrthogonal(v1a,v2a),
	    p2 = FindOrthogonal(v1a,v2b),
            m = LR4.Norm(p1),
            n = LR4.Norm(p2),
            o = LR4.Dot(p1,p2),
            d = m * n,    
            q = o/d,
            f1 = LR4.Cos(v2a,v1a),
            f2 = LR4.Cos(v2b,v1a),
            eDo = eDcos/d,
            eDd = - eDcos * q / d,
            eDm = eDd * n,
            eDn = eDd * m
          DO
            eDp1 := LR4.Mix(eDo, p2, eDm/m, p1);
            eDp2 := LR4.Mix(eDo, p1, eDn/n, p2);
            eDv1a := LR4.Mix(-f1,eDp1,-f2,eDp2);
            eDv2a := LR4.Neg(eDp1);
            eDv2b := LR4.Neg(eDp2);

            IF grad THEN
              IF vVar[ao] THEN
                eDc[ao]:= LR4.Sub(eDc[ao],LR4.Add(LR4.Add(eDv1a,eDv2a),eDv2b));
              END;
              IF vVar[ad] THEN
                eDc[ad] := LR4.Add(eDc[ad],eDv1a);
              END;
              IF vVar[ae] THEN
                eDc[ae] := LR4.Add(eDc[ae],eDv2a);
              END;
              IF vVar[be] THEN
                eDc[be] := LR4.Add(eDc[be],eDv2b);
              END
            END 
          END
        END Distribute_eDcosine;

      BEGIN 
        (* Clear dihedral angle and their derivative accumulators: *)
        FOR j := 0 TO NE-1 DO
          FOR i := 0 TO drf[j]-1 DO 
            cosine[j,i] := 0.0d0;
            eDcosine[j,i] := 0.0d0;
          END
        END;

        FOR i:= 0 TO NV-1 DO
          eDc[i] := LR4.T{0.0d0,0.0d0,0.0d0,0.0d0};
        END;
  
        (* Enumerate edges and accumulate diedral angles: *)
        FOR j := 0 TO NE-1 DO       
          IF edgeRelevant[j] THEN
            WITH
              e = edge[j],
              eun = e.vertex[0].num,
              evn = e.vertex[1].num,
              fie = faces[e.num]
            DO
	      VAR ao,a: Pair;
	          i: CARDINAL:= 0;
	      BEGIN
	        ao := fie[0].pa^;
		a := ao;
	        REPEAT
                  WITH
                   a00 = OrgV(a).num,
                   a11 = OrgV(Clock(a)).num,
                   b = Fnext(a)
                  DO
                    <* ASSERT ((a00 = eun) OR (a00 = evn)) AND
                              ((a11 = eun) OR (a11 = evn))
	            *>
                    Compute_cosine(a,b,cosine[j,i]);
                    INC(i);
                    a := b;
		  END;
		UNTIL ( a = ao )    
              END
            END   
          END            
        END;

        (* Compute energy "e" from dihedral angles, and the gradient
           "eDda": *)

        e := 0.0d0;
        FOR j := 0 TO NE-1 DO
          WITH
            ee = edge[j],
            fie = faces[ee.num],
            de = drf[ee.num]
          DO
            IF edgeRelevant[j] THEN
              FOR i := 0 TO de-1 DO
                WITH
                  f1 = fie[i],
                  af1 = f1.pa^,
                  af2 = Fnext(af1)
                DO
                  Accum_e_from_cosine(cosine[j,i],de,eDcosine[j,i]);
                  IF grad THEN
                    Distribute_eDcosine(af1,af2,eDcosine[j,i]);
                  END
                END
              END
            END
          END
        END
      END
    END
  END Eval;

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

BEGIN
END EquaAngleEnergy.