/*
  This file is part of the Astrometry.net suite.
  Copyright 2006, 2007 Dustin Lang, Keir Mierle and Sam Roweis.

  The Astrometry.net suite is free software; you can redistribute
  it and/or modify it under the terms of the GNU General Public License
  as published by the Free Software Foundation, version 2.

  The Astrometry.net suite is distributed in the hope that it will be
  useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with the Astrometry.net suite ; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
*/

#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <sys/param.h>

#include "catalog.h"
#include "fitsioutils.h"
#include "ioutils.h"
#include "starutil.h"
#include "mathutil.h"
#include "errors.h"

#define CHUNK_XYZ    0
#define CHUNK_MAG    1
#define CHUNK_MAG_ERR 2
#define CHUNK_SIG    3
#define CHUNK_PM     4
#define CHUNK_SIGPM  5
#define CHUNK_STARID 6

static fitsbin_chunk_t* xyz_chunk(catalog* cat) {
    return fitsbin_get_chunk(cat->fb, CHUNK_XYZ);
}
static fitsbin_chunk_t* mag_chunk(catalog* cat) {
    return fitsbin_get_chunk(cat->fb, CHUNK_MAG);
}
static fitsbin_chunk_t* mag_err_chunk(catalog* cat) {
    return fitsbin_get_chunk(cat->fb, CHUNK_MAG_ERR);
}
static fitsbin_chunk_t* get_chunk(catalog* cat, int i) {
    return fitsbin_get_chunk(cat->fb, i);
}

static int callback_read_header(fitsbin_t* fb, fitsbin_chunk_t* chunk) {
    qfits_header* primheader = fitsbin_get_primary_header(fb);
	catalog* cat = chunk->userdata;
	cat->numstars = qfits_header_getint(primheader, "NSTARS", -1);
	cat->healpix = qfits_header_getint(primheader, "HEALPIX", -1);
	cat->hpnside = qfits_header_getint(primheader, "HPNSIDE", 1);
	if (fits_check_endian(primheader)) {
		ERROR("Catalog file was written with wrong endianness");
        return -1;
	}
	if (cat->numstars == -1) {
		ERROR("Couldn't find NSTARS header in catalog file.");
        return -1;
	}

    chunk->nrows = cat->numstars;
	return 0;
}

static int callback_read_tagalong(fitsbin_t* fb, fitsbin_chunk_t* chunk) {
	catalog* cat = chunk->userdata;
    chunk->nrows = cat->numstars;
	return 0;
}

static catalog* new_catalog(const char* fn, anbool writing) {
	catalog* cat;
	fitsbin_chunk_t chunk;

    cat = calloc(1, sizeof(catalog));
	if (!cat) {
		fprintf(stderr, "catalog_open: malloc failed.\n");
	}

    if (writing)
        cat->fb = fitsbin_open_for_writing(fn);
    else
        cat->fb = fitsbin_open(fn);
    if (!cat->fb) {
        ERROR("Failed to create fitsbin");
        return NULL;
    }

    memset(&chunk, 0, sizeof(fitsbin_chunk_t));

    // NOTE -- the order these are added MUST match the CHUNK_XYZ, CHUNK_MAG, etc
    // ordering.

    // Star positions
    fitsbin_chunk_init(&chunk);
    chunk.tablename = "xyz";
    chunk.required = 1;
    chunk.callback_read_header = callback_read_header;
    chunk.userdata = cat;
	chunk.itemsize = DIM_STARS * sizeof(double);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Star magnitudes
    chunk.tablename = "mag";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = sizeof(float);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Star magnitude errors
    chunk.tablename = "mag_err";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = sizeof(float);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Sigmas
    chunk.tablename = "sigma_radec";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = 2 * sizeof(float);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Motions
    chunk.tablename = "proper_motion";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = 2 * sizeof(float);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Sigma Motions
    chunk.tablename = "sigma_pm";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = 2 * sizeof(float);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_reset(&chunk);

    // Ids
    chunk.tablename = "starid";
    chunk.required = 0;
    chunk.callback_read_header = callback_read_tagalong;
    chunk.userdata = cat;
    chunk.itemsize = sizeof(uint64_t);
    fitsbin_add_chunk(cat->fb, &chunk);
    fitsbin_chunk_clean(&chunk);

    return cat;
}

anbool catalog_has_mag(const catalog* cat) {
    return (cat->mag != NULL);
}

catalog* catalog_open(char* fn) {
    catalog* cat = NULL;

    cat = new_catalog(fn, FALSE);
    if (!cat)
        goto bailout;
    if (fitsbin_read(cat->fb)) {
        ERROR("catalog: fitsbin_read() failed");
		goto bailout;
	}
	cat->stars   = xyz_chunk(cat)->data;
    cat->mag     = mag_chunk(cat)->data;
    cat->mag_err = mag_err_chunk(cat)->data;
    cat->sigma_radec   = get_chunk(cat, CHUNK_SIG   )->data;
    cat->proper_motion = get_chunk(cat, CHUNK_PM    )->data;
    cat->sigma_pm      = get_chunk(cat, CHUNK_SIGPM )->data;
    cat->starids       = get_chunk(cat, CHUNK_STARID)->data;
    return cat;

 bailout:
    if (cat)
        catalog_close(cat);
    return NULL;
}

static void add_default_header_vals(catalog* cat) {
	qfits_header* hdr;
	hdr = catalog_get_header(cat);
	qfits_header_add(hdr, "AN_FILE", "OBJS", "This file has a list of object positions.", NULL);
    fits_add_endian(hdr);
	fits_add_double_size(hdr);
	qfits_header_add(hdr, "NSTARS", "0", "Number of stars in this file.", NULL);
	qfits_header_add(hdr, "HEALPIX", "-1", "Healpix covered by this catalog, with Nside=HPNSIDE", NULL);
	qfits_header_add(hdr, "HPNSIDE", "-1", "Nside of HEALPIX.", NULL);
	qfits_header_add(hdr, "COMMENT", "This is a flat array of XYZ for each catalog star.", NULL, NULL);
	qfits_header_add(hdr, "COMMENT", "  (ie, star position on the unit sphere)", NULL, NULL);
	qfits_header_add(hdr, "COMMENT", "  (stored as three native-{endian,size} doubles)", NULL, NULL);
}

catalog* catalog_open_for_writing(char* fn)  {
	catalog* cat;

	cat = new_catalog(fn, TRUE);
	if (!cat)
		goto bailout;
    cat->hpnside = 1;
    add_default_header_vals(cat);
	return cat;

 bailout:
	if (cat)
        catalog_close(cat);
	return NULL;
}

qfits_header* catalog_get_header(catalog* cat) {
    return fitsbin_get_primary_header(cat->fb);
}

int catalog_write_header(catalog* cat) {
	fitsbin_t* fb = cat->fb;

	if (fitsbin_write_primary_header(fb) ||
		fitsbin_write_chunk_header(fb, xyz_chunk(cat))) {
		ERROR("Failed to write catalog header");
		return -1;
	}
	return 0;
}

int catalog_fix_header(catalog* cat) {
	qfits_header* hdr;
	fitsbin_t* fb = cat->fb;

	hdr = catalog_get_header(cat);
	// fill in the real values...
    fits_header_mod_int(hdr, "NSTARS", cat->numstars, "Number of stars.");
    fits_header_mod_int(hdr, "HEALPIX", cat->healpix, "Healpix covered by this catalog, with Nside=HPNSIDE");
    fits_header_mod_int(hdr, "HPNSIDE", cat->hpnside, "Nside of HEALPIX.");

	if (fitsbin_fix_primary_header(fb) ||
		fitsbin_fix_chunk_header(fb, xyz_chunk(cat))) {
        ERROR("Failed to fix catalog header");
		return -1;
	}
	return 0;
}

double* catalog_get_base(catalog* cat) {
	return cat->stars;
}

double* catalog_get_star(catalog* cat, int sid) {
	if (sid >= cat->numstars) {
		fflush(stdout);
		fprintf(stderr, "catalog: asked for star %u, but catalog size is only %u.\n",
		        sid, cat->numstars);
		return NULL;
	}
	return cat->stars + sid * 3;
}

int catalog_write_star(catalog* cat, double* star) {
	if (fitsbin_write_item(cat->fb, xyz_chunk(cat), star)) {
		fprintf(stderr, "Failed to write catalog data to file %s: %s\n",
                cat->fb->filename, strerror(errno));
		return -1;
	}
	cat->numstars++;
	return 0;
}

int write_floats(catalog* cat, int chunknum, 
                  const char* name, fl* list, int nperstar) {
    int i;
    int B = 1000;
    fitsbin_chunk_t* chunk = get_chunk(cat, chunknum);
    if (!list || (fl_size(list) != cat->numstars * nperstar)) {
        ERROR("Number of %ss (%i) doesn't match number of stars (%i)",
              name, list ? fl_size(list) : 0, cat->numstars);
        return -1;
    }

    if (fitsbin_write_chunk_header(cat->fb, chunk)) {
        ERROR("Failed to write %ss header", name);
        return -1;
    }
    for (i=0; i<cat->numstars; i+=B) {
        float data[nperstar * B];
        int n = MIN(i+B, cat->numstars) - i;
        fl_copy(list, i*nperstar, nperstar*n, data);
        if (fitsbin_write_items(cat->fb, chunk, data, n)) {
            ERROR("Failed to write %s for stars %i to %i", name, i, i+n-1);
            return -1;
        }
    }
    if (fitsbin_fix_chunk_header(cat->fb, chunk)) {
        ERROR("Failed to fix %ss header", name);
        return -1;
    }
	return 0;
}

int catalog_write_mags(catalog* cat) {
    return write_floats(cat, CHUNK_MAG, "magnitude", cat->maglist, 1);
}

int catalog_write_mag_errs(catalog* cat) {
    return write_floats(cat, CHUNK_MAG_ERR, "magnitude errors", cat->magerrlist, 1);
}

int catalog_write_sigmas(catalog* cat) {
    return write_floats(cat, CHUNK_SIG, "sigma", cat->siglist, 2);
}

int catalog_write_pms(catalog* cat) {
    return write_floats(cat, CHUNK_PM, "proper motion", cat->pmlist, 2);
}

int catalog_write_sigma_pms(catalog* cat) {
    return write_floats(cat, CHUNK_SIGPM, "sigma proper motion", cat->sigpmlist, 2);
}

int catalog_write_ids(catalog* cat) {
    int i;
    char* name = "id";
    int chunknum = CHUNK_STARID;
    fitsbin_chunk_t* chunk = get_chunk(cat, chunknum);
    if (!cat->idlist || (bl_size(cat->idlist) != cat->numstars)) {
        ERROR("Number of %ss (%i) doesn't match number of stars (%i)",
              name, cat->idlist ? bl_size(cat->idlist) : 0, cat->numstars);
        return -1;
    }

    if (fitsbin_write_chunk_header(cat->fb, chunk)) {
        ERROR("Failed to write %ss header", name);
        return -1;
    }
    for (i=0; i<cat->numstars; i++)
        if (fitsbin_write_item(cat->fb, chunk, bl_access(cat->idlist, i))) {
            ERROR("Failed to write %s for star %i", name, i);
            return -1;
        }
    if (fitsbin_fix_chunk_header(cat->fb, chunk)) {
        ERROR("Failed to fix %ss header", name);
        return -1;
    }
	return 0;
}

void addfloat(fl** list, float val) {
    if (!(*list))
        *list = fl_new(256);
    fl_append(*list, val);
}

void catalog_add_mag(catalog* cat, float mag) {
    addfloat(&(cat->maglist), mag);
}
void catalog_add_mag_err(catalog* cat, float magerr) {
    addfloat(&(cat->magerrlist), magerr);
}
void catalog_add_sigmas(catalog* cat, float sra, float sdec) {
    addfloat(&(cat->siglist), sra);
    addfloat(&(cat->siglist), sdec);
}
void catalog_add_pms(catalog* cat, float sra, float sdec) {
    addfloat(&(cat->pmlist), sra);
    addfloat(&(cat->pmlist), sdec);
}
void catalog_add_sigma_pms(catalog* cat, float sra, float sdec) {
    addfloat(&(cat->sigpmlist), sra);
    addfloat(&(cat->sigpmlist), sdec);
}
void catalog_add_id(catalog* cat, uint64_t id) {
    if (!cat->idlist)
        cat->idlist = bl_new(256, sizeof(uint64_t));
    bl_append(cat->idlist, &id);
}

int catalog_close(catalog* cat) {
    int rtn;
	if (!cat) return 0;
	rtn = fitsbin_close(cat->fb);
    fl_free(cat->maglist);
    fl_free(cat->magerrlist);
    fl_free(cat->siglist);
    fl_free(cat->pmlist);
    fl_free(cat->sigpmlist);
    bl_free(cat->idlist);
	free(cat);
    return rtn;
}

