/* See bezier.h */
/* Last edited on 2021-07-18 00:10:35 by jstolfi */

#define _GNU_SOURCE
#include <stdlib.h>
#include <stdint.h>
#include <math.h>

#include <affirm.h>
#include <interval.h>
#include <box.h>
#include <bool.h>
#include <jsmath.h>
#include <filefmt.h>
#include <double_array.h>
#include <indexing_descr.h>

#include <bz_basic.h>
#include <psp_basic.h>
#include <psp_grid.h>

#include <bz_patch.h>

void bz_patch_test_map(ix_dim_t d, const double x[], double z[]);
  /* Maps a point {x[0..d-1]} of the domain to a point {z[0..1]} of {R^2}.
    Uses only coordinates {x[0..k-1]} where {k = min(d,2)}. */ 

double bz_patch_test_func ( ix_dim_t d, const double x[], ix_axis_t i );
  /* Component {i} of the dependent function of the domain point {x[0..d-1]}. */

void bz_patch_eval_m
  ( bz_patch_ddim_t m,      /* Dimension of parameter space. */
    bz_patch_rinds_t r,     /* Dimension of image space. */
    const bz_degree_t g[],  /* Degrees on each coordinate. */
    double_array_t *C,      /* Bezier control points. */
    double x[],             /* Argument vector (size {m}). */
    double_array_t *FX      /* OUT: Result array (size {r}). */
  );

/* IMPLEMENTATIONS */

const ix_size_t *bz_patch_get_dsize(bz_patch_t *P)
  { assert(P->C.ds.na == P->d + P->r);
    return P->C.ds.sz;
  }

const ix_size_t *bz_patch_get_rsize ( bz_patch_t *P )
  { assert(P->C.ds.na == P->d + P->r);
    return &(P->C.ds.sz[P->d]);
  }
  
void bz_patch_get_degrees ( bz_patch_t *P, bz_degree_t g[] )
  { assert(P->C.ds.na == P->d + P->r);
    ix_size_t *bsz = P->C.ds.sz;
    int i;
    for (i = 0; i < P->d; i ++) { assert(bsz[i] > 0); g[i] = bsz[i] - 1; }
  }

bz_patch_t bz_patch_new(bz_patch_ddim_t d, const bz_degree_t g[], bz_patch_rinds_t r, const ix_size_t rsz[])
  { /* fprintf(stderr, "+ bz_patch_new\n");  */
    /* Build the size array: */
    demand(d + r <= double_array_MAX_AXES, "too many indices in domain+range");
    ix_dim_t na = d + r;
    /* Compute total size {tsz} of control array: */
    ix_size_t bsz[na];
    { int i; for (i = 0; i < na; i++) { bsz[i] = (i < d ? g[i] + 1 : (rsz == NULL ? 1 : rsz[i-d])); } }
    bz_patch_t P;
    P.d = d;
    P.r = r;
    P.C = double_array_new(na, bsz);
    /* fprintf(stderr, "- bz_patch_new\n");  */ 
    return P;
  }

bz_patch_t bz_patch_uniform_new(bz_patch_ddim_t d, bz_degree_t g, bz_patch_rinds_t r, const ix_size_t rsz[])
  { /* fprintf(stderr, "+ bz_patch_new\n");  */
    /* Build a uniform degree array: */
    bz_degree_t gg[d];
    int i;
    for (i = 0; i < d; i++) { gg[i] = g; }
    bz_patch_t P = bz_patch_new(d, gg, r, rsz);
    /* fprintf(stderr, "- bz_patch_new\n");  */ 
    return P;
  }

bz_patch_t bz_patch_copy ( bz_patch_t *P )
  { 
    assert(P->C.ds.na == P->d + P->r);
    bz_patch_t T;
    T.d = P->d;
    T.r = P->r;
    T.C = double_array_copy(&(P->C));
    assert(T.C.ds.na == T.d + T.r);
    return T;
  }

void bz_patch_free(bz_patch_t *P)
  {
    if (P != NULL) { double_array_free_elems(&(P->C)); }
  }

double_array_t bz_patch_value_new(bz_patch_t *P)
  { assert(P->C.ds.na == P->d + P->r);
    int d = P->d;
    int r = P->r;
    ix_descr_t *D = &(P->C.ds);
    ix_size_t *vsz = &(D->sz[d]);
    return double_array_new(r, vsz);
  }

double_array_t bz_patch_get_coeff(bz_patch_t *P, const bz_patch_cpindex_t e[])
  { assert(P->C.ds.na == P->d + P->r);
    /* Make a copy of the descriptor: */
    double_array_t S = P->C;
    /* Fix the first {P.d} indices to the given values {e[0..P.d-1]}: */
    ix_descr_slice(&(S.ds), P->d, NULL, e);
    return S;
  }

void bz_patch_eval(bz_patch_t *P, const double x[], double_array_t *BX)
  { bz_patch_ddim_t  d = P->d;     /* Domain dimension. */
    bz_patch_rinds_t r = P->r;     /* Indices in each coeff and value. */
    assert(P->C.ds.na == d + r);
    
    /* Work area for the DeCasteljau algorithm: */
    double_array_t W = double_array_copy(&(P->C));
    ix_dim_t na = W.ds.na;
    ix_size_t *Wsz = W.ds.sz;
    ix_step_t *Wst = W.ds.st;
    ix_pos_t Wbp = W.ds.bp;
    assert(na == d + r);
    
    ix_count_t nv = ix_num_tuples(r, &(Wsz[d])); /* Num elems in each coeff. */
       
    auto void eval_d ( int k, ix_pos_t bp );
      /* Evaluates one slice of a patch for argument {x[k..d-1]}. Assumes that
        the slice was obtained by fixing the first {k} indices of {W} to certain
        values. The patch coefficients are in {W.e[]}, starting at base position
        {bp}. They are indexed by the last {na-k} indices of {W}, with size
        vector {Wsz[k..na-1]} and step vector {Wst[k..na-1]}.
        
        The result is left in {W.e[]} starting at position {bp}, indexed by {r}
        indices, with size vector {Wsz[d..na-1]} and step vector
        {Wst[d..na-1]}.*/
        
    auto void eval_1 ( int k, ix_pos_t bp );
      /* Applies the DeCalsteljau algorithm to a list of {Wsz[k]} Bézier coeffs.
        Assumes that the coeff with index {i} is stored in {W.e[]} starting at position 
        {bp + i*Wsz[k]}. Assumes that each coeff is indexed by {r} indices, with size
        vector {Wsz[d..na-1]} and step vector {Wst[d..dr+r-1]}. The result
        is left in {W.e[]} starting at position {bp}, indexed in the same way.
        
        The current implementation assumes that {W} is stored in lexicographic order 
        so that each coeff is a block of {nv} consecutive elements in {W.e[]}. */

    void eval_d ( int k, ix_pos_t bp )
      { /* Recursion on the argument dimension {d-k}: */
        /* fprintf(stderr, "+ bz_patch_eval  d = %d  k = %d", d, k); */
        if (k >= d)
          { /* A Bezier patch with domain {R^0} is just single coeff {R}. */
            /* Nothing to do. */
          }
        else
          { /* Recursively evaluate {sz[k]} patches of domain dimension {d-k-1}, excluding first index: */
            ix_size_t szk = Wsz[k];
            ix_step_t stk = Wst[k];
            int i;
            for (i = 0; i < szk; i++) { eval_d(k+1, bp + i*stk); }
            /* Now do a 1-dimensional DeCasteljau interpolation based on {x[0]}: */
            eval_1(k, bp);
          }
        /* fprintf(stderr, "- bz_patch_eval_d  d = %d  k = %d", d, k); */
      }

    void eval_1 ( int k, ix_pos_t bp )
      { /* The DeCasteljau algorithm. */
        double P = x[k];
        double a = 1.0 - P;
        assert(Wsz[k] > 0);
        bz_degree_t g = Wsz[k] - 1;
        ix_step_t stk = Wst[k];
        int t;
        for (t = g; t > 0; t--)
          { /* Interpolate {t+1} coeffs to get {t}: */
            int s;
            for (s = 0; s < t; s++)
              { /* Interpolate coeffs {s} and {s+1}, store in {s}: */
                double *ea = &(W.e[bp + s*stk]);
                double *eb = &(ea[stk]);
                int i;
                for (i = 0; i < nv; i++) { ea[i] = a*ea[i] + P*eb[i]; }
              }
          }
      }

    eval_d(0, Wbp);
    /* Make a descriptor for the result: */
    ix_descr_t Rd;
    Rd.na = d;
    ix_sizes_assign(r, Rd.sz, &(Wsz[d]));
    ix_steps_assign(r, Rd.st, &(Wst[d]));
    Rd.bp = Wbp;
    double_array_t R = (double_array_t) { .ds = Rd, .e = W.e };
    /* Copy result into {*BX}: */
    double_array_assign(BX, &R);
  }
  
bz_patch_t bz_patch_get_facet ( bz_patch_t *P, ix_axis_t i, interval_side_t dir )
  { assert(P->C.ds.na == P->d + P->r);
    int d = P->d;
    demand((i >= 0) && (i < d), "invalid axis");
    ix_axis_t ax[1] = { i };
    ix_index_t ix[1] = { (dir == interval_BLO ? 0 : (P->C.ds.sz[i] - 1)) };
    bz_patch_t T = (*P);
    ix_descr_slice(&(T.C.ds), 1, ax, ix);
    T.d--;
    return T;
  }

bz_patch_t bz_patch_get_face ( bz_patch_t *P, const box_signed_dir_t dir[] )
  { 
    assert(P->C.ds.na == P->d + P->r);
    int d = P->d;
    
    /* Indices along axes {ax[0..n-1]} are to be fixed to values {ix[0..n-1]}. */
    int n = 0;       
    ix_axis_t ax[d]; 
    ix_index_t ix[d];
    
    ix_axis_t i;
    for (i = 0; i < d; i++)
      { if (dir[i] != box_ZER)
          { ax[n] = i; 
            ix[n] = (dir[i] == box_NEG ? 0 : (P->C.ds.sz[i] - 1));
            n++;
          }
      }
    bz_patch_t T = (*P);
    if (n > 0) { ix_descr_slice(&(T.C.ds), n, ax, ix); T.d = d - n; }
    return T;
  }

bz_degree_t bz_patch_max_degree(bz_patch_t *P)
  { assert(P->C.ds.na == P->d + P->r);
    bz_patch_ddim_t d = P->d;
    ix_size_t *bsz = &(P->C.ds.sz[0]);
    ix_size_t szmax = 1;
    { int i; 
      for (i = 0; i < d; i++) 
        { demand(bsz[i] >= 1, "patch has empty array");
          if (bsz[i] > szmax) { szmax = bsz[i]; }
        }
    }
    return szmax - 1;
  }

void bz_patch_raise_degree(bz_patch_t *P, bz_patch_t *T)
  { /* Assumption: Linear interpolation of the indices? */
    demand(T->d == P->d, "domain index count mismatch");
    demand(T->r == P->r, "range index count mismatch");
    assert(P->C.ds.na == P->d + P->r);
    assert(T->C.ds.na == T->d + T->r);
    ix_dim_t d = P->d;
    ix_size_t *Psz = P->C.ds.sz;
    ix_size_t *Tsz = T->C.ds.sz;
    ix_step_t *Tst = T->C.ds.st;
    
    /* Build a descriptor {RC} for {T.C} cropped to the size of {P.C}: */
    double_array_t RC = T->C;
    int i;
    for (i = 0; i < d; i++) 
      { demand(Psz[i] >= 0, "invalid patch");
        demand(Tsz[i] >= Psz[i], "degree cannot be lowered");
        RC.ds.sz[i] = Psz[i];
      }
    
    /* Copy {P.C} into the part of {T.C} described by {RC}: */
    double_array_assign(&RC, &(P->C));
    
    /* Now raise the degree of {RC} along each domain dimension: */
    for (i = 0; i < d; i++) 
      { 
        bz_degree_t Pgi = Psz[i] - 1;
        bz_degree_t Tgi = Tsz[i] - 1;
        if (Pgi < Tgi)
          { 
            auto bool_t raise_one ( const ix_index_t ix[], ix_pos_t pA, ix_pos_t pB, ix_pos_t pC );
            /* Raises the degree of a univariate, single-valued Bézier 
              patch from degree {Pgi} to {Tgi}. Assumes that the {Pgi+1} coeffs are stored in
              a one-dimensional slice of {RC}, with base position {pA} and
              step {Tst[i]}.  Assumes also that there are {Tgi-Pgi}
              unused elements in {C} at the end of that slice. Leaves the 
              {Tgi+1} coeffs of the raised patch in the same slice. Arguments {pB} and {pC} are
              ignored. Always returns {FALSE}. */

            bool_t raise_one ( const ix_index_t ix[], ix_pos_t pA, ix_pos_t pB, ix_pos_t pC )
              { 
                bz_degree_t g = Pgi;
                ix_step_t st =  Tst[i];
                assert(g >= 0);
                while (g < Tgi)
                  { int k = g+1;
                    ix_pos_t pk = pA + k*st;
                    double *v1 = &(RC.e[pk]);
                    k--; pk = pk - st; 
                    double *v0 = &(RC.e[pk]);
                    (*v1) = (*v0);
                    while (k > 0)
                      { double c1 = ((double)k)/((double)(g + 1));
                        double c0 = 1.0 - c1;
                        k--; pk = pk - st; 
                        v1 = v0; v0 = &(RC.e[pk]);
                        (*v1) = c1*(*v1) + c0*(*v0);
                      }
                    g++;
                  }
                return FALSE;
              }

            /* Enumerate all other indices of {T} other than {i}, raise along axis {i}: */
            RC.ds.sz[i] = 1; /* So that {ix_descr_enum} will keep the index {i} at zero. */
            (void)ix_descr_enum( &raise_one, FALSE, &(RC.ds), NULL, NULL);
            /* Set the size of {RC} along axis {i} to the proper value: */
            RC.ds.sz[i] = Tsz[i];
          }
      }
  }

#define bz_patch_FILE_TYPE "bz_patch_t"
#define bz_patch_FILE_VERSION "2011-09-26"

void bz_patch_write ( FILE *wr, bz_patch_t *P, char *cmt )
  {
    assert(P->C.ds.na == P->d + P->r);
    
    /* Write header line: */
    filefmt_write_header(wr, bz_patch_FILE_TYPE, bz_patch_FILE_VERSION);

    /* Write comment lines, if any: */
    int ind = 0; /* Comment indentation. */
    if (cmt != NULL) { filefmt_write_comment(wr, cmt, ind, '#'); }
    
    fprintf(wr, "ddim = %d\n", P->d);
    fprintf(wr, "rinds = %d\n", P->r);

    /* !!! Should pass the {fmt} string here. */
    double_array_write(wr, &(P->C));
    
    filefmt_write_footer(wr, bz_patch_FILE_TYPE);
    fflush(wr);
  }

bz_patch_t bz_patch_read ( FILE *rd, char **cmtP )
  {
    /* Read header line: */
    filefmt_read_header(rd, bz_patch_FILE_TYPE, bz_patch_FILE_VERSION);

    /* Read the comment lines, if any: */
    char *cmt = filefmt_read_comment(rd, '#');
    if (cmtP == NULL)
      { free(cmt); }
    else
      { (*cmtP) = cmt; }

    /* Read the number of domain axes {d}: */
    uint64_t d64 = nget_uint64(rd, "ddim", 10); fget_eol(rd);
    demand(d64 <= bz_patch_MAX_DDIM, "patch domain has too many axes");
    ix_dim_t d = d64;
    
    /* Read the number of range axes {r}: */
    uint64_t r64 = nget_uint64(rd, "rinds", 10); fget_eol(rd);
    demand(r64 <= bz_patch_MAX_RINDS, "patch value has too many indices");
    ix_dim_t r = r64;
    
    /* Read the coeff array: */
    double_array_t *CP = double_array_read(rd);
    double_array_t C = (*CP);
    free(CP);
    assert(C.ds.na == d + r);
    
    filefmt_read_footer(rd, bz_patch_FILE_TYPE);
    
    return (bz_patch_t) { .d = d, .r = r, .C = C };
  }

void bz_patch_debug ( FILE *wr, char *pref, bz_patch_t *P, char *suff )
  {
    assert(P->C.ds.na == P->d + P->r);
    fprintf(wr, "%s", pref);
    /* Should pass "+10.6f" as {fmt} here. */
    bz_patch_write(wr, P, NULL);
    fprintf(wr, "%s", suff);
  }

void bz_patch_arg_debug ( FILE *wr, char *pref, bz_patch_t *P, const double x[], char *suff )
  { 
    assert(P->C.ds.na == P->d + P->r);
    int j;
    fprintf(wr, "%s[ ", pref);
    for (j = 0; j < P->d; j++,x++)
      { if (j != 0) { fprintf(wr, " "); } fprintf(wr, "%6.2f", (*x)); }
    fprintf(wr, " ]%s", suff);
  }

void bz_patch_coeff_position ( bz_patch_t *P, const ix_index_t ix[], double x[] )
  { assert(P->C.ds.na == P->d + P->r);
    int i;
    for (i = 0; i < P->d; i++)
      { ix_size_t szi = P->C.ds.sz[i];
        demand(szi > 0, "empty coeff array");
        ix_index_t ixi = ix[i];
        demand((ixi >= 0) && (ixi < szi), "invalid index");
        x[i] = ((double)ixi)/((double)szi);
      }
  }

bz_patch_t bz_patch_make_test ( bz_patch_ddim_t d, const bz_degree_t g[], bz_patch_rinds_t r, const ix_size_t rsz[] )
  { 
    fprintf(stderr, "+ bz_patch_make_test\n");
    bz_patch_t P = bz_patch_new(d, g, r, rsz);
    assert(P.C.ds.na == d + r);
    
    auto bool_t set_coeff ( const ix_index_t ix[], ix_pos_t pA, ix_pos_t pB, ix_pos_t pC );
      /* Defines the coeff of {P} with indices {ix[0..d-1]}, whose base
        position in {P.C.e} is {pA}. The arguments {pB,pC} are ignored.
        Always returns {FALSE}. */
        
    bool_t set_coeff ( const ix_index_t ix[], ix_pos_t pA, ix_pos_t pB, ix_pos_t pC )
      {
        assert(r > 0);
        double x[d]; /* Nominal position of coeff. */
        bz_patch_coeff_position(&P, ix, x);
        
        /* Compute the geometrical position {z[0..d-1]} in range space: */
        double z[3]; 
        bz_patch_test_map(d, x, z);
        
        ix_pos_t rp = 0; /* How many elements of the coeff have been set already. */

        auto bool_t set_elem ( const ix_index_t e[], ix_pos_t qA, ix_pos_t qB, ix_pos_t qC );
          /* Defines the value of element {e[0..r-1]} of the coeff of {P} with 
            indices {ix[0..d-1]}, whose position in {P.C.e} 
            is {qA}. The arguments {qB,qC} are ignored. Always returns {FALSE}. */

        bool_t set_elem ( const ix_index_t e[], ix_pos_t qA, ix_pos_t qB, ix_pos_t qC )
          { /* Compute a sequential index {rp} for the element value: */
            double val;
            if ((rp < 3) && (rp < d))
              { /* This element of the coeff is coordinate {rp} on range space: */
                val = z[rp];
              }
            else
              { /* This element of the coeff is an arbitrary attribute: */
                val = bz_patch_test_func(d, x, rp);
              }
            P.C.e[qA] = val;
            rp++;
            return FALSE;
          }
          
        (void)ix_enum(&set_elem, r, rsz, FALSE, pA, &(P.C.ds.st[d]), 0, NULL, 0, NULL);

        return FALSE;
      }
    
    (void)ix_enum(&set_coeff, d, P.C.ds.sz, FALSE, P.C.ds.bp, P.C.ds.st, 0, NULL, 0, NULL);
    fprintf(stderr, "- bz_patch_make_test\n");
    return P;
  }

/* The Archimedean perimeter parameter: */
#define PI (M_PI)

void bz_patch_test_map ( ix_dim_t d, const double x[], double z[] )
  { 
    double p[3];
    p[0] = (d > 0 ? x[0] : 0.0 );
    p[1] = (d > 1 ? x[1] : 0.0 );
    p[2] = (d > 2 ? x[2] : 0.0 );
    
    /* Apply polar-like transform to {x[0],x[1]}: */
    double theta = 2.5*M_PI*p[0];
    double rho = 1.0 + p[1] + 2*p[0]*(1-p[0]);
    z[0] = rho*cos(theta);
    z[1] = rho*sin(theta);
    
    /* Apply shearing in Z: */
    z[2] = p[2] + z[1];
  }

double bz_patch_test_func ( ix_dim_t d, const double x[], ix_axis_t i )
  { switch(i % 3)
    { double p[3];
      p[0] = (d > 0 ? x[0] : 0.0 );
      p[1] = (d > 1 ? x[1] : 0.0 );
      p[2] = (d > 2 ? x[2] : 0.0 );

      case 0:  return sin(3.0*M_PI*(p[0]+p[1]+p[2]+0.25));
      case 1:  return cos(3.0*M_PI*(p[0]-p[1]+p[2]+0.25));
      case 2:  return cos(3.0*M_PI*(p[0]-p[1]+p[2]+0.25));
      default: affirm(FALSE, "duh??"); return 0.0;
    }
  }