/*
 # This file is part of the Astrometry.net suite.
 # Licensed under a 3-clause BSD style license - see LICENSE
 */

/**
 Reads a .code or .ckdt file, projects each code onto each pair of axes,
 and histograms the results.  Writes out the histograms as Matlab literals.

 Pipe the output to a file like "hists.m", then in Matlab run the
 "codeprojections.m" script.

 HACK - I haven't looked at how code dimensionality (dimcodes)
 influences the "volume_at_value()" function.  The volume-corrected plots
 may therefore be wrong.
 */

#include <string.h>
#include <limits.h>
#include <math.h>

#include "starutil.h"
#include "codekd.h"
#include "kdtree_fits_io.h"
#include "keywords.h"
#include "boilerplate.h"

#define OPTIONS "hd"


static void print_help(char* progname)
{
    BOILERPLATE_HELP_HEADER(stderr);
    fprintf(stderr, "Usage: %s  <code kdtree>\n"
            "       [-d]: normalize by volume (produce density plots)\n\n",
            progname);
}

// 2-D hists
int** hists = NULL;
double** dhists = NULL;
int Nbins = 40;
int Dims;

// 2-D hist of {C,D}x,{C,D}y
int* xyhist = NULL;
double* dxyhist = NULL;

// 1-D hists
int* single = NULL;
double* dsingle = NULL;
int Nsingle = 100;

anbool do_density = FALSE;

double minvalue;
double scale;

static Const double volume_at_value(double x) {
    // codes in a circle live inside the circle
    //    (x-1/2)^2 + (y-1/2)^2 = 1/2
    // we are given "x" and want to find the distance
    // between the upper and lower arcs of the circle;
    // ie y(x)_upper - y(x)_lower.  Hence we don't care
    // about the y offset of the center of the circle and
    // we want twice the value y(x)_upper.  Ie, solve
    //    (x-1/2)^2 + y^2 = 1/2
    // for y, and return twice that.
    //    y = sqrt(1/2 - (x - 1/2)^2).
    //      = sqrt(1/2 - (x^2 - x + 1/4)
    //      = sqrt(-x^2 + x + 1/4)
    return 2.0 * sqrt(-x*x + x + 0.25);
}

static int value_to_bin(double val, int Nbins) {
    int bin = (int)((val - minvalue) * scale * Nbins);
    if (bin >= Nbins) {
        bin = Nbins-1;
        printf("truncating value %g\n", val);
    }
    if (bin < 0) {
        bin = 0;
        printf("truncating (up) value %g\n", val);
    }
    return bin;
}

static void add_to_single_histogram(int dim, double val) {
    int* hist = single + Nsingle * dim;
    int bin = value_to_bin(val, Nsingle);
    hist[bin]++;
    if (do_density) {
        double* dhist = dsingle + Nsingle * dim;
        dhist[bin] += 1.0 / volume_at_value(val);
    }
}

static void add_to_histogram(int dim1, int dim2, double val1, double val2) {
    int xbin, ybin;
    int* hist = hists[dim1 * Dims + dim2];
    xbin = value_to_bin(val1, Nbins);
    ybin = value_to_bin(val2, Nbins);
    hist[xbin * Nbins + ybin]++;
    if (do_density) {
        double* dhist = dhists[dim1 * Dims + dim2];
        double inc;
        if (dim1/2 == dim2/2)
            // (cx vs cy) or (dx vs dy); the other two dimensions are independent.
            inc = 1.0;
        else
            inc = 1.0 / (volume_at_value(val1) * volume_at_value(val2));
        dhist[xbin * Nbins + ybin] += inc;
    }
}

static void add_to_cd_histogram(double val1, double val2) {
    int xbin, ybin;
    xbin = value_to_bin(val1, Nbins);
    ybin = value_to_bin(val2, Nbins);
    xyhist[xbin * Nbins + ybin]++;
    if (do_density)
        dxyhist[xbin * Nbins + ybin] += 1.0 / (volume_at_value(val1) * volume_at_value(val2));
}

int main(int argc, char *argv[])
{
    int argchar;
    char *ckdtfname = NULL;
    int i, j, d, e;
    anbool circle;
    codetree* ct = NULL;
    kdtree_t* ckdt = NULL;
    int Ncodes;
    int dimcodes;

    if (argc <= 2) {
        print_help(argv[0]);
        return (OPT_ERR);
    }

    while ((argchar = getopt (argc, argv, OPTIONS)) != -1)
        switch (argchar) {
        case 'd':
            do_density = TRUE;
            break;
        case 'h':
            print_help(argv[0]);
            return (HELP_ERR);
        default:
            return (OPT_ERR);
        }

    if (optind != argc-1) {
        print_help(argv[0]);
        printf("You must give a code kdtree filename!\n");
        exit(-1);
    }
    ckdtfname = argv[optind];

    ct = codetree_open(ckdtfname);
    if (!ct) {
        fprintf(stderr, "Failed to read code kdtree file %s.\n", ckdtfname);
        exit(-1);
    }
    circle = qfits_header_getboolean(ct->header, "CIRCLE", 0);
    ckdt = ct->tree;
    Ncodes = ckdt->ndata;
    dimcodes = ckdt->ndim;

    fprintf(stderr, "Index %s the CIRCLE property.\n",
            (circle ? "has" : "does not have"));

    if (circle) {
        double margin = 0.1;
        minvalue = 0.5 - M_SQRT1_2 - (0.5 * margin);
        //scale = M_SQRT1_2 + margin;
        scale = 1.0 / (M_SQRT2 + margin);
    } else {
        double margin = 0.06;
        minvalue = 0.0 - (0.5 * margin);
        scale = 1.0 / (1.0 + margin);

        if (do_density) {
            fprintf(stderr, "Warning: this index does not have the CIRCLE property "
                    "so the -d flag has no effect.\n");
            do_density = FALSE;
        }
    }

    // Allocate memory for projection histograms
    hists  = calloc(dimcodes * dimcodes, sizeof(int*));
    dhists = calloc(dimcodes * dimcodes, sizeof(double*));

    for (d = 0; d < dimcodes; d++) {
        for (e = 0; e < d; e++) {
            hists [d*dimcodes + e] = calloc(Nbins * Nbins, sizeof(int));
            dhists[d*dimcodes + e] = calloc(Nbins * Nbins, sizeof(double));
        }
        // Since the 4x4 matrix of histograms is actually symmetric,
        // only make half
        for (; e < dimcodes; e++) {
            hists [d*dimcodes + e] = NULL;
            dhists[d*dimcodes + e] = NULL;
        }
    }

    xyhist  = calloc(Nbins * Nbins, sizeof(int));
    dxyhist = calloc(Nbins * Nbins, sizeof(double));

    single  = calloc(dimcodes * Nsingle, sizeof(int));
    dsingle = calloc(dimcodes * Nsingle, sizeof(double));

    for (i=0; i<Ncodes; i++) {
        double code[dimcodes];

        codetree_get(ct, i, code);

        for (d = 0; d < dimcodes; d++) {
            for (e = 0; e < d; e++)
                add_to_histogram(d, e, code[d], code[e]);
            add_to_single_histogram(d, code[d]);
        }
        for (d=0; d<dimcodes/2; d++)
            add_to_cd_histogram(code[2*d], code[2*d+1]);
    }

    codetree_close(ct);

    for (d = 0; d < dimcodes; d++) {
        for (e = 0; e < d; e++) {
            int* hist;
            printf("hist_%i_%i=zeros([%i,%i]);\n",
                   d, e, Nbins, Nbins);
            hist = hists[d * dimcodes + e];
            for (i = 0; i < Nbins; i++) {
                int j;
                printf("hist_%i_%i(%i,:)=[", d, e, i + 1);
                for (j = 0; j < Nbins; j++) {
                    printf("%i,", hist[i*Nbins + j]);
                }
                printf("];\n");
            }
            free(hist);
            if (do_density) {
                double* dhist;
                printf("dhist_%i_%i=zeros([%i,%i]);\n",
                       d, e, Nbins, Nbins);
                dhist = dhists[d * dimcodes + e];
                for (i = 0; i < Nbins; i++) {
                    printf("dhist_%i_%i(%i,:)=[", d, e, i + 1);
                    for (j = 0; j < Nbins; j++)
                        printf("%g,", dhist[i*Nbins + j]);
                    printf("];\n");
                }
                free(dhist);
            }
        }
        printf("hist_%i=[", d);
        for (i = 0; i < Nsingle; i++)
            printf("%i,", single[d*Nsingle + i]);
        printf("];\n");

        if (do_density) {
            printf("dhist_%i=[", d);
            for (i = 0; i < Nsingle; i++)
                printf("%g,", dsingle[d*Nsingle + i]);
            printf("];\n");
        }
    }
    printf("hist_xy=[");
    for (i=0; i<Nbins; i++) {
        for (j=0; j<Nbins; j++)
            printf("%i,", xyhist[i*Nbins+j]);
        printf(";");
    }
    printf("];\n");
    if (do_density) {
        printf("dhist_xy=[");
        for (i=0; i<Nbins; i++) {
            for (j=0; j<Nbins; j++)
                printf("%g,", dxyhist[i*Nbins+j]);
            printf(";");
        }
        printf("];\n");
    }

    free(xyhist);
    free(hists);
    free(single);
    if (do_density) {
        free(dxyhist);
        free(dhists);
        free(dsingle);
    }

    fprintf(stderr, "Done!\n");

    return 0;
}