#include <v_lib.h>
#include "decomp.h"
#include "iplmath.h"
#include "messages.h"
#include "istrings.h"
#include "util.h"

V_LOCAL btree_p V_CALL
bt_img_decomp_r(matrix_p data,
    u32_t left,u32_t top,u32_t right,u32_t bottom,
    vec_t threshold,
    u32_t depth)
{
    u32_t vsplit, hsplit;
    vec_t min, max;
    btree_p node;
    btdecomp_node_t node_data;
    vboolean vspan = left == right; /* vertical span of pixels */
    vboolean hspan = top == bottom; /* horizontal span */
    node_data.left = left;
    node_data.top = top;
    node_data.right = right;
    node_data.bottom = bottom;
    node_data.depth = depth;
    if(hspan && vspan)
    {
        /* one pixel: leaf */
        node_data.value = data->cell[top][left];
        return btStore(&node_data,sizeof(node_data),NULL,0);
    }
    matSubmatrixMinMax(data,left,top,right,bottom,&min,&max);
    node_data.value = (min + max) * (vec_t)0.5;
    node = btStore(&node_data,sizeof(node_data),NULL,0);
    if((max - min) < threshold) /* difference too small: leaf */
        return node;
    depth++;
    /* choose split direction */
    if((right-left) > (bottom-top))
    {
        /* split into left and right submatrices */
        hsplit = left + ((right - left) / 2);
        node->children[0] =
            bt_img_decomp_r(data,left,top,hsplit,bottom,threshold,depth);
        node->children[0]->parent = node;
        node->children[1] =
            bt_img_decomp_r(data,hsplit+1,top,right,bottom,threshold,depth);
        node->children[1]->parent = node;
    }
    else
    {
        /* split into upper and lower submatrices */
        vsplit = top + ((bottom - top) / 2);
        node->children[0] =
            bt_img_decomp_r(data,left,top,right,vsplit,threshold,depth);
        node->children[0]->parent = node;
        node->children[1] =
            bt_img_decomp_r(data,left,vsplit+1,right,bottom,threshold,depth);
        node->children[1]->parent = node;
    }
    return node;
}

V_LOCAL void
bt_img_describe_r(bstack_p info_stack,btree_p root)
{
    btdecomp_node_p nd;
    btdecomp_leaf_info_t info;
    if(!root->children[0]) /* leaf */
    {
        nd = (btdecomp_node_p)root->data;
        info.depth = nd->depth;
        info.left = nd->left;
        info.top = nd->top;
        info.right = nd->right;
        info.bottom = nd->bottom;
        info.area = (info.right - info.left + 1) * (info.bottom - info.top + 1);
        info.value = nd->value;
        stackPush(info_stack,&info);
        return;
    }
    bt_img_describe_r(info_stack,root->children[0]);
    bt_img_describe_r(info_stack,root->children[1]);
}

bt_decomp_p
iplDecomposeImage(matrix_p m,vec_t threshold,
    u32_t histogram_divs)
{
    u32_t i, d_offs, num_leafs, *counts;
    btdecomp_leaf_info_p base, ptr;
    bt_decomp_p self = (bt_decomp_p)memAllocZ(sizeof(bt_decomp_t));
    self->root = bt_img_decomp_r(m,0,0,m->cols - 1,m->rows - 1,threshold,0);
    self->source_rows = m->rows;
    self->source_cols = m->cols;
    stackInit(&self->info_stack,sizeof(btdecomp_leaf_info_t));
    bt_img_describe_r(&self->info_stack,self->root);
    num_leafs = stackCount(&self->info_stack);
    base = (btdecomp_leaf_info_p)stackGetStorage(&self->info_stack);
    self->min_depth = MAXLONG;
    self->max_depth = 0;
    self->min_area = m->rows * m->cols;
    self->max_area = 0;
    self->min_value = MAXVEC;
    self->max_value = -MAXVEC;
    for(i = 0;i < num_leafs;i++)
    {
        ptr = base + i;
        if(self->min_depth > ptr->depth)
            self->min_depth = ptr->depth;
        if(self->max_depth < ptr->depth)
            self->max_depth = ptr->depth;
        if(self->min_area > ptr->area)
            self->min_area = ptr->area;
        if(self->max_area < ptr->area)
            self->max_area = ptr->area;
        if(self->min_value > ptr->value)
            self->min_value = ptr->value;
        if(self->max_value < ptr->value)
            self->max_value = ptr->value;
    }
    self->depth_levels = self->max_depth - self->min_depth + 1;
    self->leafs_at_depth = 
        (u32_p)memAllocZ(sizeof(u32_t) * self->depth_levels);
    counts = (u32_p)memAllocZ(sizeof(u32_t) * self->depth_levels);
    for(i = 0;i < num_leafs;i++)
    {
        ptr = base + i;
        self->leafs_at_depth[ptr->depth - self->min_depth]++;
    }
    self->values_at_depth =
        (vector_p*)memAlloc(sizeof(vector_p) * self->depth_levels);
    self->histogram_at_depth =
        (vector_p*)memAlloc(sizeof(vector_p) * self->depth_levels);
    self->histogram_space_at_depth =
        (vector_p*)memAlloc(sizeof(vector_p) * self->depth_levels);
    for(i = 0;i < self->depth_levels;i++)
    {
        if(self->leafs_at_depth[i])
            self->values_at_depth[i] = vecAlloc(self->leafs_at_depth[i]);
        else
            self->values_at_depth[i] = 0;
    }
    for(i = 0;i < num_leafs;i++)
    {
        ptr = base + i;
        d_offs = ptr->depth - self->min_depth;
        self->values_at_depth[d_offs]->elements[counts[d_offs]] =
            ptr->value;
        counts[d_offs]++;
    }
    for(i = 0;i < self->depth_levels;i++)
    {
        if(self->values_at_depth[i])
        {
            self->histogram_at_depth[i] = vecAlloc(histogram_divs);
            self->histogram_space_at_depth[i] = vecAlloc(histogram_divs);
            vecHistogram(self->histogram_at_depth[i],self->values_at_depth[i],
                false,self->histogram_space_at_depth[i]);
        }
        else
        {
            self->histogram_at_depth[i] = NULL;
            self->histogram_space_at_depth[i] = NULL;
        }
    }
    memFree(counts);
    return self;
}
    
void
iplSaveDecompositionDescription(char *fname,
    bt_decomp_p decomp,
    vboolean save_values,
    vboolean save_histograms,
    vboolean relative_frequencies)
{
    u32_t i;
    vec_p ptr, sptr, endptr;
    vec_t freq_scale = 1;
    vhandle_t out = fileOpenWrite(fname);
    V_ASSERT(out != LIBV_INVALID_HANDLE);
    V_ASSERT(decomp);
    filePrintf(out,"%s (%u,%u)\n",ipl_key_source_dimension,
        decomp->source_rows,decomp->source_cols);
    filePrintf(out,"%s (%u,%u)\n",ipl_key_depth_range,
        decomp->min_depth,decomp->max_depth);
    filePrintf(out,"%s (%u,%u)\n",ipl_key_area_range,
        decomp->min_area,decomp->max_area);
    filePrintf(out,"%s (%g,%g)\n",ipl_key_value_range,
        decomp->min_value,decomp->max_value);
    if(!save_values && !save_histograms)
    {
        for(i = 0;i < decomp->depth_levels;i++)
            filePrintf(out,"%u\t%u\n",i + decomp->min_depth,
                decomp->leafs_at_depth[i]);
        fileClose(out);
        return;
    }
    /* 20.04.2011: at 1st pass, save separate table of (depth,numleafs) */
    filePrintf(out,"%s\t%s\n",ipl_key_at_depth,ipl_key_leafs);
    for(i = 0;i < decomp->depth_levels;i++)
    {
        filePrintf(out,"%u\t%u\n",i + decomp->min_depth,decomp->leafs_at_depth[i]);
    }
    for(i = 0;i < decomp->depth_levels;i++)
    {
        filePrintf(out,"%s %u\n",ipl_key_at_depth,i + decomp->min_depth);
        filePrintf(out,"%s %u\n",ipl_key_leafs,decomp->leafs_at_depth[i]);
        if(save_values)
        {
            filePrintf(out,"%s\n",ipl_key_values);
            if(decomp->values_at_depth[i])
            {
                ptr = decomp->values_at_depth[i]->elements;
                endptr = ptr + decomp->values_at_depth[i]->size;
                for(;ptr < endptr;ptr++)
                    filePrintFloatStr(out,dotchar,*ptr,ipl_eol_string);
            }
        }
        if(save_histograms)
        {
            if(relative_frequencies && decomp->values_at_depth[i])
            {
                freq_scale = (vec_t)1.0 / decomp->leafs_at_depth[i];
            }
            else
                freq_scale = 1;
            if(decomp->histogram_at_depth[i])
            {
                filePrintf(out,"%s\n",ipl_key_histogram);
                ptr = decomp->histogram_at_depth[i]->elements;
                sptr = decomp->histogram_space_at_depth[i]->elements;
                endptr = ptr + decomp->histogram_at_depth[i]->size;
                for(;ptr < endptr;ptr++,sptr++)
                {
                    filePrintFloatStr(out,dotchar,
                        *sptr,ipl_tab_string);
                    filePrintFloatStr(out,dotchar,
                        *ptr * freq_scale,ipl_eol_string);
                }
            }
        }
    }
    fileClose(out);
}

void
iplFreeDecomposition(bt_decomp_p decomp)
{
    u32_t i;
    if(!decomp)
        return;
    btDeleteAllNodes(decomp->root,NULL);
    stackClear(&decomp->info_stack);
    for(i = 0;i < decomp->depth_levels;i++)
    {
        if(decomp->values_at_depth[i])
            memFree(decomp->values_at_depth[i]);
        if(decomp->histogram_at_depth[i])
            memFree(decomp->histogram_at_depth[i]);
        if(decomp->histogram_space_at_depth[i])
            memFree(decomp->histogram_space_at_depth[i]);
    }
    memFree(decomp->values_at_depth);
    memFree(decomp->histogram_at_depth);
    memFree(decomp->histogram_space_at_depth);
    memFree(decomp->leafs_at_depth);
    memFree(decomp);
}

V_LOCAL void
bt_img_draw_decomp_r(matrix_p surface,btree_p root)
{
    btdecomp_node_p nd;
    if(!root->children[0]) /* leaf */
    {
        nd = (btdecomp_node_p)root->data;
        matSetRect(surface,nd->left,nd->top,nd->right,nd->bottom,nd->value);
        return;
    }
    bt_img_draw_decomp_r(surface,root->children[0]);
    bt_img_draw_decomp_r(surface,root->children[1]);
}

V_LOCAL void
bt_img_draw_scale_decomp_r(matrix_p surface,btree_p root,
    u32_t depth_base,vec_t depth_to_color[])
{
    btdecomp_node_p nd;
    if(!root->children[0]) /* leaf */
    {
        nd = (btdecomp_node_p)root->data;
        matSetRect(surface,nd->left,nd->top,nd->right,nd->bottom,
            depth_to_color[nd->depth - depth_base]);
        return;
    }
    bt_img_draw_scale_decomp_r(surface,root->children[0],
        depth_base,depth_to_color);
    bt_img_draw_scale_decomp_r(surface,root->children[1],
        depth_base,depth_to_color);
}

void
iplDrawDecomposition(matrix_p surface,bt_decomp_p decomp)
{
    V_ASSERT(surface->rows == decomp->source_rows);
    V_ASSERT(surface->cols == decomp->source_cols);
    bt_img_draw_decomp_r(surface,decomp->root);
}

void
iplDrawDecompositionScale(matrix_p surface,bt_decomp_p decomp,
    vector_p scale)
{
    V_ASSERT(surface->rows == decomp->source_rows);
    V_ASSERT(surface->cols == decomp->source_cols);
    V_ASSERT(scale->size == decomp->depth_levels);
    bt_img_draw_scale_decomp_r(surface,decomp->root,
        decomp->min_depth,scale->elements);
}

void
iplDrawDecompositionLinearScale(matrix_p surface,bt_decomp_p decomp,
    vec_t min_value,vec_t max_value)
{
    vector_p scale;
    V_ASSERT(surface->rows == decomp->source_rows);
    V_ASSERT(surface->cols == decomp->source_cols);
    scale = vecAlloc(decomp->depth_levels);
    vecLinSpaced(scale,min_value,max_value);
    bt_img_draw_scale_decomp_r(surface,decomp->root,
        decomp->min_depth,scale->elements);
    memFree(scale);
}
