/* See dg_tree.h */
/* Last edited on 2009-08-21 18:34:31 by stolfilocal */

#include "dg_tree.h"
#include "mdg_grid.h"

/* ALLOCATION */

dg_tree_node_t *dg_tree_node_new(dg_tree_node_t *pa, dg_tree_node_index_t index)
  { void *v = notnull(malloc(sizeof(dg_tree_node_t)), "no mem for dg_tree_node_t");
    dg_tree_node_t *n = (dg_tree_node_t*)v;
    /* fprintf(stderr, "+ dg_tree_node_new %d\n", (int)v); */ 
    n->pa = pa;
    LOCH(*n) = NULL;
    HICH(*n) = NULL;
    n->data = NULL;
    n->index = index;
    n->nodes = 1;
    /* fprintf(stderr, "- dg_tree_node_new\n"); */
    return n;
  }

void dg_tree_node_free(dg_tree_node_t *node)
  { affirm(LOCH(*node) == NULL, "node has ch[0]");
    affirm(HICH(*node) == NULL, "node has ch[1]");
    affirm(node->data == NULL, "node has data");
    free(node);
  }
  
void dg_free_subtree(dg_tree_node_t* root)
  { if (root != NULL)
      { dg_free_subtree(LOCH(*root)); LOCH(*root) = NULL;
        dg_free_subtree(HICH(*root)); HICH(*root) = NULL;
        dg_tree_node_free(root);
      }
  }
                    
/* MODIFYING A BINARY TREE */

void dg_tree_node_split(dg_tree_node_t *n)
  { affirm(LOCH(*n) == NULL, "node has ch[0]");
    affirm(HICH(*n) == NULL, "node has ch[1]");
    LOCH(*n) = dg_tree_node_new(n, 2 * n->index);
    HICH(*n) = dg_tree_node_new(n, 2 * n->index + 1);
    n->nodes = 3;
  }

void dg_tree_node_unsplit(dg_tree_node_t *n)
  { affirm((LOCH(*n) != NULL) && (HICH(*n) != NULL), "bad unnsplit");
    affirm(n->nodes == 3, "bad unnsplit");
    dg_tree_node_free(LOCH(*n)); LOCH(*n) = NULL;
    dg_tree_node_free(HICH(*n)); HICH(*n) = NULL;
    n->nodes = 1;
  }

void dg_tree_node_enclose_children(dg_tree_node_t *n, dg_tree_node_t *ch[])
  { if (n == NULL)
      { /* No node, no children: */
        ch[0] = ch[1] = NULL;
      }
    else if ((LOCH(*n) == NULL) && (HICH(*n) == NULL))
      { /* Leaf node, children are self: */
        ch[0] = ch[1] = n;
      }
    else
      { /* Get the children of {n}, even if one is NULL: */
        ch[0] = LOCH(*n);
        ch[1] = HICH(*n);
      }
  }

bool_t dg_tree_pack_is_leaf (dg_tree_vec_t *N)
  { int NC = N->ne;
    mdg_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;
    mdg_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];
    mdg_cell_index_t c0 = N0->index;
    return (c0 % 2 == 0);
  }

void dg_tree_pack_check (dg_tree_vec_t *N, mdg_dim_t d, mdg_grid_size_t psz[], int r)
  { uint64_t NC = N->ne;
    if (psz != NULL)
      { uint64_t NA = mdg_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) 
          { mdg_cell_index_t c = Npkx->index;
            mdg_rank_t rk = mdg_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
  ( mdg_dim_t d,
    mdg_grid_size_t psz[],
    dg_tree_vec_t *N,
    mdg_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];
    mdg_grid_pos_t e[mdg_MAX_GRID_DIM]; /* Relative position of a cell in {N}. */
    mdg_grid_pos_t f[mdg_MAX_GRID_DIM]; /* Relative position of a daughter cell in {NCH} */
    mdg_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, 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}: */
                    mdg_pack_slot_index_t pky = ix_packed_position(d, f, psz, ix_order_F);
                    NCH->e[pky] = chi;
                   }
              }
          }
      }
    /* Paranoia: */
    dg_tree_pack_check(NCH, d, psz, -1);
  }

void dg_tree_pack_split
  ( mdg_dim_t d,
    mdg_grid_size_t psz[],
    dg_tree_vec_t *N,
    mdg_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];
    mdg_grid_pos_t e[mdg_MAX_GRID_DIM]; /* Relative position of a cell in {N}. */
    mdg_grid_pos_t f[mdg_MAX_GRID_DIM]; /* Relative position of a daughter cell in {NS} */
    mdg_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, 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}: */
            mdg_pack_slot_index_t pky = ix_packed_position(d, f, psz, ix_order_F);
            NS->e[pky] = chi;
          }
      }
  }

void dg_tree_pack_enum
  ( mdg_dim_t d, 
    dg_tree_t G,
    mdg_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 = mdg_pack_slot_count(d, psz);
    /* Create a node pack with copies of the root cell: */
    dg_tree_vec_t NR = dg_tree_vec_new(NC);
    { int pkx; for (pkx = 0; pkx < NC; pkx++) { NR.e[pkx] = G; } }
    /* Now recurse: */
    
    auto void do_enum(dg_tree_vec_t *N, mdg_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, mdg_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_tree_pack_child(d, psz, N, (r % d), 0, &NLO);
            dg_tree_pack_child(d, psz, N, (r % d), 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, mdg_dim_t d, mdg_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 
          { mdg_cell_print(wr, d, Ni->index); }
      }
    fprintf(wr, " ]");
  }

/* MANIPULATION OF TREE VECTORS */

vec_typeimpl(dg_tree_vec_t,dg_tree_vec,dg_tree_t);