/* See figs2d.h. */
/* Last edited on 2024-12-21 11:20:59 by stolfi */

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <values.h>

#include <affirm.h>
#include <pswr.h>
#include <pswr_iso.h>
#include <pswr_color.h>
#include <jsrandom.h>

#include <dg_grid.h>
#include <dg_tree.h>
#include <dg_locus.h>
#include <dg_pulse.h>
#include <dg_tent.h>
#include <dg_tent_basis.h>
#include <dg_plot.h>
#include <dg_makegrid.h>

#include <figtools.h>
#include <plot2d.h>

#include <figs2d.h>

#define DIM 2

/* Colors for color-banding */
static double figs2d_neg_R, figs2d_neg_G, figs2d_neg_B;
static double figs2d_zer_R, figs2d_zer_G, figs2d_zer_B;
static double figs2d_pos_R, figs2d_pos_G, figs2d_pos_B;

void figs2d_set_negative_color ( double R, double G, double B )
  { figs2d_neg_R = R; figs2d_neg_G = G; figs2d_neg_B = B; }
  
void figs2d_set_zero_color ( double R, double G, double B )
  { figs2d_zer_R = R; figs2d_zer_G = G; figs2d_zer_B = B; }
  
void figs2d_set_positive_color ( double R, double G, double B )
  { figs2d_pos_R = R; figs2d_pos_G = G; figs2d_pos_B = B; }
  
void figs_2d_full
  ( PlotOptions *o, 
    char *tag,
    interval_t rootCell[],
    dg_rank_t max_rank
  )
  {
    char levTag[300];
    dg_rank_t r;

    for (r = 0; r <= max_rank; r++)
      { sprintf(levTag, "%s-%02d", tag, r);
        fig_2d_full_level(o, levTag, rootCell, r);
      }
  }

void figs_2d_tree
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[],
    dg_rank_t max_rank
  )
  {
    /* bz_patch_t b = bz_make_test_patch(DIM, 2, 1); */
    char levTag[300];
    dg_rank_t r;

    sprintf(levTag, "%s-lvs", tag);

    fig_2d_tree_leaves(o, levTag, G, rootCell);
    for (r = 0; r <= max_rank; r++)
      { sprintf(levTag, "%s-%02d", tag, r);
        fig_2d_tree_level(o, levTag, G, rootCell, r);
      }
  }

void fig_2d_full_level
  ( PlotOptions *o, 
    char *tag, 
    interval_t rootCell[], 
    dg_rank_t rank
  )
  {
    auto void do_plot(dg_rank_t r, interval_t cell[]);
    
    void do_plot(dg_rank_t r, interval_t cell[])
      { if (r == rank)
          { pswr_rectangle(ps, 
              LO(cell[0]), HI(cell[0]), 
              LO(cell[1]), HI(cell[1]), 
              TRUE, TRUE
            );
          }
        else
          { interval_t save = cell[r % DIM];
            cell[r % DIM] = interval_split(&save, 0);
            do_plot(r+1, cell);
            cell[r % DIM] = interval_split(&save, 1);
            do_plot(r+1, cell);
            cell[r % DIM] = save;
          }
      }

    double scale = o->scale;
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    double *rgb = &(o->fill.c[0]);
    pswr_set_fill_color(ps, rgb[0], rgb[1], rgb[2]);
    do_plot(0, rootCell);
    finish_figure(o);
  }

void fig_2d_tree_leaves
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[]
  )
  {
    double scale = o->scale;
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    plot_2d_tree_leaves(ps, o, G, rootCell);
    finish_figure(o);
  }

void fig_2d_tree_level
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_rank_t rank
  )
  { bool_t plotAll;
    
    auto void do_plot(dg_rank_t r, dg_tree_node_t *n, interval_t cell[]);
    
    void do_plot(dg_rank_t r, dg_tree_node_t *n, interval_t cell[])
      { if ((n == NULL) && (! plotAll))
          { /* Do nothing */ }
        else if (r == rank)
          { /* Target level, plot the cell: */
            if (plotAll)
              { pswr_rectangle(ps, 
                  LO(cell[0]), HI(cell[0]), 
                  LO(cell[1]), HI(cell[1]), 
                  FALSE, TRUE
                );
              }
            else
              { double *rgb = &(o->fill.c[0]);
                bool_t internal = ((n != NULL) && (LOCH(*n) != NULL));
                double dim = (internal ? 0.75 : 1.0);
                pswr_set_fill_color(ps, dim*rgb[0], dim*rgb[1], dim*rgb[2]);
                pswr_rectangle(ps, 
                  LO(cell[0]), HI(cell[0]), 
                  LO(cell[1]), HI(cell[1]),
                  TRUE, TRUE
                );
              }
          }
      else if ((r < rank) && ((n != NULL) || plotAll))
          { /* Recurse: */
            interval_t save = cell[r % DIM];
            cell[r % DIM] = interval_split(&save, 0);
            do_plot(r+1, (n != NULL ? LOCH(*n) : NULL), cell);
            cell[r % DIM] = interval_split(&save, 1);
            do_plot(r+1, (n != NULL ? HICH(*n) : NULL), cell);
            cell[r % DIM] = save;
          }
      }

    double scale = o->scale;
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    /* First draw all cells of that level in light color: */
    plotAll = TRUE; 
    pswr_set_pen(ps, 0.5, 0.5, 0.5,  0.15,  0.0, 0.0);
    do_plot(0, G, rootCell);
    /* Now draw only the existing cells of that level: */
    plotAll = FALSE; 
    pswr_set_pen(ps, 0.0, 0.0, 0.0,  0.15,  0.0, 0.0);
    do_plot(0, G, rootCell);
    finish_figure(o);
  }

void fig_2d_bezier_array(PlotOptions *o, char *tag, dg_degree_t g)
  { /* Get an arbitrary rectangular cell: */
    interval_t cell[DIM];
    dg_cell_box_canonical(DIM, dg_ROOT_CELL, cell);
    
    /* Start the figure: */
    double mag = (o->eps ? 1.0 : 5.0);
    start_figure(o, tag, tag, &(cell[0]), mag*25.4);
    
    /* Draw the cell: */
    pswr_set_pen(ps, 0,0,0,  0.15,  0.0, 0.0);
    double *rgb = &(o->fill.c[0]);
    pswr_set_fill_color(ps, rgb[0], rgb[1], rgb[2]);
    pswr_rectangle(ps, 
      LO(cell[0]), HI(cell[0]), 
      LO(cell[1]), HI(cell[1]), 
      TRUE, TRUE
    );
    
    /* Draw its Bézier array: */
    pswr_set_pen(ps, 0,0,0,  0.15,  0.0, 0.0);
    plot_2d_bezier_array(ps, o, cell, g, mag);

    /* Finish the figure: */
    finish_figure(o);
  }

void fig_2d_bezier_constraints
  ( PlotOptions *o, 
    char *tag, 
    dg_cont_t c, 
    dg_degree_t g
  )
  { /* Get a large cell: */
    interval_t cell[DIM];
    dg_cell_box_canonical(DIM, dg_ROOT_CELL, cell);
    
    /* Start the figure: */
    double mag = (o->eps ? 2.0 : 5.0);
    start_figure(o, tag, tag, &(cell[0]), mag*25.4);
    
    /* Draw the four sub-cells, linkages, and Bézier nodes: */
    int pass, i, j;
    double xstep = (HI(cell[0]) - LO(cell[0]))/2;
    double ystep = (HI(cell[1]) - LO(cell[1]))/2;
    interval_t subCell[DIM];
    for (pass = 0; pass <= 2; pass++)
      { /* Pass 0 = cells, 1 = linkages, 2 = nodes. */
        for (i = 0; i < 2; i++)
          { LO(subCell[0]) = LO(cell[0]) + i*xstep;
            HI(subCell[0]) = LO(cell[0]) + (i+1)*xstep;
            for (j = 0; j < 2; j++) 
              { LO(subCell[1]) = LO(cell[1]) + j*ystep;
                HI(subCell[1]) = LO(cell[1]) + (j+1)*ystep;
                if (pass == 0)
                  { /* Draw the sub-cell: */
                    double *rgb = &(o->fill.c[0]);
                    double dim = ((i+j)%2 == 0 ? 0.90 : 1.0);
                    pswr_set_pen(ps, 0,0,0,  0.15,  0.0, 0.0);
                    pswr_set_fill_color(ps, dim*rgb[0], dim*rgb[1], dim*rgb[2]);
                    pswr_rectangle(ps, 
                      LO(subCell[0]), HI(subCell[0]), 
                      LO(subCell[1]), HI(subCell[1]), 
                      TRUE, TRUE
                    );
                  }
                if (pass == 1)
                  { /* Draw the continuity constraints: */
                    pswr_set_pen(ps, 0,0,0,  mag*0.20,  0.0, 0.0);
                    plot_2d_bezier_linkages(ps, o, subCell, c, g);
                  }
                if (pass == 2)
                  { /* Draw the Bézier controls: */
                    pswr_set_pen(ps, 0,0,0,  0.15,  0.0, 0.0);
                    plot_2d_bezier_array(ps, o, subCell, g, mag);
                  }
              }
          }
      }
    /* Finish the figure: */
    finish_figure(o);
  }

void figs_2d_basis
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_pulse_family_t fam[], 
    bool_t superior
  )
  { 
    char locusTag[300];
    int i, j;
    dg_tent_vec_t bas = dg_tent_basis_get(DIM, fam, G, superior);
    
    for (j = 0; j < DIM; j++)
      { affirm(fam[j].g >= 2*fam[j].c + 1, "overconstrained splines not supported yet"); }

    /* !!! dg_locus_vec_t loc = dg_tent_basis_loci(&bas, fam);  */
    /* !!! fig_2d_loci(o, tag, G, rootCell, loc); */

    for (i = 0; i < bas.nel; i++)
      { dg_tent_t t = bas.el[i];
        dg_locus_t E = loc.el[i];
        fprintf(stderr, "[%03d] ", i); 
        dg_locus_print(stderr, DIM, E);
        fprintf(stderr, "\n"); 
        sprintf(locusTag, "%s-%06d-%d", tag, (int)E.cell, E.norm);
        fig_2d_locus_and_star(o, locusTag, G, rootCell, E);
        bool_t show;
        if (loc.nel <= 100)
          { show = TRUE; }
        else
          { int step = loc.nel/20;
            show = ((i < 10) || (i % step == 0) || (i >= loc.nel-10));
          }
        if (show)
          { figs_2d_locus_tents
              ( o, locusTag, 
                G, rootCell, E, 
                pkind, c, g
              );
          }
      }
      
    free(loc.el);
  }

void figs_2d_locus_tents
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[],
    dg_locus_t E, 
    dg_pulse_kind_t pkind, 
    dg_cont_t c,
    dg_degree_t g
  )
  { 
    char elemTag[300];
    /* Get the tents associated to this locus: */
    dg_locus_vec_t loc = dg_locus_vec_new(1);
    loc.nel = 1; loc.el[0] = E;
    dg_tent_vec_t tv = dg_tent_basis(DIM, DG_TK_H, c, g, loc);
    fprintf(stderr, "locus ");
    dg_locus_print(stderr, DIM, E);
    fprintf(stderr, " created %d tents\n", tv.nel);
    int i;
    for (i = 0; i < tv.nel; i++)
      { dg_tent_t t = tv.el[i];
        sprintf(elemTag, "%s-%02d", tag, i);
        fig_2d_tent(o, elemTag, G, rootCell, pkind, c, g, t);
      }
    free(loc.el); free(tv.el);
  }

void figs_2d_mother_pulse_catalog
  ( PlotOptions *o, 
    dg_pulse_kind_t pkind,
    dg_cont_t cMin,
    dg_cont_t cMax, 
    dg_degree_t gOver
  )
  {
    char pkch = dg_pulse_kind_to_char(pkind);
    dg_cont_t c;
    for (c = cMin; c <= cMax; c++)
      { /* The variable {gMin} holds the min degree valid for {pkind,c} found so far: */
        dg_degree_t gMin = dg_PULSE_MAX_DEGREE + 1;  /* For now. */
        /* Loop over all degrees {g}: */
        dg_degree_t g;
        for (g = 0; g <= gMin + gOver; g++)
          { dg_pulse_family_t fam = dg_pulse_family(pkind, c, g);
            if (fam.nmp > 0)
              { /* Found a valid degree {g} for the pair {pkind,c}. */
                if (g < gMin) { /* First valid {g}: */ gMin = g; }
                /* Get the maximum pulse support size {mszMax} for this family: */
                dg_grid_size_t mszMax = dg_pulse_max_supp_count(&fam);
                /* Loop over grid sizes to show effect of self-overlap: */
                dg_rank_t r = 0; /* Grid level. */
                while(TRUE) 
                  { /* Get the grid size {gsz} along each axis: */
                    dg_grid_size_t gsz[DIM];
                    dg_grid_size(DIM, r, gsz);
                    /* Get the min grid size {gszMin} along any axis: */
                    dg_grid_size_t gszMin = ipow(2, r/DIM);
                    /* Use the same arg and value ranges for all plots with same {c,gsz}: */
                    interval_t fr[3];
                    compute_2d_common_tent_ranges(c, gsz, 1.0, 1.0, /*out*/ fr);
                    char *tag = jsprintf("d2-k%c-g%d-c%d-r%02d", pkch, g, c, r);
                    fprintf(stderr, "considering %s\n", tag);
                    dg_pulse_family_t fam = dg_pulse_family(pkind, c, g); 
                    figs_2d_mother_vertex_tent_set(o, tag, &fam, gsz, fr);
                    free(tag);
                    if (gszMin > mszMax) { break; }
                    r++;
                  }
              }
          }
      }
  }

void figs_2d_mother_vertex_tent_set
  ( PlotOptions *o, 
    char *tag,               /* Tag for figure/file name. */
    dg_pulse_family_t *fam,  /* The tent mother pulse family. */
    dg_grid_size_t gsz[],    /* X- and Y-domain periods (number of grid intervals). */
    interval_t fr[]          /* X-, Y-, and function ranges of interest. */
  )
  { 
    char elemTag[300];
    /* Number of tents associated with a vertex: */
    int NT = ipow(g-2*c, 2); ???
    fprintf(stderr, "plotting %d tents\n", NT);
    dg_tent_mother_index_t tix;
    for (tix = 0; tix < NT; tix++)
      { sprintf(elemTag, "%s-%02d", tag, tix);
        fig_2d_mother_vertex_tent(o, elemTag, pkind, c, g, tix, gsz, fr);
      }
  }

void fig_2d_loci
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_locus_vec_t loc
  )
  { int i;
    /* Start the figure: */
    double mag = (o->eps ? 5.0 : 5.0);
    start_figure(o, tag, tag, &(rootCell[0]), mag*25.4);
    plot_2d_tree_leaves(ps, o,  G, rootCell);
    for (i = 0; i < loc.nel; i++)
      { dg_locus_t E = loc.el[i];
        fprintf(stderr, "[%03d] ", i); 
        dg_locus_print(stderr, DIM, E);
        fprintf(stderr, "\n"); 
        plot_2d_locus(ps, o, E, G, rootCell);
      }
    finish_figure(o);
  }

void fig_2d_locus_and_star
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_locus_t E
  )
  { double scale = o->scale;
    double hstep = (o->eps ? 1.0 : 2.0)/scale; /* Hatchline spacing */
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    plot_2d_tree_leaves(ps, o,  G, rootCell);
    plot_2d_locus_star(ps, o, E, G, rootCell, hstep);
    if (dg_locus_dimension(DIM, E) < 2)
      { plot_2d_locus(ps, o, E, G, rootCell); }
    finish_figure(o);
  }

void fig_2d_tent
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_pulse_family_t fam[], 
    dg_pulse_t p[],
    dg_tent_t t
  )
  { 
    double scale = o->scale;
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    dg_rank_t plotDepth = compute_2d_plot_depth(rootCell, scale, o->meshSize);
    plot_2d_tree_leaves(ps, o,  G, rootCell);
    /* Estimate maximum function value: */
    interval_t fRange = compute_nd_tent_range(2, fam, p, t);
    double fMax = HI(fRange);
    double fStep = fMax/5.0; /* pswr_round_to_nice(fMax)/5.0; */
    double fStart = fStep/2;
    double tmp = 1.0;
    while ((tmp = tmp/2) >= fMax) { fStep /= 2; }
    
    int kMin = pswr_inf_isoline(fStart, fStep, -fMax);
    int kMax = pswr_sup_isoline(fStart, fStep, +fMax);
    int N;
    double *R, *G, *B;

    pswr_make_color_table
      ( fStart, fStep, kMin, kMax,
        figs2d_neg_R, figs2d_neg_G, figs2d_neg_B,
        figs2d_zer_R, figs2d_zer_G, figs2d_zer_B,
        figs2d_pos_R, figs2d_pos_G, figs2d_pos_B,
        &N, &R, &G, &B
      );
    plot_2d_tent(ps, o, 
      pkind, c, g, t, 
      G, rootCell, 0, 
      plotDepth, 
      fStart, fStep, kMin, kMax, R, G, B
    );
    finish_figure(o);
    free(R); free(G); free(B);
  }
 
void fig_2d_mother_vertex_tent
  ( PlotOptions *o, 
    char *tag, 
    dg_pulse_kind_t pkind,
    dg_cont_t c,
    dg_degree_t g, 
    dg_tent_mother_index_t tix,    /* Index of tent in set. */
    dg_grid_size_t gsz[],    /* X- and Y-domain periods (number of grid intervals). */
    interval_t fr[]     /* X-, Y-, and function ranges of interest. */
  )
  { 
    double scale = (o->eps ? 1.0 : 3.0)*25.4;
    start_figure(o, tag, tag, &(fr[0]), scale);
    dg_rank_t plotDepth = compute_2d_plot_depth(&(fr[0]), scale, o->meshSize);
    /* Estimate maximum function value: */
    interval_t fRange = fr[2];
    double fMax = HI(fRange);
    double fStep = fMax/5.0; /* pswr_round_to_nice(fMax)/5.0; */
    double fStart = fStep/2;
    double tmp = 1.0;
    while ((tmp = tmp/2) >= fMax) { fStep /= 2; }
    
    int kMin = pswr_inf_isoline(fStart, fStep, -fMax);
    int kMax = pswr_sup_isoline(fStart, fStep, +fMax);
    int N;
    double *R, *G, *B;

    pswr_make_color_table
      ( fStart, fStep, kMin, kMax,
        figs2d_neg_R, figs2d_neg_G, figs2d_neg_B,
        figs2d_zer_R, figs2d_zer_G, figs2d_zer_B,
        figs2d_pos_R, figs2d_pos_G, figs2d_pos_B,
        &N, &R, &G, &B
      );
    plot_2d_mother_vertex_tent(ps, o, 
      pkind, c, g, tix, gsz, 
      &(fr[0]), 
      plotDepth, 
      fStart, fStep, kMin, kMax, R, G, B
    );
    finish_figure(o);
    free(R); free(G); free(B);
  }

void fig_2d_sample_spline
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_pulse_kind_t pkind,
    dg_cont_t c,
    dg_degree_t g, 
    bool_t superior
  )
  { affirm(g >= 2*c + 1, "overconstrained splines not supported yet");
    bool_t vtxOnly = (g == 2*c + 1);
    dg_locus_vec_t loc = dg_find_critical_loci(DIM, G, vtxOnly, superior);
    dg_tent_vec_t tv = dg_tent_basis(DIM, DG_TK_H, c, g, loc);
    fprintf(stderr, " created %d tents\n", tv.nel);
    double_vec_t coef = double_vec_new(tv.nel); /* Spline coef vector */
    int i;
    for (i = 0; i < tv.nel; i++)
      { dg_tent_t t = tv.el[i];
        dg_rank_t r = dg_cell_rank(t.cell);
        double magn = pow(0.5, ((double)r)/4.0); 
        coef.el[i] = magn*(2.0*drandom()-1.0);
      }
    interval_t fr = (interval_t) {{ -1.0, +1.0 }};
    fig_2d_spline(o, tag, G, rootCell, pkind, c, g, tv, coef, fr);
    free(loc.el); free(tv.el);
  }

void fig_2d_spline
  ( PlotOptions *o, 
    char *tag, 
    dg_tree_node_t *G,
    interval_t rootCell[], 
    dg_pulse_kind_t pkind,
    dg_cont_t c,
    dg_degree_t g, 
    dg_tent_vec_t tv,
    double_vec_t coef,
    interval_t fr
  )
  { 
    double scale = (o->eps ? 3.0 : 5.0)*25.4;
    start_figure(o, tag, tag, &(rootCell[0]), scale);
    dg_rank_t plotDepth = compute_2d_plot_depth(rootCell, scale, o->meshSize);
    /* Estimate maximum function value: */
    double fMin = LO(fr);
    double fMax = HI(fr);
    double fMaxAbs = (fabs(fMax) > fabs(fMin) ? fabs(fMax) : fabs(fMin));
    double fStep = fMaxAbs/5.0; /* pswr_round_to_nice(fMax)/5.0; */
    double fStart = fStep/2;
    double tmp = 1.0;
    while ((tmp = tmp/2) >= fMax) { fStep /= 2; }
    
    int kMin = pswr_inf_isoline(fStart, fStep, fMin);
    int kMax = pswr_sup_isoline(fStart, fStep, fMax);
    int N;
    double *R, *G, *B;

    pswr_make_color_table
      ( fStart, fStep, kMin, kMax,
        figs2d_neg_R, figs2d_neg_G, figs2d_neg_B,
        figs2d_zer_R, figs2d_zer_G, figs2d_zer_B,
        figs2d_pos_R, figs2d_pos_G, figs2d_pos_B,
        &N, &R, &G, &B
      );
    plot_2d_spline(ps, o, 
      pkind, c, g, tv, coef, 
      G, rootCell, 0, 
      plotDepth, 
      fStart, fStep, kMin, kMax, R, G, B
    );
    finish_figure(o);
    free(R); free(G); free(B);
  }