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

#include <assert.h>
#include <math.h>
#include <string.h>
#include <stdint.h>

#include "verify2.h"
#include "permutedsort.h"
#include "mathutil.h"
#include "keywords.h"
#include "log.h"
#include "sip-utils.h"
#include "healpix.h"

#define DEBUGVERIFY 1
#if DEBUGVERIFY
//#define debug(args...) fprintf(stderr, args)
#else
#define debug(args...)
#endif

#include "fitsioutils.h"
#include "errors.h"
#include "cairoutils.h"	

static void add_gaussian_to_image(double* img, int W, int H,
                                  double cx, double cy, double sigma,
                                  double scale,
                                  double nsigma, int boundary) {
    int x, y;
    if (boundary == 0) {
        // truncate.
        for (y = MAX(0, cy - nsigma*sigma); y <= MIN(H-1, cy + nsigma * sigma); y++) {
            for (x = MAX(0, cx - nsigma*sigma); x <= MIN(W-1, cx + nsigma * sigma); x++) {
                img[y*W + x] += scale * exp(-(square(y-cy)+square(x-cx)) / (2.0 * square(sigma)));
            }
        }

    } else if (boundary == 1) {
        // mirror.
        int mx, my;
        for (y=MAX(-(H-1), floor(cy - nsigma * sigma));
             y<=MIN(2*H-1, ceil(cy + nsigma * sigma)); y++) {
            if (y < 0)
                my = -1 - y;
            else if (y >= H)
                my = 2*H - 1 - y;
            else
                my = y;
            for (x=MAX(-(W-1), floor(cx - nsigma * sigma));
                 x<=MIN(2*W-1, ceil(cx + nsigma * sigma)); x++) {
                if (x < 0)
                    mx = -1 - x;
                else if (x >= W)
                    mx = 2*W - 1 - x;
                else
                    mx = x;
                img[my*W + mx] += scale * exp(-(square(y-cy)+square(x-cx)) / (2.0 * square(sigma)));
            }
        }
    }
}

void verify_get_all_matches(const double* refxys, int NR,
                            const double* testxys, const double* testsigma2s, int NT,
                            double effective_area,
                            double distractors,
                            double nsigma,
                            double limit,
                            il*** p_reflist,
                            dl*** p_problist) {
    double* refcopy;
    kdtree_t* rtree;
    int Nleaf = 10;
    int i,j;
    double logd;
    double logbg;
    double loglimit;

    il** reflist;
    dl** problist;

    reflist  = calloc(NT, sizeof(il*));
    problist = calloc(NT, sizeof(dl*));

    // Build a tree out of the index stars in pixel space...
    // kdtree scrambles the data array so make a copy first.
    refcopy = malloc(2 * NR * sizeof(double));
    memcpy(refcopy, refxys, 2 * NR * sizeof(double));
    rtree = kdtree_build(NULL, refcopy, NR, 2, Nleaf, KDTT_DOUBLE, KD_BUILD_SPLIT);

    logbg = log(1.0 / effective_area);
    logd  = log(distractors / effective_area);
    loglimit = log(distractors / effective_area * limit);

    for (i=0; i<NT; i++) {
        const double* testxy;
        double sig2;
        kdtree_qres_t* res;

        testxy = testxys + 2*i;
        sig2 = testsigma2s[i];

        logverb("\n");
        logverb("test star %i: (%.1f,%.1f), sigma: %.1f\n", i, testxy[0], testxy[1], sqrt(sig2));

        // find all ref stars within nsigma.
        res = kdtree_rangesearch_options(rtree, testxy, sig2*nsigma*nsigma,
                                         KD_OPTIONS_SORT_DISTS | KD_OPTIONS_SMALL_RADIUS);

        if (res->nres == 0) {
            kdtree_free_query(res);
            continue;
        }

        reflist[i] = il_new(4);
        problist[i] = dl_new(4);

        for (j=0; j<res->nres; j++) {
            double d2;
            int refi;
            double loggmax, logfg;

            d2 = res->sdists[j];
            refi = res->inds[j];

            // peak value of the Gaussian
            loggmax = log((1.0 - distractors) / (2.0 * M_PI * sig2 * NR));
            // value of the Gaussian
            logfg = loggmax - d2 / (2.0 * sig2);

            if (logfg < loglimit)
                continue;

            logverb("  ref star %i, dist %.2f, sigmas: %.3f, logfg: %.1f (%.1f above distractor, %.1f above bg, %.1f above keep-limit)\n",
                    refi, sqrt(d2), sqrt(d2 / sig2), logfg, logfg - logd, logfg - logbg, logfg - loglimit);

            il_append(reflist[i], refi);
            dl_append(problist[i], logfg);
        }

        kdtree_free_query(res);
    }
    kdtree_free(rtree);
    free(refcopy);

    *p_reflist  = reflist;
    *p_problist = problist;
}


/*
 unsigned char* img = malloc(4*W*H);

 // draw images of index and field densities.
 double* idensity = calloc(W * H, sizeof(double));
 double* fdensity = calloc(W * H, sizeof(double));

 double iscale = 2. * sqrt((double)(W * H) / (NI * M_PI));
 double fscale = 2. * sqrt((double)(W * H) / (NF * M_PI));
 logverb("NI = %i; iscale = %g\n", NI, iscale);
 logverb("NF = %i; fscale = %g\n", NF, fscale);
 logverb("computing density images...\n");
 for (i=0; i<NI; i++)
 add_gaussian_to_image(idensity, W, H, indexpix[i*2 + 0], indexpix[i*2 + 1], iscale, 1.0, 3.0, 1);
 for (i=0; i<NF; i++)
 add_gaussian_to_image(fdensity, W, H, starxy_getx(vf->field, i), starxy_gety(vf->field, i), fscale, 1.0, 3.0, 1);

 double idmax=0, fdmax=0;
 for (i=0; i<(W*H); i++) {
 idmax = MAX(idmax, idensity[i]);
 fdmax = MAX(fdmax, fdensity[i]);
 }
 for (i=0; i<(W*H); i++) {
 unsigned char val = 255.5 * idensity[i] / idmax;
 img[i*4+0] = val;
 img[i*4+1] = val;
 img[i*4+2] = val;
 img[i*4+3] = 255;
 }
 cairoutils_write_png("idensity.png", img, W, H);
 for (i=0; i<(W*H); i++) {
 unsigned char val = 255.5 * fdensity[i] / fdmax;
 img[i*4+0] = val;
 img[i*4+1] = val;
 img[i*4+2] = val;
 img[i*4+3] = 255;
 }
 cairoutils_write_png("fdensity.png", img, W, H);

 free(idensity);
 free(fdensity);
 */


/*
 // index star weights.
 double* iweights = malloc(NI * sizeof(double));
 for (i=0; i<NI; i++)
 iweights[i] = 1.0;
 // don't count index stars that are part of the matched quad.
 for (i=0; i<NI; i++)
 for (j=0; j<dimquads; j++)
 if (starids[i] == mo->star[j])
 iweights[i] = 0.0;

 double ranksigma = 20.0;
 int Nrankprobs = (int)(ranksigma * 5);
 double* rankprobs = malloc(Nrankprobs * sizeof(double));
 for (i=0; i<Nrankprobs; i++)
 rankprobs[i] = exp(-(double)(i*i) / (2. * ranksigma * ranksigma));

 double qc[2], qr2;

 get_quad_center(vf, mo, qc, &qr2);
 */

/*
 // create the probability distribution map for this field star.
 double* pmap = malloc(W * H * sizeof(double));
 // background rate...
 for (j=0; j<(W*H); j++)
 pmap[j] = distractors / (fieldW*fieldH);

 double normrankprob;
 int dr;
 normrankprob = 0.0;
 for (j=0; j<NI; j++) {
 dr = abs(i - j);
 if (dr >= Nrankprobs)
 continue;
 normrankprob += rankprobs[dr];
 }
 for (j=0; j<NI; j++) {
 double r2, sig2, sig;

 dr = abs(i - j);
 if (dr >= Nrankprobs)
 continue;

 r2 = distsq(indexpix + j*2, qc, 2);
 sig2 = get_sigma2_at_radius(verify_pix2, r2, qr2);
 sig = sqrt(sig2);

 for (y = MAX(0, indexpix[j*2 + 1] - 5.0*sig);
 y <= MIN(H-1, indexpix[j*2 + 1] + 5.0 * sig); y++) {
 for (x = MAX(0, indexpix[j*2 + 0] - 5.0*sig);
 x <= MIN(W-1, indexpix[j*2 + 0] + 5.0 * sig); x++) {
 pmap[y*W + x] += iweights[j] * (rankprobs[dr] / normrankprob) *
 exp(-(square(indexpix[j*2+0]-x) + square(indexpix[j*2+1]-y)) / (2.0 * sig2)) *
 (1.0 / (2.0 * M_PI * sig2)) * (1.0 - distractors);
 }
 }
 }

 double maxval = 0.0;
 for (j=0; j<(W*H); j++)
 maxval = MAX(maxval, pmap[j]);

 double psum = 0.0;
 for (j=0; j<(W*H); j++)
 psum += pmap[j];
 logverb("Probability sum: %g\n", psum);

 char* fn;

 printf("Writing probability image %i\n", i);
 printf("maxval = %g\n", maxval);

 for (j=0; j<(W*H); j++) {
 unsigned char val = (unsigned char)MAX(0, MIN(255, 255.5 * sqrt(pmap[j] / maxval)));
 img[4*j + 0] = val;
 img[4*j + 1] = val;
 img[4*j + 2] = val;
 img[4*j + 3] = 255;
 }

 free(pmap);

 cairo_t* cairo;
 cairo_surface_t* target;

 target = cairo_image_surface_create_for_data(img, CAIRO_FORMAT_ARGB32, W, H, W*4);
 cairo = cairo_create(target);
 cairo_set_line_join(cairo, CAIRO_LINE_JOIN_ROUND);
 cairo_set_antialias(cairo, CAIRO_ANTIALIAS_GRAY);

 // draw the quad.
 cairo_set_source_rgba(cairo, 0.0, 0.0, 1.0, 0.5);

 double starxy[DQMAX*2];
 double angles[DQMAX];
 int perm[DQMAX];
 int DQ = dimquads;

 for (k=0; k<DQ; k++)
 starxy_get(vf->field, mo->field[k], starxy + k*2);
 for (k=0; k<DQ; k++)
 angles[k] = atan2(starxy[k*2 + 1] - qc[1], starxy[k*2 + 0] - qc[0]);
 permutation_init(perm, DQ);
 permuted_sort(angles, sizeof(double), compare_doubles_asc, perm, DQ);
 for (k=0; k<DQ; k++) {
 double px, py;
 px = starxy[perm[k]*2+0];
 py = starxy[perm[k]*2+1];
 if (k == 0)
 cairo_move_to(cairo, px, py);
 else
 cairo_line_to(cairo, px, py);
 }
 cairo_close_path(cairo);
 cairo_stroke(cairo);

 // draw the source point.
 cairo_set_source_rgb(cairo, 1.0, 0.0, 0.0);
 cairoutils_draw_marker(cairo, CAIROUTIL_MARKER_CROSSHAIR, fxy[0], fxy[1], 3.0);
 cairo_stroke(cairo);

 */


/*
 // find index stars within n sigma...
 double nsigma = 5.0;
 double r2 = sigma2s[i] * nsigma*nsigma;
 int opts = KD_OPTIONS_COMPUTE_DISTS;

 kdtree_qres_t* res = kdtree_rangesearch_options(itree, fxy, r2, opts);

 logverb("found %i index stars within range.\n", res->nres);

 if (res->nres == 0)
 goto nextstar;

 double bestprob = 0.0;
 int bestarg = -1;

 for (j=0; j<res->nres; j++) {
 int ii = res->inds[j];
 dr = abs(i - ii);
 logverb("  index star %i: rank diff %i\n", ii, i-ii);
 if (dr >= Nrankprobs)
 continue;
 double prob = iweights[ii] * rankprobs[dr] *
 exp(-res->sdists[j] / (2.0 * sigma2s[i]));
 logverb("  -> prob %g\n", prob);
 if (prob > bestprob) {
 bestprob = prob;
 bestarg = j;
 }
 }

 if (bestarg == -1) {
 // FIXME -- distractor?
 goto nextstar;
 }

 int besti = res->inds[bestarg];
 double lostweight = bestprob / iweights[besti]; 
 // exp(-res->sdists[bestarg] / (2.0 * sigma2s[i]));
 iweights[besti] = MAX(0.0, iweights[besti] - lostweight);
 logverb("matched index star %i: dropping weight by %g; new weight %g.\n", besti, lostweight, iweights[besti]);
		
 bestprob *= (1.0 / (2.0 * M_PI * sigma2s[i])) / normrankprob * (1.0 - distractors);
 logverb("bestprob %g, vs background %g\n", bestprob, (1.0 / (double)(W * H)));
		
 kdtree_free_query(res);

 cairo_set_source_rgb(cairo, 0.0, 1.0, 0.0);
 cairoutils_draw_marker(cairo, CAIROUTIL_MARKER_CIRCLE, indexpix[2*besti+0], indexpix[2*besti+1], 3.0);
 cairo_stroke(cairo);


 nextstar:
 cairoutils_argb32_to_rgba(img, W, H);
 cairo_surface_destroy(target);
 cairo_destroy(cairo);
 asprintf(&fn, "logprob-%04i.png", i);
 cairoutils_write_png(fn, img, W, H);
 free(img);


 if (i >= 9)
 break;
 }

 kdtree_free(itree);
 free(ipixcopy);

 */







