/* See {VarWindingEnergy.h} */
#include <VarWindingEnergy.h>

#define _GNU_SOURCE
#include <stdio.h>
#include <math.h>
#include <r3.h>
#include <r4.h>
#include <Triangulation.h>


#include <Octf.h>
#include <Triangulation.h>
#include <Energy.h>



TYPE
double bool_t = bool_t;
double bool_vec_t = ARRAY OF bool_t;
double NAT = uint;
double NATS = uint_vec_t;
  
REVEAL
double T = Public BRANDED OBJECT
      double K;              /* The energy normalization factor */
      ElemTableRec_t *top;     /* The map elements. */
      termOK: REF bool_vec_t;        /* The terms that should be included in the energy. */
      term: REF Terms;          /* Nodes in each term. */
      rDpw, sDpw: REF Coords3D_t; /* Work areas for "Eval". */
    OVERRIDES
      init = Init;
      defTop = DefTop;
      defVar = DefVar;
      eval = Eval;
      name = Name;
    }
    
TYPE
double Term = RECORD  /* Node numbers included in some energy term. */
      NAT un, vn;   /* Endpoints of some @{edge->?}. */
      wn: REF NATS;  /* Tips of walls incident to that @{edge->?}. */
      ex: REF bool_vec_t; /* "ex[i]" == cell between wn[i] and wn[i-1] exists. */
    }
double Terms = ARRAY OF Term;

CONST Debug == FALSE;

SELF_Init(SELF_T *erg)
  { return erg;
  } /* END Init */

void SELF_DefTop(SELF_T *erg, ElemTableRec_t *top)
  { erg->K = 1.0/((double)top->NE);      
    erg->top = top;
    erg->termOK = bool_vec_new(top->NE); 
    erg->term = CollectTerms(top);
    /* Size of "rDpw", "sDpw" is maximum @{edge->?} degree: */
    ??? maxDeg = Max@{Edge->?}Degree(erg->term^);
    {
      erg->rDpw = NEW(REF Coords3D_t, maxDeg);
      erg->sDpw = NEW(REF Coords3D_t, maxDeg);
    }
    /* Just in case the client forgets to call "defVar": */
    for (i = 0; i < top->NE; i++){ erg->termOK[i] = FALSE; }
  } /* END DefTop */
  
Terms *CollectTerms(ElemTableRec_t *top)==
  /* Collects relevant nodes for each @{edge->?}. */
  Place_t *b; k: NAT;
  { ??? NE = top->NE;
    ??? termr = NEW(REF Terms, NE), term == termr^;
    {
      for (i = 0; i < NE; i++)
      {
        ??? e = top->edge[i];
        ??? a = e.pa;
        ??? deg = DegreeOfWall(a);
        ??? wnr = NEW(REF NATS, deg), wn == wnr^;
        ??? exr = bool_vec_new(deg), ex == exr^;
        Node_t un = OrgV(a)->num;
        Node_t vn = OrgV(NextE(a))->num;
        {
          b = a; k = 0;
          do 
          {
            wn[k] = OrgV(PrevE(b))->num;
            ex[k] = PnegP(b)!=NULL;
            b = NextF(b); k++
          } while (b != a);
          assert(k == deg);
          term[i] = Term{un = un, vn = vn, wn = wnr, ex = exr}
        }
      }
      return termr;
    }
  } /* END CollectTerms */
  
PROCEDURE Max@{Edge->?}Degree(*term: Terms): NAT ==
      maxDeg: NAT = 0;
  { for (i = 0 TO (term.nel-1)){
      maxDeg = MAX(maxDeg, ((term[i].wn^).nel))
    }
    return maxDeg;
  } /* END Max@{Edge->?}Degree */
  
void SELF_DefVar(SELF_T *erg, bool_vec_t *variable)

  { /*
      Decide which @{edge->?}s are relevant to "VarWinding" energy.
      A @{edge->?} is relevant iff its star includes at least one 
      variable node.*/
    ??? termOK = erg->termOK^;
    ??? @{edge->?} = erg->top->edge^;
    {
      assert((variable.nel) == erg->top->node.nel);
      fprintf(stderr, "VarWindingEnergy: relevant @{edge->?}s == { ");
      for (i = 0 TO (termOK.nel-1))
      {
        ??? e = @{edge->?}[i];
        {
          assert(e->num == i); 
          termOK[i] = e->exists) && (@{Edge->?}StarIsVariable(e.pa, variable);
          if (termOK[i]){ fprintf(stderr, " " & Fmt.Int(i)); }
        }
      }
      fprintf(stderr, " }\n");
    }
  } /* END DefVar */
  
PROCEDURE @{Edge->?}StarIsVariable(Place_t @p, bool_vec_t *variable): bool_t ==
  Place_t *b;
  { Node_t u = OrgV(a);
    Node_t v = OrgV(NextE(a));
    {
      if ((variable[u->num]) || (variable[v->num]))
      {
        return TRUE;
      }
      else
      {
        b = a;
        do 
        {
          Node_t w = OrgV(PrevE(b));
    {
            if (variable[w->num]){ return TRUE;}
          }
          b = NextF(b)
        } while (b != a);
        return FALSE;
      }
    }
  } /* END @{Edge->?}StarIsVariable */

void SELF_Eval(
    SELF_T erg; 
    Coords_t *c; 
    double *e; 
    bool_t grad;
    Gradient *eDc;
 )
  { int NV = erg->top->node.nel;
    ??? NE = erg->top->NE;
    ??? termOK = erg->termOK^;
    ??? term = erg->term^;
    ??? K = erg->K;
    ??? rDpw = erg->rDpw^;
    ??? sDpw = erg->sDpw^;
    {
      e = 0.0;
      /* Clear gradient accumulators */
      for (i = 0; i < NV; i++){ eDc[i] = (r4_t){0.0, ..}; }

      /* Enumerate @{edge->?}s and compute @{edge->?} sums: */
      for (i = 0; i < NE; i++)
      {
        if (termOK[i])
        {
          AddTerm(term[i], c, K, grad, e, eDc, rDpw, sDpw)
        }
      }
    }
  } /* END Eval */

void AddTerm(
    Term *term;
    Coords_t *c;
    double K;
    bool_t grad;
    VAR /*IO*/ e: double;
    VAR /*IO*/ eDc: Gradient;
    VAR rDpw, sDpw: Coords3D_t;  /* Work areas. */
 )
  /* 
    Adds one term of the energy to "e".  If "grad" is set,
    also adds the gradient of that term to "eDc". */
      
    double r, s;          /* Radial and perimetral sums for term. */
    rDpu, rDpv: r3_t;       /* Derivatives of "r" wrt @{edge->?} endpoints. */
    sDpu, sDpv: r3_t;       /* Derivatives of "s" wrt @{edge->?} endpoints. */
    double t;             /* An energy term, corresponding to a single @{edge->?}. */
    double tDr, tDs;      /* Derivatives of "t" w.r.t. "r" and "s". */
    NAT gaps;               /* Degree of "e", and num of missing cells. */
    double SNorm, RNorm;  /* Normalization factors for "r" and "s". */
  CONST
  double Epsilon = 1.0d-20;               /* To avoid divide by zero. */
  double Pi = M_PI;       /* Pi */
  double HfPi = 1.5707963267948966;     /* Pi/2 */
  { ??? un = term.un;
    ??? pu == SUBARRAY(c[un], 0, 3),
    ??? vn = term.vn,  pv == SUBARRAY(c[vn], 0, 3);
    ??? h = r3_sub(pv, pu);
    ??? h2 = LR3.Dot(h, h);
    double hm = sqrt(h2);
    ??? wn = term.wn^;
    ??? ex = term.ex^;
    double N = FLOAT((wn.nel), double);
    {

      void Compute_r_s_Perimeter()
        /*
          Returns in "r" and "s" the sums corresponding to the given
          term of the energy ("Perimeter" formula). Also, if "grad" is
          set, computes the gradients "rDpu", "rDPv", "rDpw[k]" of "r"
          relative to the 3D coordinates of relevant nodes, and
          similarly for "s".
          
          Also defines the geometry-independent constants "RNorm" and
          "SNorm", such that a regular @{edge->?} star of height "h" and
          radius "r" has "r == RNorm·h^2·b^2" and "s == SNorm·h^2·b^2". */
        uint *kx, xn; px: r3_t; gaps: NAT;
        {
          r = Epsilon; s = Epsilon;
          gaps = 0;
          if (grad)
          {
            rDpu = (r3_t){0.0,..}; rDpv = (r3_t){0.0,..};
            sDpu = (r3_t){0.0,..}; sDpv = (r3_t){0.0,..};
            for (i = 0 TO (wn.nel-1))
            {
              rDpw[i] = (r3_t){0.0,..};
              sDpw[i] = (r3_t){0.0,..}
            }
          } 
          kx = (wn.nel-1); 
          xn = wn[kx];
          px = SUBARRAY(c[xn], 0, 3);
          for (ky = 0 TO (wn.nel-1))
          {
            ??? yn = wn[ky];
            ??? py = SUBARRAY(c[yn], 0, 3);
            double fy = r3_sub(py, pu);
            double ry = LR3Extras.Cross(h, fy);
            double rb = LR3.Dot(ry,ry);
              
            double gy = r3_sub(py, px);
            double sy = LR3Extras.Cross(h, gy);
            double sb = LR3.Dot(sy,sy)
            {
              r = r + rb;
              if (ex[ky]){ s = s + sb }else{ gaps++; }
              if (grad)
              {
                ??? rbDry = r3_scale(2.0, ry);
                double rbDfy = LR3Extras.Cross(rbDry, h);
                double rbDh = LR3Extras.Cross(fy, rbDry)
                {
                  if (Debug)
                  {
                    PrL("rb == ", rb); PrEOL();
                    PrV("rbDfy == ", rbDfy); PrEOL();
                    PrV("rbDh == ", rbDh);  PrEOL();
                  }
                  
                  rDpw[ky] = r3_add(rDpw[ky], rbDfy);
                  rDpu = r3_sub(rDpu, r3_add(rbDfy, rbDh));
                  rDpv = r3_add(rDpv, rbDh);
                  if (ex[ky])
                  {
                    ??? sbDsy = r3_scale(2.0, sy);
                    double sbDgy = LR3Extras.Cross(sbDsy, h);
                    double sbDh = LR3Extras.Cross(gy, sbDsy);
                    {                    
                      if (Debug)
                      {
                        PrL("sb == ", sb); PrEOL();
                        PrV("sbDgy == ", sbDgy); PrEOL();
                        PrV("sbDh == ", sbDh);  PrEOL();
                      }
                      sDpw[ky] = r3_add(sDpw[ky], sbDgy);
                      sDpw[kx] = r3_sub(sDpw[kx], sbDgy);
                      sDpu = r3_sub(sDpu, sbDh);
                      sDpv = r3_add(sDpv, sbDh);
                    }
                  }
                }
              }
             kx = ky; xn = yn; px = py; 
            }
          }
          if (Debug)
          {
            PrL("r == ", r); PrEOL();
            PrV("rDpu == ", rDpu); PrEOL();
            PrV("rDpv == ", rDpv);  PrEOL();
            PrL("s == ", s); PrEOL();
            PrV("sDpu == ", sDpu); PrEOL();
            PrV("sDpv == ", sDpv);  PrEOL();  PrEOL();
          }
          RNorm = N;
          if (gaps == 0)
          {
            /* @{Edge->?} is interior; ideally, winding number should be ±1 */
            ??? sn = Math.sin(Pi/N);
    {
              SNorm = 4.0*N*sn*sn
            }
          }
          else
          {
            assert(gaps < (wn.nel));
            /* @{Edge->?} is on border; winding number should be ±1/2, with fencepost correction */
            ??? G = ((double)gaps);
              ??? sn = Math.sin(HfPi/(N-G));
            {
              SNorm = 4.0*(N-G)*sn*sn
            }
          }
        } Compute_r_s_Perimeter;

      void Compute_r_s_Area()
        /*
          Returns in "r" and "s" the sums corresponding to the given
          term of the energy ("Area" formula). Also, if "grad" is
          set, computes the gradients "rDpu", "rDPv", "rDpw[k]" of "r"
          relative to the 3D coordinates of relevant nodes, and
          similarly for "s".
          
          Also defines the geometry-independent constants "RNorm" and
          "SNorm", such that a regular @{edge->?} star of height "h" and
          radius "r" has "r == RNorm·h^2·b^2" and "s == SNorm·h^2·b^2". */
        uint *kx, xn; px: r3_t; gaps: NAT;
        {
          r = Epsilon; s = Epsilon;
          gaps = 0;
          if (grad)
          {
            rDpu = (r3_t){0.0,..}; rDpv = (r3_t){0.0,..};
            sDpu = (r3_t){0.0,..}; sDpv = (r3_t){0.0,..};
            for (i = 0 TO (wn.nel-1))
            {
              rDpw[i] = (r3_t){0.0,..};
              sDpw[i] = (r3_t){0.0,..}
            }
          } 
          kx = (wn.nel-1); 
          xn = wn[kx];
          px = SUBARRAY(c[xn], 0, 3);
          for (ky = 0 TO (wn.nel-1))
          {
            ??? yn = wn[ky];
            ??? py = SUBARRAY(c[yn], 0, 3);
            double fy = r3_sub(py, pu);
            double ry = LR3Extras.Cross(h, fy);
            double rb = LR3.Dot(ry,ry);
              
            double dy = LR3Extras.Cross(fx,fy);
            double sb = LR3.Dot(h, dy);
            {
              r = r + rb;
              if (ex[ky]){ s = s + sb }else{ gaps++; }
              if (grad)
              {
                ??? rbDry = r3_scale(2.0, ry);
                double rbDfy = LR3Extras.Cross(rbDry, h);
                double rbDh = LR3Extras.Cross(fy, rbDry);
                {
                  if (Debug){ PrL("rb == ", rb); PrEOL(); PrV("rbDfy == ", rbDfy); PrEOL(); PrV("rbDh == ", rbDh);  PrEOL(); }
                  rDpw[ky] = r3_add(rDpw[ky], rbDfy);
                  rDpu = r3_sub(rDpu, r3_add(rbDfy, rbDh));
                  rDpv = r3_add(rDpv, rbDh);
                  if (ex[ky])
                  {
                    ??? sbDdy = fx;
                    ??? sbDry = LR3Extras.Cross(h, sbDdy);
                    double sbDfy = - sbDry * h ×;
                    double sbDfx =;
                    double sbDgy = LR3Extras.Cross(sbDsy, h);
                    double sbDh = LR3Extras.Cross(gy, sbDsy)
                    {                    
                      if (Debug)
                      {
                        PrL("sb == ", sb); PrEOL();
                        PrV("sbDgy == ", sbDgy); PrEOL();
                        PrV("sbDh == ", sbDh);  PrEOL();
                      }
                      sDpw[ky] = r3_add(sDpw[ky], sbDgy);
                      sDpw[kx] = r3_sub(sDpw[kx], sbDgy);
                      sDpu = r3_sub(sDpu, sbDh);
                      sDpv = r3_add(sDpv, sbDh);
                    }
                  }
                }
              }
             kx = ky; xn = yn; px = py; 
            }
          }
          if (Debug)
          {
            PrL("r == ", r); PrEOL();
            PrV("rDpu == ", rDpu); PrEOL();
            PrV("rDpv == ", rDpv);  PrEOL();
            PrL("s == ", s); PrEOL();
            PrV("sDpu == ", sDpu); PrEOL();
            PrV("sDpv == ", sDpv);  PrEOL();  PrEOL();
          }
        } Compute_r_s_Area;

      void Compute_RNorm_SNorm_Perimeter()
        /* 
          Defines "RNorm,SNorm" which are the
          expected values of "r,s", for a regular @{edge->?} star,
          divided by "h^2·b^2" where "h" is the length of the
          @{edge->?} and "b" is the radius of the star. */
          
        {
          RNorm = N;
          if (gaps == 0)
          {
            /* @{Edge->?} is interior; ideally, winding number should be ±1 */
            ??? sn = Math.sin(Pi/N);
    {
              SNorm = 4.0*N*sn*sn
            }
          }
          else
          {
            assert(gaps < (wn.nel));
            /* @{Edge->?} is on border; winding number should be ±1/2, with fencepost correction */
            ??? G = ((double)gaps);
              ??? sn = Math.sin(HfPi/(N-G));
            {
              SNorm = 4.0*(N-G)*sn*sn
            }
          }
        } Compute_RNorm_SNorm;

      void Compute_t_from_r_s()
        /*
          Returns in "t" the energy term corresponding to an @{edge->?} with
          sums "r" and "s"  Also, if "grad" is true, stores in "tDr"
          and "tDs" the derivatives of the term relative to those sums. */
        {
            ??? R = r/RNorm;
              ??? S = s/SNorm;
            ??? f = R/S;
            ??? g = S/R;
            ??? y = f + g;
            ??? z = y*y;
            {
              t = K * (z - 4.0);
              if (Debug)
              {
                PrL("R == ", R); PrL("S == ", S); PrL("t == ", t); PrEOL();
              }
              if (grad)
              {
                ??? zDy = 2.0 * y;
                ??? zDr = zDy*(1.0/S - g/R)/RNorm;
                double zDs = zDy*(1.0/R - f/S)/SNorm;
                {
                  tDr = K * zDr;
                  tDs = K * zDs
                }
              }
              else
              {
                tDr = 0.0; tDs = 0.0;
              }
            }
          } Compute_t_from_r_s;
     
      void Distribute_tDr_tDs()
        /*
          Accumulates in "eDc" the gradient of an energy term "t" relative to
          the nodes in the star of the corresponding @{edge->?} "Edg(a)", given 
          the derivatives "tDr" and "tDs" of the term relative to the edge
          sums "r" and "s". */

        void AddGradientToNode(vn: NAT; rDpv, sDpv: r3_t)
          /* Adds to "eDc" the contribution due to node "vn". */
          {
            ??? eDcv = eDc[vn];
    {
              eDcv[0] = eDcv[0] + tDr * rDpv[0] + tDs * sDpv[0];
              eDcv[1] = eDcv[1] + tDr * rDpv[1] + tDs * sDpv[1];
              eDcv[2] = eDcv[2] + tDr * rDpv[2] + tDs * sDpv[2]
            }
          } AddGradientToNode;

        {
          AddGradientToNode(un, rDpu, sDpu);
          AddGradientToNode(vn, rDpv, sDpv);
          for (k = 0 TO (wn.nel-1))
          {
            AddGradientToNode(wn[k], rDpw[k], sDpw[k])
          }
        } Distribute_tDr_tDs;
              
      {
        Compute_r_s();
        Compute_t_from_r_s();
        e = e + t;
        if (grad){ Distribute_tDr_tDs();}
      }  
    }            
  } /* END AddTerm */

PROCEDURE Name(<*UNUSED);erg: T) : char *==
  { return "VarWinding()";
  } /* END Name */

void PrI(prefix: char *; x: INTEGER)

  { fprintf(stderr, " ");
    fprintf(stderr, prefix);
    fprintf(stderr, Fmt.Pad(Fmt.Int(x), 12));
  } /* END PrI */

void PrL(prefix: char *; x: double)

  { fprintf(stderr, " ");
    fprintf(stderr, prefix);
    fprintf(stderr, 
      Fmt.Pad(Fmt.LongReal(x, style = Fmt.Style.Fix, prec = 8), 22)
   );
  } /* END PrL */

void PrV(prefix: char *; x: r3_t)

  { fprintf(stderr, " ");
    fprintf(stderr, prefix); 
    for (i = 0; i < 3; i++)
    {
      fprintf(stderr, 
        Fmt.Pad(Fmt.LongReal(x[i], style = Fmt.Style.Fix, prec = 8), 21) & " "
     )
    }
  } /* END PrV */

void PrEOL()

  { fprintf(stderr, "\n");
  } /* END PrEOL */

{
} VarWindingEnergy.