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

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

#include "qfits_rw.h"
#include "qfits_tools.h"

#include "os-features.h"
#include "sip_qfits.h"
#include "an-bool.h"
#include "fitsioutils.h"
#include "errors.h"
#include "log.h"
#include "ioutils.h"

#include "anqfits.h"

sip_t* sip_from_string(const char* str, int slen, sip_t* dest) {
    qfits_header* hdr;
    sip_t* rtn;
    if (slen == 0) {
        slen = strlen(str);
    }
    hdr = qfits_header_read_hdr_string((const unsigned char*)str, slen);
    if (!hdr) {
        ERROR("Failed to parse a FITS header from the given string");
        return NULL;
    }
    rtn = sip_read_header(hdr, dest);
    qfits_header_destroy(hdr);
    return rtn;
}

sip_t* sip_read_tan_or_sip_header_file_ext(const char* wcsfn, int ext, sip_t* dest, anbool forcetan) {
    sip_t* rtn;
    if (forcetan) {
        sip_t sip;
        memset(&sip, 0, sizeof(sip_t));
        if (!tan_read_header_file_ext(wcsfn, ext, &(sip.wcstan))) {
            ERROR("Failed to parse TAN header from file %s, extension %i", wcsfn, ext);
            return NULL;
        }
        if (!dest)
            dest = malloc(sizeof(sip_t));
        memcpy(dest, &sip, sizeof(sip_t));
        return dest;
    } else {
        rtn = sip_read_header_file_ext(wcsfn, ext, dest);
        if (!rtn)
            ERROR("Failed to parse SIP header from file %s, extension %i", wcsfn, ext);
        return rtn;
    }
}

int sip_write_to(const sip_t* sip, FILE* fid) {
    qfits_header* hdr;
    int res;
    if ((sip->a_order == 0) && (sip->b_order == 0) &&
        (sip->ap_order == 0) && (sip->bp_order == 0))
        return tan_write_to(&(sip->wcstan), fid);
    hdr = sip_create_header(sip);
    if (!hdr) {
        ERROR("Failed to create FITS header from WCS");
        return -1;
    }
    res = qfits_header_dump(hdr, fid);
    qfits_header_destroy(hdr);
    return res;
}

int sip_write_to_file(const sip_t* sip, const char* fn) {
    FILE* fid;
    int res;
    if ((sip->a_order == 0) && (sip->b_order == 0) &&
        (sip->ap_order == 0) && (sip->bp_order == 0))
        return tan_write_to_file(&(sip->wcstan), fn);
		
    fid = fopen(fn, "wb");
    if (!fid) {
        SYSERROR("Failed to open file \"%s\" to write WCS header", fn);
        return -1;
    }
    res = sip_write_to(sip, fid);
    if (res) {
        ERROR("Failed to write FITS header to file \"%s\"", fn);
        return -1;
    }
    if (fclose(fid)) {
        SYSERROR("Failed to close file \"%s\" after writing WCS header", fn);
        return -1;
    }
    return 0;
}

int tan_write_to(const tan_t* tan, FILE* fid) {
    qfits_header* hdr;
    int res;
    hdr = tan_create_header(tan);
    if (!hdr) {
        ERROR("Failed to create FITS header from WCS");
        return -1;
    }
    res = qfits_header_dump(hdr, fid);
    qfits_header_destroy(hdr);
    return res;
}

int tan_write_to_file(const tan_t* tan, const char* fn) {
    FILE* fid;
    int res;
    fid = fopen(fn, "wb");
    if (!fid) {
        SYSERROR("Failed to open file \"%s\" to write WCS header", fn);
        return -1;
    }
    res = tan_write_to(tan, fid);
    if (res) {
        ERROR("Failed to write FITS header to file \"%s\"", fn);
        return -1;
    }
    if (fclose(fid)) {
        SYSERROR("Failed to close file \"%s\" after writing WCS header", fn);
        return -1;
    }
    return 0;
}

static void wcs_hdr_common(qfits_header* hdr, const tan_t* tan) {
    qfits_header_add(hdr, "WCSAXES", "2", NULL, NULL);
    qfits_header_add(hdr, "EQUINOX", "2000.0", "Equatorial coordinates definition (yr)", NULL);
    qfits_header_add(hdr, "LONPOLE", "180.0", NULL, NULL);
    qfits_header_add(hdr, "LATPOLE", "0.0", NULL, NULL);

    fits_header_add_double(hdr, "CRVAL1", tan->crval[0], "RA  of reference point");
    fits_header_add_double(hdr, "CRVAL2", tan->crval[1], "DEC of reference point");
    fits_header_add_double(hdr, "CRPIX1", tan->crpix[0], "X reference pixel");
    fits_header_add_double(hdr, "CRPIX2", tan->crpix[1], "Y reference pixel");
    qfits_header_add(hdr, "CUNIT1", "deg", "X pixel scale units", NULL);
    qfits_header_add(hdr, "CUNIT2", "deg", "Y pixel scale units", NULL);

    fits_header_add_double(hdr, "CD1_1", tan->cd[0][0], "Transformation matrix");
    fits_header_add_double(hdr, "CD1_2", tan->cd[0][1], "");
    fits_header_add_double(hdr, "CD2_1", tan->cd[1][0], "");
    fits_header_add_double(hdr, "CD2_2", tan->cd[1][1], "");

    if (tan->imagew > 0.0)
        fits_header_add_double(hdr, "IMAGEW", tan->imagew, "Image width,  in pixels.");
    if (tan->imageh > 0.0)
        fits_header_add_double(hdr, "IMAGEH", tan->imageh, "Image height, in pixels.");
}

int sip_get_image_size(const qfits_header* hdr, int* pW, int* pH) {
    int W, H;
    W = qfits_header_getint(hdr, "IMAGEW", 0);
    debug("sip_get_image_size: IMAGEW = %i\n", W);
    H = qfits_header_getint(hdr, "IMAGEH", 0);
    debug("sip_get_image_size: IMAGEH = %i\n", H);
    if (W == 0 || H == 0) {
        // no IMAGE[WH].  Check for fpack-compressed image.
        int eq;
        char* str = fits_get_dupstring(hdr, "XTENSION");
        //printf("XTENSION: '%s'\n", str);
        // qfits_header_getstr turns the string double-quotes to single-quotes
        eq = streq(str, "BINTABLE");
        free(str);
        if (eq) {
            // ZNAXIS1 =                 2046 / length of data axis 1
            // ZNAXIS2 =                 4094 / length of data axis 2
            if (!W) {
                W = qfits_header_getint(hdr, "ZNAXIS1", 0);
                debug("sip_get_image_size: ZNAXIS1 = %i\n", W);
            }
            if (!H) {
                H = qfits_header_getint(hdr, "ZNAXIS2", 0);
                debug("sip_get_image_size: ZNAXIS2 = %i\n", H);
            }
        }
        if (!W) {
            W = qfits_header_getint(hdr, "NAXIS1", 0);
            debug("sip_get_image_size: NAXIS1 = %i\n", W);
        }
        if (!H) {
            H = qfits_header_getint(hdr, "NAXIS2", 0);
            debug("sip_get_image_size: NAXIS2 = %i\n", H);
        }
    }
    if (pW) *pW = W;
    if (pH) *pH = H;
    return 0;
}

static void add_polynomial(qfits_header* hdr, const char* format,
                           int order, const double* data, int datastride) {
    int i, j;
    char key[64];
    for (i=0; i<=order; i++)
        for (j=0; (i+j)<=order; j++) {
            //if (i+j < 1)
            //	continue;
            //if (drop_linear && (i+j < 2))
            //	continue;
            sprintf(key, format, i, j);
            fits_header_add_double(hdr, key, data[i*datastride + j], "");
        }
}

void sip_add_to_header(qfits_header* hdr, const sip_t* sip) {
    wcs_hdr_common(hdr, &(sip->wcstan));
    if (sip->wcstan.sin) {
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE2", "DEC--SIN-SIP", "SIN projection + SIP distortions", NULL);
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE1", "RA---SIN-SIP", "SIN projection + SIP distortions", NULL);
    } else {
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE2", "DEC--TAN-SIP", "TAN (gnomic) projection + SIP distortions", NULL);
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE1", "RA---TAN-SIP", "TAN (gnomic) projection + SIP distortions", NULL);
    }

    fits_header_add_int(hdr, "A_ORDER", sip->a_order, "Polynomial order, axis 1");
    add_polynomial(hdr, "A_%i_%i", sip->a_order, (double*)sip->a, SIP_MAXORDER);

    fits_header_add_int(hdr, "B_ORDER", sip->b_order, "Polynomial order, axis 2");
    add_polynomial(hdr, "B_%i_%i", sip->b_order, (double*)sip->b, SIP_MAXORDER);

    fits_header_add_int(hdr, "AP_ORDER", sip->ap_order, "Inv polynomial order, axis 1");
    add_polynomial(hdr, "AP_%i_%i", sip->ap_order, (double*)sip->ap, SIP_MAXORDER);

    fits_header_add_int(hdr, "BP_ORDER", sip->bp_order, "Inv polynomial order, axis 2");
    add_polynomial(hdr, "BP_%i_%i", sip->bp_order, (double*)sip->bp, SIP_MAXORDER);
}

qfits_header* sip_create_header(const sip_t* sip) {
    qfits_header* hdr = qfits_table_prim_header_default();
    sip_add_to_header(hdr, sip);
    return hdr;
}

void tan_add_to_header(qfits_header* hdr, const tan_t* tan) {
    wcs_hdr_common(hdr, tan);
    if (tan->sin) {
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE2", "DEC--SIN", "SIN projection", NULL);
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE1", "RA---SIN", "SIN projection", NULL);
    } else {
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE2", "DEC--TAN", "TAN (gnomic) projection", NULL);
        qfits_header_add_after(hdr, "WCSAXES", "CTYPE1", "RA---TAN", "TAN (gnomic) projection", NULL);
    }
}

qfits_header* tan_create_header(const tan_t* tan) {
    qfits_header* hdr = qfits_table_prim_header_default();
    tan_add_to_header(hdr, tan);
    return hdr;
}

static void* read_header_file(const char* fn, int ext, anbool only, void* dest,
                              void* (*readfunc)(const qfits_header*, void*)) {
    qfits_header* hdr;
    void* result;
    if (only) {
        hdr = anqfits_get_header_only(fn, ext);
    } else {
        hdr = anqfits_get_header2(fn, ext);
    }
    if (!hdr) {
        ERROR("Failed to read FITS header from file \"%s\" extension %i", fn, ext);
        return NULL;
    }
    result = readfunc(hdr, dest);
    if (!result) {
        ERROR("Failed to parse WCS header from file \"%s\" extension %i", fn, ext);
    }
    qfits_header_destroy(hdr);
    return result;
}

// silly little dispatch function to avoid casting - I like a modicum of type safety
static void* call_sip_read_header(const qfits_header* hdr, void* dest) {
    return sip_read_header(hdr, dest);
}
sip_t* sip_read_header_file(const char* fn, sip_t* dest) {
    return read_header_file(fn, 0, FALSE, dest, call_sip_read_header);
}
sip_t* sip_read_header_file_ext(const char* fn, int ext, sip_t* dest) {
    return read_header_file(fn, ext, TRUE, dest, call_sip_read_header);
}
sip_t* sip_read_header_file_ext_only(const char* fn, int ext, sip_t* dest) {
    return read_header_file(fn, ext, TRUE, dest, call_sip_read_header);
}

static void* call_tan_read_header(const qfits_header* hdr, void* dest) {
    return tan_read_header(hdr, dest);
}
tan_t* tan_read_header_file(const char* fn, tan_t* dest) {
    return read_header_file(fn, 0, FALSE, dest, call_tan_read_header);
}
tan_t* tan_read_header_file_ext(const char* fn, int ext, tan_t* dest) {
    return read_header_file(fn, ext, FALSE, dest, call_tan_read_header);
}
tan_t* tan_read_header_file_ext_only(const char* fn, int ext, tan_t* dest) {
    return read_header_file(fn, ext, TRUE, dest, call_tan_read_header);
}


static anbool read_polynomial(const qfits_header* hdr, const char* format,
                              int order, double* data, int datastride,
                              anbool skip_linear, anbool skip_zero) {
    int i, j;
    char key[64];
    double nil = -HUGE_VAL;
    double val;
    for (i=0; i<=order; i++)
        for (j=0; (i+j)<=order; j++) {
            if (skip_zero && (i+j < 1))
                continue;
            if (skip_linear && (i+j < 2))
                continue;
            sprintf(key, format, i, j);
            val = qfits_header_getdouble(hdr, key, nil);
            if (val == nil) {
                // don't warn if linear terms are "missing"
                if (i+j >= 2) {
                    ERROR("SIP: warning: key \"%s\" not found; setting to zero.", key);
                }
                val=0.0;
            }
            data[i*datastride + j] = val;
        }
    return TRUE;
}

sip_t* sip_read_header(const qfits_header* hdr, sip_t* dest) {
    sip_t sip;
    char* str;
    const char* key;
    const char* expect;
    const char* expect2;
    anbool is_sin;
    anbool is_tan;
    anbool skip_linear;
    anbool skip_zero;
    char pretty[FITS_LINESZ];

    memset(&sip, 0, sizeof(sip_t));

    key = "CTYPE1";
    expect  = "RA---TAN-SIP";
    expect2 = "RA---SIN-SIP";
    str = qfits_header_getstr(hdr, key);
    str = qfits_pretty_string_r(str, pretty);
    if (!str) {
        ERROR("SIP header: no key \"%s\"", key);
        return NULL;
    }
    is_tan = (strncmp(str, expect, strlen(expect)) == 0);
    is_sin = (strncmp(str, expect2, strlen(expect2)) == 0);
    if (!(is_tan || is_sin)) {
        if (!tan_read_header(hdr, &(sip.wcstan))) {
            ERROR("SIP: failed to read TAN header");
            return NULL;
        }
        goto gohome;
    }

    key = "CTYPE2";
    if (is_sin) {
        expect = "DEC--SIN-SIP";
    } else {
        expect = "DEC--TAN-SIP";
    }
    str = qfits_header_getstr(hdr, key);
    str = qfits_pretty_string_r(str, pretty);
    if (!str || strncmp(str, expect, strlen(expect))) {
        ERROR("SIP header: incorrect key \"%s\": expected \"%s\", got \"%s\"", key, expect, str);
        return NULL;
    }

    if (!tan_read_header(hdr, &sip.wcstan)) {
        ERROR("SIP: failed to read TAN header");
        return NULL;
    }

    sip.a_order  = qfits_header_getint(hdr, "A_ORDER", -1);
    sip.b_order  = qfits_header_getint(hdr, "B_ORDER", -1);
    sip.ap_order = qfits_header_getint(hdr, "AP_ORDER", 0);
    sip.bp_order = qfits_header_getint(hdr, "BP_ORDER", 0);

    if ((sip.a_order == -1) || 
        (sip.b_order == -1)) {
        ERROR("SIP: failed to read polynomial orders (A_ORDER=%i, B_ORDER=%i, -1 means absent)\n",
              sip.a_order, sip.b_order);
        return NULL;
    }
    if ((sip.ap_order == 0) || 
        (sip.bp_order == 0)) {
        logverb("Warning: SIP: failed to read polynomial orders (A_ORDER=%i, B_ORDER=%i (-1 means absent), AP_ORDER=%i, BP_ORDER=%i, (0 means absent)\n",
                sip.a_order, sip.b_order, sip.ap_order, sip.bp_order);
    }

    if ((sip.a_order > SIP_MAXORDER) || 
        (sip.b_order > SIP_MAXORDER) || 
        (sip.ap_order > SIP_MAXORDER) || 
        (sip.bp_order > SIP_MAXORDER)) {
        ERROR("SIP: polynomial orders (A=%i, B=%i, AP=%i, BP=%i) exceeds maximum of %i",
              sip.a_order, sip.b_order, sip.ap_order, sip.bp_order, SIP_MAXORDER);
        return NULL;
    }

    skip_linear = FALSE;
    skip_zero = FALSE;

    if (!read_polynomial(hdr, "A_%i_%i",  sip.a_order,  (double*)sip.a,  SIP_MAXORDER, skip_linear, skip_zero) ||
        !read_polynomial(hdr, "B_%i_%i",  sip.b_order,  (double*)sip.b,  SIP_MAXORDER, skip_linear, skip_zero) ||
        (sip.ap_order > 0 && !read_polynomial(hdr, "AP_%i_%i", sip.ap_order, (double*)sip.ap, SIP_MAXORDER, FALSE, FALSE)) ||
        (sip.bp_order > 0 && !read_polynomial(hdr, "BP_%i_%i", sip.bp_order, (double*)sip.bp, SIP_MAXORDER, FALSE, FALSE))) {
        ERROR("SIP: failed to read polynomial terms");
        return NULL;
    }

 gohome:
    if (!dest)
        dest = malloc(sizeof(sip_t));

    memcpy(dest, &sip, sizeof(sip_t));
    return dest;
}

static int check_tan_ctypes(char* ct1, char* ct2, anbool* is_sin) {
    const char* ra  = "RA---TAN";
    const char* dec = "DEC--TAN";
    const char* ra2  = "RA---SIN";
    const char* dec2 = "DEC--SIN";
    int NC = 8;
    *is_sin = FALSE;
    if (!ct1 || !ct2)
        return -1;
    if (strlen(ct1) < NC || strlen(ct2) < NC)
        return -1;
    if ((strncmp(ct1, ra, NC) == 0) && (strncmp(ct2, dec, NC) == 0))
        return 0;
    if ((strncmp(ct1, dec, NC) == 0) && (strncmp(ct2, ra, NC) == 0))
        return 1;

    if ((strncmp(ct1, ra2, NC) == 0) && (strncmp(ct2, dec2, NC) == 0)) {
        *is_sin = TRUE;
        return 0;
    }
    if ((strncmp(ct1, dec2, NC) == 0) && (strncmp(ct2, ra2, NC) == 0)) {
        *is_sin = TRUE;
        return 1;
    }
    return -1;
}

tan_t* tan_read_header(const qfits_header* hdr, tan_t* dest) {
    tan_t tan;
    double nil = -1e300;
    char* ct1;
    char* ct2;
    int swap;
    int W, H;
    anbool is_sin;

    memset(&tan, 0, sizeof(tan_t));

    ct1 = fits_get_dupstring(hdr, "CTYPE1");
    ct2 = fits_get_dupstring(hdr, "CTYPE2");
    swap = check_tan_ctypes(ct1, ct2, &is_sin);
    if (swap == -1) {
        ERROR("TAN header: expected CTYPE1 = RA---TAN, CTYPE2 = DEC--TAN "
              "(or vice versa), or RA---SIN, DEC--SIN or vice versa; "
              "got CTYPE1 = \"%s\", CYTPE2 = \"%s\"\n",
              ct1, ct2);
    }
    free(ct1);
    free(ct2);
    if (swap == -1)
        return NULL;

    sip_get_image_size(hdr, &W, &H);
    tan.imagew = W;
    tan.imageh = H;

    {
        const char* keys[] = { "CRVAL1", "CRVAL2", "CRPIX1", "CRPIX2",
                               "CD1_1", "CD1_2", "CD2_1", "CD2_2" };
        double* vals[] = { &(tan.crval[0]), &(tan.crval[1]),
                           &(tan.crpix[0]), &(tan.crpix[1]),
                           &(tan.cd[0][0]), &(tan.cd[0][1]),
                           &(tan.cd[1][0]), &(tan.cd[1][1]) };
        int i;
        for (i=0; i<4; i++) {
            *(vals[i]) = qfits_header_getdouble(hdr, keys[i], nil);
            if (*(vals[i]) == nil) {
                ERROR("TAN header: missing or invalid value for \"%s\"", keys[i]);
                return NULL;
            }
        }
        // Try CD
        int gotcd = 1;
        char* complaint = NULL;
        for (i=4; i<8; i++) {
            *(vals[i]) = qfits_header_getdouble(hdr, keys[i], nil);
            if (*(vals[i]) == nil) {
                asprintf_safe(&complaint, "TAN header: missing or invalid value for key \"%s\"", keys[i]);
                gotcd = 0;
                break;
            }
        }
        if (!gotcd) {
            double cdelt1,cdelt2;
            // Try CDELT
            char* key = "CDELT1";
            cdelt1 = qfits_header_getdouble(hdr, key, nil);
            if (cdelt1 == nil) {
                ERROR("%s; also tried but didn't find \"%s\"", complaint, key);
                free(complaint);
                return NULL;
            }
            key = "CDELT2";
            cdelt2 = qfits_header_getdouble(hdr, key, nil);
            if (cdelt2 == nil) {
                ERROR("%s; also tried but didn't find \"%s\"", complaint, key);
                free(complaint);
                return NULL;
            }
            tan.cd[0][0] = cdelt1;
            tan.cd[0][1] = 0.0;
            tan.cd[1][0] = 0.0;
            tan.cd[1][1] = cdelt2;
        }
    }

    if (swap == 1) {
        double tmp;
        tmp = tan.crval[0];
        tan.crval[0] = tan.crval[1];
        tan.crval[1] = tmp;
        // swap CD1_1 <-> CD2_1
        tmp = tan.cd[0][0];
        tan.cd[0][0] = tan.cd[1][0];
        tan.cd[1][0] = tmp;
        // swap CD1_2 <-> CD2_2
        tmp = tan.cd[0][1];
        tan.cd[0][1] = tan.cd[1][1];
        tan.cd[1][1] = tmp;
    }

    tan.sin = is_sin;

    if (!dest)
        dest = malloc(sizeof(tan_t));
    memcpy(dest, &tan, sizeof(tan_t));
    return dest;
}
