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

#include <math.h>
#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <limits.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <assert.h>

#include "os-features.h"
#include "build-index.h"
#include "boilerplate.h"
#include "errors.h"
#include "log.h"
#include "quad-utils.h"
#include "uniformize-catalog.h"
#include "startree.h"
#include "codetree.h"
#include "unpermute-quads.h"
#include "unpermute-stars.h"
#include "bl.h"
#include "ioutils.h"
#include "rdlist.h"
#include "kdtree.h"
#include "hpquads.h"
#include "sip.h"
#include "sip_qfits.h"
#include "codefile.h"
#include "codekd.h"
#include "merge-index.h"
#include "fitsioutils.h"
#include "permutedsort.h"

static void add_boilerplate(index_params_t* p, qfits_header* hdr) {
}

static int step_hpquads(index_params_t* p,
						codefile_t** p_codes, quadfile_t** p_quads,
						char** p_codefn, char** p_quadfn, 
						startree_t* starkd, const char* skdtfn,
						sl* tempfiles) {
	codefile_t* codes = NULL;
	quadfile_t* quads = NULL;
	char* quadfn = NULL;
	char* codefn = NULL;

	if (p->inmemory) {
		codes = codefile_open_in_memory();
		quads = quadfile_open_in_memory();
		if (hpquads(starkd, codes, quads, p->Nside,
					p->qlo, p->qhi, p->dimquads, p->passes, p->Nreuse, p->Nloosen,
					p->indexid, p->scanoccupied,
					p->hpquads_sort_data, p->hpquads_sort_func, p->hpquads_sort_size,
					p->args, p->argc)) {
			ERROR("hpquads failed");
			return -1;
		}
		if (!quadfile_nquads(quads)) {
			logmsg("Did not create any quads.  Perhaps your catalog does not have enough stars?\n");
			return -1;
		}
		if (quadfile_switch_to_reading(quads)) {
			ERROR("Failed to switch quadfile to read-mode");
			return -1;
		}
		if (codefile_switch_to_reading(codes)) {
			ERROR("Failed to switch codefile to read-mode");
			return -1;
		}

	} else {
		quadfn = create_temp_file("quad", p->tempdir);
		sl_append_nocopy(tempfiles, quadfn);
		codefn = create_temp_file("code", p->tempdir);
		sl_append_nocopy(tempfiles, codefn);

		if (hpquads_files(skdtfn, codefn, quadfn, p->Nside,
						  p->qlo, p->qhi, p->dimquads, p->passes, p->Nreuse, p->Nloosen,
						  p->indexid, p->scanoccupied, 
						  p->hpquads_sort_data, p->hpquads_sort_func, p->hpquads_sort_size,
						  p->args, p->argc)) {
			ERROR("hpquads failed");
			return -1;
		}
		
	}

	if (p_codes) *p_codes = codes;
	if (p_quads) *p_quads = quads;
	if (p_codefn) *p_codefn = codefn;
	if (p_quadfn) *p_quadfn = quadfn;
		
	return 0;
}

static int step_codetree(index_params_t* p,
						 codefile_t* codes, codetree_t** p_codekd,
						 const char* codefn, char** p_ckdtfn,
						 sl* tempfiles) {
	codetree_t* codekd = NULL;
	char* ckdtfn=NULL;

	if (p->inmemory) {
		logmsg("Building code kdtree from %i codes\n", codes->numcodes);
		logmsg("dim: %i\n", codefile_dimcodes(codes));
		codekd = codetree_build(codes, 0, 0, 0, 0, p->args, p->argc);
		if (!codekd) {
			ERROR("Failed to build code kdtree");
			return -1;
		}
		if (codefile_close(codes)) {
			ERROR("Failed to close codefile");
			return -1;
		}

	} else {
		ckdtfn = create_temp_file("ckdt", p->tempdir);
		sl_append_nocopy(tempfiles, ckdtfn);

		if (codetree_files(codefn, ckdtfn, 0, 0, 0, 0, p->args, p->argc)) {
			ERROR("codetree failed");
			return -1;
		}
	}
	
	if (p_codekd) *p_codekd = codekd;
	if (p_ckdtfn) *p_ckdtfn = ckdtfn;
	return 0;
}

static int step_unpermute_quads(index_params_t* p,
								quadfile_t* quads2, codetree_t* codekd,
								quadfile_t** p_quads3, codetree_t** p_codekd2,
								const char* quad2fn, const char* ckdtfn,
								char** p_quad3fn, char** p_ckdt2fn,
								sl* tempfiles) {
	quadfile_t* quads3 = NULL;
	codetree_t* codekd2 = NULL;
	char* quad3fn=NULL;
	char* ckdt2fn=NULL;

	logmsg("Unpermute-quads...\n");
	if (p->inmemory) {
		quads3 = quadfile_open_in_memory();
		if (unpermute_quads(quads2, codekd, quads3, &codekd2, p->args, p->argc)) {
			ERROR("Failed to unpermute-quads");
			return -1;
		}
		// unpermute-quads makes a shallow copy of the tree, so don't just codetree_close(codekd)...
		free(codekd->tree->perm);
		free(codekd->tree);
		codekd->tree = NULL;
		codetree_close(codekd);

		if (quadfile_switch_to_reading(quads3)) {
			ERROR("Failed to switch quads3 to read-mode");
			return -1;
		}
		if (quadfile_close(quads2)) {
			ERROR("Failed to close quadfile quads2");
			return -1;
		}

	} else {
		ckdt2fn = create_temp_file("ckdt2", p->tempdir);
		sl_append_nocopy(tempfiles, ckdt2fn);
		quad3fn = create_temp_file("quad3", p->tempdir);
		sl_append_nocopy(tempfiles, quad3fn);
		logmsg("Unpermuting quads from %s and %s to %s and %s\n", quad2fn, ckdtfn, quad3fn, ckdt2fn);
		if (unpermute_quads_files(quad2fn, ckdtfn,
								  quad3fn, ckdt2fn, p->args, p->argc)) {
			ERROR("Failed to unpermute-quads");
			return -1;
		}
	}

	if (p_quads3) *p_quads3 = quads3;
	if (p_codekd2) *p_codekd2 = codekd2;
	if (p_quad3fn) *p_quad3fn = quad3fn;
	if (p_ckdt2fn) *p_ckdt2fn = ckdt2fn;
	return 0;
}

static int step_merge_index(index_params_t* p,
							codetree_t* codekd2, quadfile_t* quads3,
							startree_t* starkd2,
							index_t** p_index,
							const char* ckdt2fn, const char* quad3fn,
							const char* skdt2fn, const char* indexfn) {
	index_t* index = NULL;

	if (p->inmemory) {
		qfits_header* hdr;

		index = index_build_from(codekd2, quads3, starkd2);
		if (!index) {
			ERROR("Failed to create index from constituent parts");
			return -1;
		}
		hdr = quadfile_get_header(index->quads);
		if (hdr)
			add_boilerplate(p, hdr);

		/* When closing:
		 kdtree_free(codekd2->tree);
		 codekd2->tree = NULL;
		 */
		*p_index = index;

	} else {
		quadfile_t* quad;
		codetree_t* code;
		startree_t* star;
		qfits_header* hdr;

		logmsg("Merging %s and %s and %s to %s\n", quad3fn, ckdt2fn, skdt2fn, indexfn);
		/*
		 if (merge_index_files(quad3fn, ckdt2fn, skdt2fn, indexfn)) {
		 ERROR("Failed to merge-index");
		 return -1;
		 }
		 */
		if (merge_index_open_files(quad3fn, ckdt2fn, skdt2fn,
								   &quad, &code, &star)) {
			ERROR("Failed to open index files for merging");
			return -1;
		}
		hdr = quadfile_get_header(quad);
		if (hdr)
			add_boilerplate(p, hdr);
		if (merge_index(quad, code, star, indexfn)) {
			ERROR("Failed to write merged index");
			return -1;
		}
		codetree_close(code);
		startree_close(star);
		quadfile_close(quad);
	}
	return 0;
}

static void step_delete_tempfiles(index_params_t* p, sl* tempfiles) {
	if (p->delete_tempfiles) {
		int i;
		for (i=0; i<sl_size(tempfiles); i++) {
			char* fn = sl_get(tempfiles, i);
			logverb("Deleting temp file %s\n", fn);
			if (unlink(fn))
				SYSERROR("Failed to delete temp file \"%s\"", fn);
		}
	}
}

int build_index_shared_skdt(const char* skdtfn,
							startree_t* starkd, index_params_t* p,
							index_t** p_index, const char* indexfn) {
	// assume we've got a final (ie, post-unpermute-stars) skdt
	// we use that skdt's stars to uniformize, along with a column pulled
	// from its tag-along data.  This yields a permutation array, which
	// we feed to hpquads to alter its idea of how to sort stars that could
	// be used to build quads.
	// Then we do unpermute-quads, skip unpermute-stars, and feed
	// the original skdt in to merge-index.
	// TODO - tweak the index-reading code to allow multiple indices
	// sharing an skdt in one FITS file.
	// --quads_X table
	// --header card in the .quads HDU saying what the name of its skdt is.
	double* sortdata = NULL;
	//int* uniperm = NULL;
	int rtn = -1;
	codefile_t* codes = NULL;
	quadfile_t* quads = NULL;
	char* quadfn=NULL;
	char* codefn=NULL;
	codetree_t* codekd = NULL;
	char* ckdtfn=NULL;

	startree_t* starkd2 = NULL;
	quadfile_t* quads2 = NULL;
	char* quad2fn=NULL;

	quadfile_t* quads3 = NULL;
	codetree_t* codekd2 = NULL;
	char* quad3fn=NULL;
	char* ckdt2fn=NULL;

	sl* tempfiles;

	if (!p->UNside)
		p->UNside = p->Nside;

	assert(p->Nside);

	if (p->inmemory && !p_index) {
		ERROR("If you set inmemory, you must set p_index");
		return -1;
	}
	if (!p->inmemory && !indexfn) {
		ERROR("If you set !inmemory, you must set indexfn");
		return -1;
	}
	assert(starkd->tree);
	assert(starkd->tree->perm == NULL);
	assert(p->sortcol);
	if (!p->sortcol) {
		ERROR("You must set the sort column\n");
		return -1;
	}

    tempfiles = sl_new(4);

	// OR, should we just strictly sort on the tag-along sortcol in hpquads?
	// Yes, probably.  Uniformization isn't going to change the ordering much
	// in deciding which stars to use in a quad -- they're all nearby.

	logverb("Grabbing tag-along column \"%s\" for sorting...\n", p->sortcol);
	sortdata = startree_get_data_column(starkd, p->sortcol, NULL, startree_N(starkd));
	if (!sortdata) {
		ERROR("Failed to find sort column data for sorting catalog");
		goto cleanup;
	}

	/*
	 logverb("Uniformizing...\n");
	 uniperm = uniformize_catalog_get_permutation(skdt, sortdata,
	 p->bighp, p->bignside, p->margin,
	 p->UNside, p->dedup, p->sweeps,
	 p->args, p->argc);
	 if (!uniperm) {
	 ERROR("Failed to find uniformization permutation array");
	 goto cleanup;
	 }
	 p->hpquads_sort = uniperm;
	 p->hpquads_sortfunc = compare_ints_asc;
	 */

	p->hpquads_sort_data = sortdata;
	p->hpquads_sort_func = (p->sortasc ? compare_doubles_asc : compare_doubles_desc);
	p->hpquads_sort_size = sizeof(double);

	// hpquads
	if (step_hpquads(p, &codes, &quads, &codefn, &quadfn,
					 starkd, skdtfn, tempfiles))
		return -1;

	// codetree
	if (step_codetree(p, codes, &codekd,
					  codefn, &ckdtfn, tempfiles))
		return -1;

	// no unpermute-stars...
	quads2 = quads;
	quad2fn = quadfn;
	starkd2 = starkd;

	// unpermute-quads...
	if (step_unpermute_quads(p, quads2, codekd, &quads3, &codekd2,
							 quad2fn, ckdtfn, &quad3fn, &ckdt2fn, tempfiles))
		return -1;

	// merge-index...
	if (step_merge_index(p, codekd2, quads3, starkd2, p_index,
						 ckdt2fn, quad3fn, skdtfn, indexfn))
		return -1;

	step_delete_tempfiles(p, tempfiles);

	sl_free2(tempfiles);

	rtn = 0;

 cleanup:
	//free(uniperm);
	free(sortdata);
	return rtn;
}

int build_index(fitstable_t* catalog, index_params_t* p,
				index_t** p_index, const char* indexfn) {

	fitstable_t* uniform;

	// star kdtree
	startree_t* starkd = NULL;
	fitstable_t* startag = NULL;

	// hpquads
	codefile_t* codes = NULL;
	quadfile_t* quads = NULL;

	// codetree
	codetree_t* codekd = NULL;

	// unpermute-stars
	startree_t* starkd2 = NULL;
	quadfile_t* quads2 = NULL;
	fitstable_t* startag2 = NULL;

	// unpermute-quads
	quadfile_t* quads3 = NULL;
	codetree_t* codekd2 = NULL;

	//index_t* index = NULL;

	sl* tempfiles;
	char* unifn=NULL;
	char* skdtfn=NULL;
	char* quadfn=NULL;
	char* codefn=NULL;
	char* ckdtfn=NULL;
	char* skdt2fn=NULL;
	char* quad2fn=NULL;
	char* quad3fn=NULL;
	char* ckdt2fn=NULL;

	if (!p->UNside)
		p->UNside = p->Nside;

	assert(p->Nside);

	if (p->inmemory && !p_index) {
		ERROR("If you set inmemory, you must set p_index");
		return -1;
	}
	if (!p->inmemory && !indexfn) {
		ERROR("If you set !inmemory, you must set indexfn");
		return -1;
	}

    tempfiles = sl_new(4);

	if (p->inmemory)
		uniform = fitstable_open_in_memory();
	else {
		unifn = create_temp_file("uniform", p->tempdir);
		sl_append_nocopy(tempfiles, unifn);
		uniform = fitstable_open_for_writing(unifn);
	}
	if (!uniform) {
		ERROR("Failed to open output table %s", unifn);
		return -1;
	}

	if (uniformize_catalog(catalog, uniform, p->racol, p->deccol,
						   p->sortcol, p->sortasc, p->brightcut,
						   p->bighp, p->bignside, p->margin,
						   p->UNside, p->dedup, p->sweeps, p->args, p->argc)) {
		return -1;
	}

	if (fitstable_fix_primary_header(uniform)) {
		ERROR("Failed to fix output table");
		return -1;
	}

	if (p->inmemory) {
		if (fitstable_switch_to_reading(uniform)) {
			ERROR("Failed to switch uniformized table to read-mode");
			return -1;
		}
	} else {
		if (fitstable_close(uniform)) {
			ERROR("Failed to close output table");
			return -1;
		}
	}
	fitstable_close(catalog);

	// startree
	if (!p->inmemory) {
		skdtfn = create_temp_file("skdt", p->tempdir);
		sl_append_nocopy(tempfiles, skdtfn);

		logverb("Reading uniformized catalog %s...\n", unifn);
		uniform = fitstable_open(unifn);
		if (!uniform) {
			ERROR("Failed to open uniformized catalog");
			return -1;
		}
	}

	// DEBUG -- print RA,Dec from uniform catalog.
	if (log_get_level() > LOG_VERB) {
		tfits_type dubl = fitscolumn_double_type();
		double* ra;
		double* dec;
		int i,N;
		ra = fitstable_read_column(uniform, p->racol, dubl);
		dec = fitstable_read_column(uniform, p->deccol, dubl);
		N = fitstable_nrows(uniform);
		logdebug("Checking %i columns of 'uniform' catalog\n", N);
		logdebug("  RA column: \"%s\"; Dec column: \"%s\"\n", p->racol, p->deccol);
		assert(ra && dec);
		for (i=0; i<N; i++)
			logdebug("  %i RA,Dec %g,%g\n", i, ra[i], dec[i]);
		free(ra);
		free(dec);
	}


	{
		int Nleaf = 25;
		int datatype = KDT_DATA_U32;
		int treetype = KDT_TREE_U32;
		int buildopts = KD_BUILD_SPLIT;

		logverb("Building star kdtree from %i stars\n", fitstable_nrows(uniform));
		starkd = startree_build(uniform, p->racol, p->deccol, datatype, treetype,
								buildopts, Nleaf, p->args, p->argc);
		if (!starkd) {
			ERROR("Failed to create star kdtree");
			return -1;
		}

		if (p->jitter > 0.0) {
			startree_set_jitter(starkd, p->jitter);
		}

		if (!p->inmemory) {
			logverb("Writing star kdtree to %s\n", skdtfn);
			if (startree_write_to_file(starkd, skdtfn)) {
				ERROR("Failed to write star kdtree");
				return -1;
			}
			startree_close(starkd);
		}

		if (startree_has_tagalong_data(uniform)) {
			logverb("Adding star kdtree tag-along data...\n");
			if (p->inmemory) {
				startag = fitstable_open_in_memory();
			} else {
				startag = fitstable_open_for_appending(skdtfn);
				if (!startag) {
					ERROR("Failed to re-open star kdtree file %s for appending", skdtfn);
					return -1;
				}
			}
			if (startree_write_tagalong_table(uniform, startag, p->racol, p->deccol, NULL, TRUE)) {
				ERROR("Failed to write tag-along table");
				return -1;
			}
			if (p->inmemory) {
				if (fitstable_switch_to_reading(startag)) {
					ERROR("Failed to switch star tag-along data to read-mode");
					return -1;
				}
				starkd->tagalong = startag;
			} else {
				if (fitstable_close(startag)) {
					ERROR("Failed to close star kdtree tag-along data");
					return -1;
				}
			}
		}
	}
	fitstable_close(uniform);

	// hpquads
	if (step_hpquads(p, &codes, &quads, &codefn, &quadfn, 
					 starkd, skdtfn,
					 tempfiles))
		return -1;

	// codetree
	if (step_codetree(p, codes, &codekd,
					  codefn, &ckdtfn, tempfiles))
		return -1;

	// unpermute-stars
	logmsg("Unpermute-stars...\n");
	if (p->inmemory) {
		quads2 = quadfile_open_in_memory();
		if (unpermute_stars(starkd, quads, &starkd2, quads2,
							TRUE, FALSE, p->args, p->argc)) {
			ERROR("Failed to unpermute-stars");
			return -1;
		}
		if (quadfile_close(quads)) {
			ERROR("Failed to close in-memory quads");
			return -1;
		}
		if (quadfile_switch_to_reading(quads2)) {
			ERROR("Failed to switch quads2 to read-mode");
			return -1;
		}
		if (startag) {
			startag2 = fitstable_open_in_memory();
			startag2->table = fits_copy_table(startag->table);
			startag2->table->nr = 0;
			startag2->header = qfits_header_copy(startag->header);
			if (unpermute_stars_tagalong(starkd, startag2)) {
				ERROR("Failed to unpermute-stars tag-along data");
				return -1;
			}
			starkd2->tagalong = startag2;
		}

		// unpermute-stars makes a shallow copy of the tree, so don't just startree_close(starkd)...
		free(starkd->tree->perm);
		free(starkd->tree);
		starkd->tree = NULL;
		startree_close(starkd);

	} else {
		skdt2fn = create_temp_file("skdt2", p->tempdir);
		sl_append_nocopy(tempfiles, skdt2fn);
		quad2fn = create_temp_file("quad2", p->tempdir);
		sl_append_nocopy(tempfiles, quad2fn);

		logmsg("Unpermuting stars from %s and %s to %s and %s\n", skdtfn, quadfn, skdt2fn, quad2fn);
		if (unpermute_stars_files(skdtfn, quadfn, skdt2fn, quad2fn,
								  TRUE, FALSE, p->args, p->argc)) {
			ERROR("Failed to unpermute-stars");
			return -1;
		}
	}

	// unpermute-quads
	if (step_unpermute_quads(p, quads2, codekd, &quads3, &codekd2,
							 quad2fn, ckdtfn, &quad3fn, &ckdt2fn, tempfiles))
		return -1;


	// index
	if (step_merge_index(p, codekd2, quads3, starkd2, p_index,
						 ckdt2fn, quad3fn, skdt2fn, indexfn))
		return -1;

	// FIXME -- close codekd2, quads3, starkd2?

	step_delete_tempfiles(p, tempfiles);

	sl_free2(tempfiles);
	return 0;
}


int build_index_files(const char* infn, int ext, const char* indexfn,
					  index_params_t* p) {
	fitstable_t* catalog;
	logmsg("Reading %s...\n", infn);
    if (ext)
        catalog = fitstable_open_extension_2(infn, ext);
    else
        catalog = fitstable_open(infn);
    if (!catalog) {
        ERROR("Couldn't read catalog %s", infn);
		return -1;
    }
	logmsg("Got %i stars\n", fitstable_nrows(catalog));

	if (p->inmemory) {
		index_t* index;
		if (build_index(catalog, p, &index, NULL)) {
			return -1;
		}
		logmsg("Writing to file %s\n", indexfn);
		if (merge_index(index->quads, index->codekd, index->starkd, indexfn)) {
			ERROR("Failed to write index file");
			return -1;
		}
		kdtree_free(index->codekd->tree);
		index->codekd->tree = NULL;
		index_close(index);

	} else {
		if (build_index(catalog, p, NULL, indexfn)) {
			return -1;
		}
	}

	return 0;
}

int build_index_shared_skdt_files(const char* starkdfn, const char* indexfn,
								  index_params_t* p) {
	startree_t* skdt = NULL;

	logmsg("Reading %s...\n", starkdfn);
	skdt = startree_open(starkdfn);
    if (!skdt) {
        ERROR("Couldn't read star kdtree from \"%s\"", starkdfn);
		return -1;
    }
	logmsg("Got %i stars\n", startree_N(skdt));

	if (p->inmemory) {
		index_t* index;
		if (build_index_shared_skdt(starkdfn, skdt, p, &index, NULL)) {
			return -1;
		}
		logmsg("Writing to file %s\n", indexfn);
		if (merge_index(index->quads, index->codekd, index->starkd, indexfn)) {
			ERROR("Failed to write index file \"%s\"", indexfn);
			return -1;
		}
		// FIXME?  Why close codekd independently?
		kdtree_free(index->codekd->tree);
		index->codekd->tree = NULL;
		index_close(index);

	} else {
		if (build_index_shared_skdt(starkdfn, skdt, p, NULL, indexfn)) {
			return -1;
		}
	}
	return 0;
}


void build_index_defaults(index_params_t* p) {
	memset(p, 0, sizeof(index_params_t));
	p->sweeps = 10;
	p->racol = "RA";
	p->deccol = "DEC";
	p->passes = 16;
	p->Nreuse = 8;
	p->Nloosen = 20;
	p->dimquads = 4;
	p->sortasc = TRUE;
	p->brightcut = -HUGE_VAL;
	// default to all-sky
	p->bighp = -1;
	//p->inmemory = TRUE;
	p->delete_tempfiles = TRUE;
	p->tempdir = "/tmp";
}