/* See dg_tree_pack.h */
/* Last edited on 2020-10-03 21:31:06 by jstolfi */

#include <stdlib.h>
#include <assert.h>

#include <affirm.h>
#include <ix.h>
#include <vec.h>

#include <dg_grid.h>
#include <dg_pack.h>
#include <dg_tree.h>
#include <dg_tree_pack.h>

bool_t dg_tree_pack_is_leaf (dg_tree_vec_t *N)
  { int NC = N->ne;
    dg_pack_slot_index_t pkx;
    bool_t leaf = FALSE;
    for (pkx = 0; pkx < NC; pkx++)
      { dg_tree_node_t *Npkx = N->e[pkx];
        if (Npkx == NULL) { return FALSE; }
        if ((LOCH(*Npkx) == NULL) || (HICH(*Npkx) == NULL)) { leaf = TRUE; }
      }
    return leaf;
  }

bool_t dg_tree_pack_is_even(dg_tree_vec_t *N)
  { int NC = N->ne;
    dg_pack_slot_index_t pkx;
    for (pkx = 0; pkx < NC; pkx++)
      { dg_tree_node_t *Npkx = N->e[pkx];
        if (Npkx == NULL) { return FALSE; }
      }
    dg_tree_node_t *N0 = N->e[pkx];
    dg_cell_id_t c0 = N0->id;
    return (c0 % 2 == 0);
  }

void dg_tree_pack_check (dg_tree_vec_t *N, dg_dim_t d, dg_grid_size_t psz[], int r)
  { uint64_t NC = N->ne;
    if (psz != NULL)
      { uint64_t NA = dg_pack_slot_count(d, psz);
        demand(NA == NC, "wrong cell count");
      }
    uint64_t pkx;
    for (pkx = 0; pkx < NC; pkx++) 
      { dg_tree_t Npkx = N->e[pkx]; 
        if (Npkx != NULL) 
          { dg_cell_id_t c = Npkx->id;
            dg_rank_t rk = dg_cell_rank(c); 
            if (r == -1) { r = rk; } else { demand(r == rk, "wrong rank"); }
            /* Should check position of cell {c} against {pkx}. */
          }
      }
  }

void dg_tree_pack_child
  ( dg_dim_t d,
    dg_grid_size_t psz[],
    dg_tree_vec_t *N,
    dg_axis_t i,
    int which,
    dg_tree_vec_t *NCH
  )
  { int NC = N->ne;
    
    /* Allocate child pack: */
    (*NCH) = dg_tree_vec_new(NC);
    
    /* Paranoia: */
    /* dg_tree_pack_check(N, d, psz, -1); */
    
    /* Enumerate cells of parent pack, get their relative positions {e}: */
    dg_tree_node_t *ch[2];
    dg_grid_pos_t e[dg_dim_MAX]; /* Relative position of a cell in {N}. */
    dg_grid_pos_t f[dg_dim_MAX]; /* Relative position of a daughter cell in {NCH} */
    dg_pack_slot_index_t pkx;
    for (pkx = 0; pkx < NC; pkx++)
      { /* Get relative position {e} of node {N[pkx]} in pack: */
        ix_packed_indices(d, pkx, 0, psz, ix_order_F, e);
        /* Is any child of this node inside {NCH}? */
        if (2*e[i] < psz[i] + which)
          { /* Get children of node: */
            dg_tree_node_t *n = N->e[pkx];
            if (n == NULL)
              { ch[0] = NULL; ch[1] = NULL; }
            else
              { ch[0] = LOCH(*n); ch[1] = HICH(*n); }
            /* Store them in {NCH}, if appropriate: */
            ix_assign(d, f, e);
            int ich;
            for (ich = 0; ich < 2; ich++)
              { dg_tree_node_t *chi = ch[ich];
                /* Get relative position {f} of {chi} in {NCH}: */
                f[i] = 2*e[i] + ich - which;
                if ((f[i] >= 0) && (f[i] < psz[i]))
                  { /* Compute slot index of {chi} in {NCH}: */
                    dg_pack_slot_index_t pky = ix_packed_position(d, f, 0, psz, ix_order_F);
                    NCH->e[pky] = chi;
                   }
              }
          }
      }
    /* Paranoia: */
    /* dg_tree_pack_check(NCH, d, psz, -1); */
  }

void dg_tree_pack_split
  ( dg_dim_t d,
    dg_grid_size_t psz[],
    dg_tree_vec_t *N,
    dg_axis_t i,
    dg_tree_vec_t *NS
  )
  { int NC = N->ne;
    
    /* Allocate child pack: */
    (*NS) = dg_tree_vec_new(2*NC);
    
    /* Enumerate cells of parent pack, get their relative positions {e}: */
    dg_tree_node_t *ch[2];
    dg_grid_pos_t e[dg_dim_MAX]; /* Relative position of a cell in {N}. */
    dg_grid_pos_t f[dg_dim_MAX]; /* Relative position of a daughter cell in {NS} */
    dg_pack_slot_index_t pkx;
    for (pkx = 0; pkx < NC; pkx++)
      { /* Get relative position {e} of node {N[pkx]} in pack: */
        ix_packed_indices(d, pkx, 0, psz, ix_order_F, e);
        /* Get children of node: */
        dg_tree_node_t *n = N->e[pkx];
        if (n == NULL)
          { ch[0] = ch[1] = NULL; }
        else
          { ch[0] = LOCH(*n); ch[1] = HICH(*n); }
        /* Store them in {NS}: */
        ix_assign(d, f, e);
        int ich;
        for (ich = 0; ich < 2; ich++)
          { dg_tree_node_t *chi = ch[ich];
            /* Get relative position {f} of {chi} in {NS}: */
            f[i] = 2*e[i] + ich;
            /* Compute slot index of {chi} in {NS}: */
            dg_pack_slot_index_t pky = ix_packed_position(d, f, 0, psz, ix_order_F);
            NS->e[pky] = chi;
          }
      }
  }

void dg_tree_pack_enum
  ( dg_dim_t d, 
    dg_tree_t G,
    dg_grid_size_t psz[], 
    dg_tree_pack_op_t op
  )
  { /* Check root: */
    demand(G != NULL, "null root");
    /* How many trees in a pack? */
    uint64_t NC = dg_pack_slot_count(d, psz);
    /* Create a node pack with copies of the root cell: */
    vec_size_t NCv = (vec_size_t)NC;
    assert(NC == (uint64_t)NCv); /* Check for size overflow. */
    dg_tree_vec_t NR = dg_tree_vec_new(NCv);
    { int pkx; for (pkx = 0; pkx < NC; pkx++) { NR.e[pkx] = G; } }
    /* Now recurse: */
    
    auto void do_enum(dg_tree_vec_t *N, dg_rank_t r);
      /* Enumerates all valid leaf tree packs descending from {N}.
        Assumes that all non-null slots of the pack are from level {r},
        and that the lower slot in the pack is not null. */
      
    do_enum(&NR, 0);
    free(NR.e);
    
    void do_enum(dg_tree_vec_t *N, dg_rank_t r)
      {
        /* Visit if appropriate: */
        bool_t is_leaf = dg_tree_pack_is_leaf(N);
        op(N, is_leaf);
        if (! is_leaf)
          { /* Get lower tree of pack: */
            dg_tree_node_t *n = N->e[0];
            assert(n != NULL);
            assert((LOCH(*n) != NULL) && (HICH(*n) != NULL));
            /* Get children blocks: */
            dg_tree_vec_t NLO, NHI;
            dg_axis_t ay = dg_split_axis(d, r);
            dg_tree_pack_child(d, psz, N, ay, 0, &NLO);
            dg_tree_pack_child(d, psz, N, ay, 1, &NHI);
            do_enum(&NLO, r+1);
            free(NLO.e);
            do_enum(&NHI, r+1);
            free(NHI.e);
          }
      }
  }

void dg_tree_pack_print(FILE *wr, dg_dim_t d, dg_grid_size_t psz[], dg_tree_vec_t *N)
  { int i;
    fprintf(wr, "[");
    for (i = 0; i < N->ne; i++)
      { dg_tree_node_t *Ni = N->e[i];
        fprintf(wr, " ");
        if (Ni == NULL) 
          { fprintf(wr, "Ø"); }
        else 
          { dg_cell_print(wr, d, Ni->id); }
      }
    fprintf(wr, " ]");
  }

/* MANIPULATION OF TREE VECTORS */

vec_typeimpl(dg_tree_vec_t,dg_tree_vec,dg_tree_t);