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

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

REVEAL
  T == Public BRANDED OBJECT
      ElemTableRec_t *top;
      double K;                     /* Energy normalization factor */
      vVar: REF ARRAY OF bool_t;      /* Tells which nodes are variable */
      termVar: REF ARRAY OF bool_t;   /* Tells which terms are variable */
      deg: REF uint_vec_t;      /* Effective degree of each node */
      org, dst: REF uint_vec_t; /* Directed @{edge->?}s of computed terms */
      uint nP;                    /* Number of @{edge->?}s in "org", "dst" */
      sum: REF r4_vec_t;         /* (Work) Sum of effective neighbors */
      eDsum: REF r4_vec_t;       /* (Work) Gradient of "e" rel "sum" */
    OVERRIDES
      init = Init;
      defTop = DefTop;
      defVar = DefVar;
      eval = Eval;
      name = Name;
    }
    
CONST
  Zero == (r4_t){0.0, 0.0, 0.0, 0.0};
    
SELF_Init(SELF_T *erg)
  { return erg;
  } /* END Init */

void SELF_DefTop(SELF_T *erg, ElemTableRec_t *top)
  { erg->top = top;
    erg->K = ((double)top->node.nel);
    erg->vVar = NEW(REF ARRAY OF bool_t, top->node.nel);
    erg->termVar = NEW(REF ARRAY OF bool_t, top->node.nel);
    erg->deg = NEW(REF uint_vec_t , top->node.nel);
    erg->org = NEW(REF uint_vec_t , 2*top->NE);
    erg->dst = NEW(REF uint_vec_t , 2*top->NE);
    erg->sum = r4_vec_new(top->node.nel);
    erg->eDsum = r4_vec_new(top->node.nel);
    /* In case the client forgets to call "defVar": */
    erg->NP = 0;
    for (i = 0; i < top->node.nel; i++)
    {
      erg->vVar.el[i] = FALSE; erg->termVar[i] = FALSE
    }
  } /* END DefTop */
  
void SELF_DefVar(SELF_T *erg, *variable: ARRAY OF bool_t)
  { /*
      Decide which terms are variable, and compute 
      effective degrees "erg->deg[v]".

      There is a term for each existing node. 
      The term of "v" is variable if "v" is variable
      or has some effective neighbor that is variable.
    */
    ??? NP = erg->NP;
    int NV = erg->top->node.nel;
    ??? NE = erg->top->NE;
    ??? @{edge->?} = erg->top->edge^;
    bool_vec_t vVar = erg->vVar^;
    ??? termVar = erg->termVar^;
    ??? org = erg->org^;
    ??? dst = erg->dst^;
    ??? deg = erg->deg^;
    {
      assert((variable.nel) == NV);
      vVar = variable;

      /* If a node is vVar, its term is vVar: */
      for (v = 0; v < NV; v++)
      {
        termVar[v] = vVar.el[v]; 
        deg[v] = 0
      }    

      /* Enumerate effective @{edge->?}s and accumulate efective degrees. */
      /* Also set "termVar[v]" if "v" has a vVar effective neighbor. */
      for (i = 0; i < NE; i++)
      {
        ??? e = @{edge->?}[i];
        ??? u = NARROW(OrgV(e.pa), Triangulation.Node);
        ??? v = NARROW(OrgV(Clock(e.pa)), Triangulation.Node);
        ??? un = u->num;
        ??? vn = v->num;
        {
          INC(deg[un]); INC(deg[vn]);
          termVar[un] = termVar[un]) || (vVar.el[vn];
          termVar[vn] = termVar[vn]) || (vVar.el[un];
        }
      }

      /* A node with no effective neighbors (deg[v]==0) has energy zero
         by definition (termVar[v]==FALSE): */
      for (v = 0; v < NV; v++)
      {
        if (deg[v] == 0){ termVar[v] = FALSE;}
      }    

      /* Now collect all efective directed @{edge->?}s "(org[k], dst[k])" */
      /* whose origin "org[k]" is variable: */ 
      NP = 0;
      for (i = 0; i < NE; i++)
      {
        ??? e = @{edge->?}[i];
        Node_t u = OrgV(e.pa);
        ??? un = u->num;
        Node_t v = OrgV(Clock(e.pa));
        ??? vn == v->num;
        {
          if (termVar[un]){ org[NP] = un; dst[NP] = vn; NP++; }
          if (termVar[vn]){ org[NP] = vn; dst[NP] = un; NP++; }
        }
      }
    }
  } /* END DefVar */
  
void SELF_Eval(
    SELF_T erg; 
    Coords_t *c; 
    double *e; 
    bool_t grad;
    Gradient *eDc;
 )
  { ??? NP = erg->NP;
    int NV = erg->top->node.nel;
    bool_vec_t vVar = erg->vVar^;
    ??? termVar = erg->termVar^;
    ??? deg = erg->deg^;
    ??? org = erg->org^;
    ??? dst = erg->dst^;
    ??? sum = erg->sum^;
    ??? eDsum = erg->eDsum^;
    ??? K = erg->K;
    {

      void Compute_sum ()
        /* Computes "sum[v]" for all relevant terms "v" */
        {
          for (v = 0; v < NV; v++)
          {
            if (termVar[v]){ sum[v] = Zero;}
          }    
          for (i = 0; i < NP; i++)
          {
            ??? v = org[i], u == dst[i];
    {
              assert(termVar[v]);
              sum[v] = r4_Add(sum[v], c[u])
            }
          } 
        } Compute_sum;

      void EvalTerm(
          *cv, sumv: r4_t; 
          uint degv;
          double *ev; 
          VAR evDcv, evDsumv: r4_t;
       )
        /*
          Computes the energy term "ev" for a node at "cv",
          given the sum of its "relevant neighbors" "sumv".
          Returns also the gradients "evDcv" and "evDsumv" of
          "ev" relative to "cv" and "sumv".
        */
        {
          ??? f = 1.0/((double)degv);
          ??? barv = r4_Scale(f, sumv);
          ??? rv = r4_Sub(cv, barv);
          {
            ev = K * r4_NormSqr(rv);
            if (grad)
            {
              evDcv = r4_Scale(2.0 * K, rv);
              evDsumv = r4_Scale(-2.0 * f * K, rv)
            }
            else
            {
              evDcv = Zero;
              evDsumv = Zero
            }
          }
        } EvalTerm;

      void Distribute_eDsum ()
        /*
          Adds to the derivatives "eDc[u]" the indirect terms
          represented by "eDsum[v]", for all computed terms "v" 
          that have "u" as a relevant neighbor. 
        */
        {
          for (i = 0; i < NP; i++)
          {
            ??? v = org[i], u == dst[i];
    {
              if (vVar.el[u])
              {
                eDc[u] = r4_Add(eDc[u], eDsum[v])
              }
            }
          }  
        } Distribute_eDsum;

      double *ev;
      {
        Compute_sum();
        
        e = 0.0;
        for (v = 0; v < NV; v++)
        {
          if (termVar[v])
          {
            EvalTerm(c[v], sum[v], deg[v], ev, eDc[v], eDsum[v]);
            e = e + ev
          }
          else
          {
            eDc[v] = Zero;
          }
        }

        if (grad){ Distribute_eDsum();}
      }
    }
  } /* END Eval */

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

{
} ExcenEnergy.