from astropy.wcs import WCS
from astropy.io import fits
from astropy.table import Table
from astropy.utils.exceptions import AstropyWarning

import numpy as np
import scipy
import matplotlib.pyplot as plt
import warnings
import os


def update_WCS(coords, refs, crval, CD):
    # Solve equations:
    X_obs = list()
    X_ref = list()
    for obs, ref in zip(coords, refs):
        x0, y0 = obs
        xr, yr = ref
        # For x:
        X_obs.append([1., 0., x0, y0, 0., 0.])
        X_ref.append(xr)
        # For y:
        X_obs.append([0., 1., 0., 0., x0, y0])
        X_ref.append(yr)

    X_obs = np.array(X_obs)
    X_ref = np.array(X_ref)

    p_opt = np.linalg.lstsq(X_obs, X_ref, rcond=None)[0]
    cx, cy, a, b, c, d = p_opt

    R = np.array([[a, b], [c, d]])
    offset = np.array([cx, cy])

    # Calculate updated solution:
    crval_new = offset + R.dot(crval)
    CD_new = R.dot(CD)
    wcs_new = {
        'CRVAL1': crval_new[0],
        'CRVAL2': crval_new[1],
        'CD1_1': CD_new[0, 0],
        'CD1_2': CD_new[0, 1],
        'CD2_1': CD_new[1, 0],
        'CD2_2': CD_new[1, 1]}
    return wcs_new


def mad(x):
    return np.median(np.abs(x-np.median(x)))


def get_WCS(hdr):
    p0 = np.array([hdr['CRVAL1'], hdr['CRVAL2']])
    x0 = np.array([hdr['CRPIX1'], hdr['CRPIX2']])
    CD = np.array([[hdr['CD1_1'], hdr['CD1_2']],
                   [hdr['CD2_1'], hdr['CD2_2']]])
    proj = np.array([np.cos(p0[1]/180*np.pi), 1.])
    return x0, p0, CD, proj


def pixel_to_radec(coords_pix, header):
    # -- Convert SEP pixel coordinates to WCS using header
    x0, p0, CD, proj = get_WCS(header)
    tmp = coords_pix - (x0-1)
    coords = p0 + CD.dot(tmp.T).T / proj
    return coords


def match_catalogs(coords, refs):
    matched = list()
    matches_in_coords = []
    used_index = list()
    for xy in coords:
        index = np.argmin(np.sum((refs - xy)**2, axis=1))
        # if index not in used_index:
        matched.append(refs[index])
        used_index.append(index)
        matches_in_coords.append(xy)
    return np.array(matches_in_coords), np.array(matched)


def get_gaia_catalog(ra, dec, radius=4., limit=2000, catalog_fname='', database='dr3'):
    """
    ra and dec: units of degrees
    radius: units of arcmin
    limit: max number of targets to retrieve
    """
    from astroquery.gaia import Gaia
    query_args = {'limit': limit, 'ra': ra, 'dec': dec, 'radius': radius/60., 'dr': database}
    query = """SELECT TOP {limit} ra, dec, phot_g_mean_mag, bp_rp FROM gaia{dr}.gaia_source
    WHERE CONTAINS(POINT('ICRS', gaia{dr}.gaia_source.ra, gaia{dr}.gaia_source.dec),
                   CIRCLE('ICRS', {ra}, {dec}, {radius}))=1;""".format(**query_args)
    job = Gaia.launch_job_async(query, dump_to_file=True,
                                output_format='csv', verbose=False,
                                output_file=catalog_fname)
    result = job.get_results()
    return result




def correct_wcs(img_fname, sep_fname, output='', fig_fname='', min_num=6, G_lim=5, q_lim=0.8, kde_factor=0.1,
                p_kde=0.5, debug=False, plot=False):
    """
    WCS calibration using Gaia

    Sources are matched using the initial WCS solution from the header. Outliers are rejected
    based on projected distance. The WCS solution is then matched to match the Gaia positions.

    Parameters
    ----------
    img_fname : string
        ALFOSC image filename (.fits)

    sep_fname : string
        Filename of the source extraction table (_phot.fits)

    output : string
        Filename of the WCS corrected image. Autogenerated by default.

    fig_fname : string
        Filename of the diagnostic figure. Autogenerated by default.

    min_num : int  [default=6]
        Minimum number of targets required.

    G_lim : float  [default=5]
        Bright limit of Gaia sources. Reject targets brighter than G < G_lim as these are
        very likely saturated in the ALFOSC image.

    q_lim : float  [default=0.8]
        Reject elliptical sources with axis ratio < `q_lim`.
        Axis ratio is defined as minor/major.

    kde_factor : float  [default=0.1]
        Smoothing factor for the kernel density estimator used for source filtering
        in the distance-orientation space.
    
    p_kde : float  [default=0.5]
        Threshold for projected vector filtering. The clustering of sources in the distance-orientation space
        is estimated using a kernel density estimator with a smoothing scale of `kde_factor`.
        A given source is rejected if its KDE probability is less than this threshold.
    
    plot : bool  [default=False]
        Plot a diagnostic plot of the identified sources after WCS alignment?

    Returns
    -------
    output_msg : string
        Log of messages from the function call.
    """
    msg = list()

    img = fits.getdata(img_fname)
    hdr = fits.getheader(img_fname)
    msg.append("          - Loaded image: %s" % img_fname)

    # Prepare output filenames:
    base, ext = os.path.splitext(os.path.basename(img_fname))
    dirname = os.path.dirname(img_fname)

    if fig_fname == '':
        fig_fname = 'wcs_solution_' + base + '.pdf'
        fig_fname = os.path.join(dirname, fig_fname)

    image_radius = np.max([hdr['NAXIS1'], hdr['NAXIS2']]) / 2 * 0.95
    image_scale = np.sqrt(hdr['CD1_1']**2 + hdr['CD1_2']**2)
    deg_to_arcmin = 60.
    radius = image_scale * image_radius * deg_to_arcmin
    gaia_cat_name = 'gaia_source_%.2f%+.2f_%.1f.csv' % (hdr['CRVAL1'], hdr['CRVAL2'], radius)
    gaia_cat_name = os.path.join(dirname, gaia_cat_name)
    gaia_dr = 'dr3'
    if os.path.exists(gaia_cat_name):
        ref_cat = Table.read(gaia_cat_name)
        msg.append("          - Loading Gaia source catalog: %s" % gaia_cat_name)
        msg.append("          - Position: (ra, dec) = (%.5f ; %+.5f)  within %.1f arcmin" % (hdr['CRVAL1'],
                                                                                             hdr['CRVAL2'],
                                                                                             radius))
    else:
        # Download Gaia positions:
        msg.append("          - Downloading Gaia source catalog... (%s)" % gaia_dr.upper())
        msg.append("          - Position: (ra, dec) = (%.5f ; %+.5f)  within %.1f arcmin" % (hdr['CRVAL1'],
                                                                                             hdr['CRVAL2'],
                                                                                             radius))
        try:
            ref_cat = get_gaia_catalog(hdr['CRVAL1'], hdr['CRVAL2'],
                                       radius=radius, catalog_fname=gaia_cat_name, database=gaia_dr)
            msg.append("          - Saving Gaia source catalog: %s" % gaia_cat_name)
        except:
            msg.append(" [ERROR]  - Could not reach Gaia server! Check your internet connection.")
            msg.append("")
            return "\n".join(msg)

    # reject brightest sources:
    bright_cut = ref_cat['phot_g_mean_mag'] > G_lim
    msg.append("          - Rejecting sources brighter than G < %.1f mag" % G_lim)

    # Get the source extraction catalog:
    sep_cat = Table.read(sep_fname)
    msg.append("          - Loaded image source catalog: %s" % sep_fname)
    ### Reject brightest and faintest 10%:
    # lower, upper = np.percentile(sep_cat['flux_auto'], [10, 90])
    # cut_flux = (sep_cat['flux_auto'] > lower) & (sep_cat['flux_auto'] < upper)
    # #sep_cat = sep_cat[cut_flux]

    # Pixel coordinates:
    sep_coords = np.array([sep_cat['x'], sep_cat['y']]).T

    # Convert to WCS:
    sep_wcs = pixel_to_radec(sep_coords, hdr)
    axis_ratio = sep_cat['b']/sep_cat['a']
    
    # Sort by source magnitude:
    index = np.argsort(sep_cat['mag_auto'])
    sep_wcs = sep_wcs[index]

    # Select only 'round' sources:
    axis_mask = axis_ratio > q_lim
    msg.append("          - Rejecting sources with axis ratio < %.2f" % q_lim)

    ref_coords = np.array([ref_cat['ra'], ref_cat['dec']]).T
    ref_coords = ref_coords[bright_cut]
    N_gaia = len(ref_coords)

    sep_wcs = sep_wcs[axis_mask]
    sep_wcs = sep_wcs[:(N_gaia+10)]

    if len(ref_coords) < len(sep_wcs):
        matched_ref, matched_sep = match_catalogs(ref_coords, sep_wcs)
    else:
        matched_sep, matched_ref = match_catalogs(sep_wcs, ref_coords)

    # Reject outliers:
    msg.append("          - Performing initial cross identification")
    msg.append("          - Found %s source%s" % (len(matched_ref),
                                                  '' if len(matched_ref)==1 else 's'))
    dist = np.sqrt(np.sum((matched_ref - matched_sep)**2, axis=1))*3600
    dxdy = np.abs(matched_sep - matched_ref)
    theta = np.arctan2(dxdy[:, 1], dxdy[:, 0])
    values = np.vstack([dist, theta])
    kernel = scipy.stats.gaussian_kde(values, bw_method=kde_factor)
    
    ymin = 0.
    ymax = np.max(theta)*1.25
    xmin = 0.
    xmax = np.max(dist)*1.25
    x = np.linspace(xmin, xmax, 200)
    y = np.linspace(ymin, ymax, 200)
    xx, yy = np.meshgrid(x, y)
    positions = np.vstack([xx.ravel(), yy.ravel()])
    pdf = np.reshape(kernel(positions).T, xx.shape)
    p0 = pdf.max()
    mask = (kernel(values) / p0) > p_kde
    msg.append("          - Rejecting outliers with p-val > %.2f" % p_kde)
    msg.append("          - Initial average WCS offset: %.3f arcsec" % (np.median(dist[mask])))
    
    if debug:
        plt.close('all')
        pdf = pdf/p0
        peak = (pdf == pdf.max()).nonzero()
        peak_dist = x[peak[1]]
        peak_theta = y[peak[0]]
        plt.contourf(xx, yy, pdf, cmap='Blues')
        plt.contour(xx, yy, pdf, [p_kde, 1.0], colors=['r', 'r'])
        plt.plot(dist, theta, 'k+')
        plt.plot(dist[mask], theta[mask], 'ro', alpha=0.5)
        plt.axvline(peak_dist, color='r', ls=':', alpha=0.5)
        plt.axhline(peak_theta, color='r', ls=':', alpha=0.5)
        plt.xlabel("Offset (arcsec)")
        plt.ylabel("Offset orientation, $\\theta$ (radians)")
        plt.show()
    
    matched_sep = matched_sep[mask]
    matched_ref = matched_ref[mask]

    if np.sum(mask) < min_num:
        msg.append(" [ERROR]  - Not enough targets found in the image!")
        msg.append(" [ERROR]  - Found %i but expected at least %i" % (np.sum(mask), min_num))
        msg.append("")
        return "\n".join(msg)

    # Fit the WCS transformation:    
    _, crval, CD, _ = get_WCS(hdr)
    msg.append("          - Fitting WCS transformation using %i sources" % len(matched_ref))
    wcs_keys = update_WCS(matched_sep, matched_ref, crval, CD)
    hdr.update(wcs_keys)
    hdr.add_comment("PyNOT: WCS calibration using Gaia %s" % gaia_dr.upper())
    hdr['RADESYSa'] = 'ICRS'
    msg.append("          - Updating WCS information")


    # -- Plot solution:
    if plot:
        plt.close('all')
        warnings.simplefilter('ignore', category=AstropyWarning)
        wcs_new = WCS(hdr)
        fig = plt.figure()
        ax = fig.add_subplot(111, projection=wcs_new)
        med_val = np.median(img)
        ax.imshow(img, vmin=med_val-1*mad(img), vmax=med_val+10*mad(img),
                origin='lower', cmap=plt.cm.gray_r)
        ax.scatter(matched_ref[:, 0], matched_ref[:, 1],
                transform=ax.get_transform('fk5'), edgecolor='r', facecolor='none', s=50)
        if debug:
            ax.scatter(matched_sep[:, 0], matched_sep[:, 1],
                    transform=ax.get_transform('fk5'), edgecolor='green', facecolor='none', s=10, alpha=0.7)
            for i, r in zip(matched_sep, matched_ref):
                ax.plot([i[0], r[0]], [i[1], r[1]], 'k', alpha=0.5, transform=ax.get_transform('fk5'))
        ax.set_xlabel("Right Ascension")
        ax.set_ylabel("Declination")
        fig.savefig(fig_fname)
        msg.append(" [OUTPUT] - Saving WCS solution figure: %s" % fig_fname)


    # -- Calculate Dispersion in WCS Solution:
    sep_wcs = pixel_to_radec(sep_coords, hdr)
    ref_coords = np.array([ref_cat['ra'], ref_cat['dec']]).T
    ref_coords = ref_coords[bright_cut]
    if len(ref_coords) < len(sep_wcs):
        matched_ref, matched_sep = match_catalogs(ref_coords, sep_wcs)
    else:
        matched_sep, matched_ref = match_catalogs(sep_wcs, ref_coords)
    dist = np.sqrt(np.sum((matched_ref - matched_sep)**2, axis=1)) * 3600
    wcs_resid = mad(dist) * 1.48
    wcs_offset = np.median(dist)
    msg.append("          - WCS precision: %.3f arcsec" % wcs_resid)
    msg.append("          - average WCS offset: %.3f arcsec" % wcs_offset)
    msg.append("[WARNING] - Large WCS offset!")


    # Update photometric table:
    sep_cat['ra'] = sep_wcs[:, 0]
    sep_cat['dec'] = sep_wcs[:, 1]
    sep_cat.write(sep_fname, format='fits', overwrite=True)
    msg.append(" [OUTPUT] - Updating source identification table: %s" % sep_fname)


    # -- Update image:
    if output == '':
        with fits.open(img_fname, mode='update') as hdu_list:
            if 'DATA' not in hdu_list:
                hdu_list[0].header = hdr
            else:
                hdu_list['DATA'].header = hdr
            hdu_list['ERR'].header.update(wcs_keys)
        output = img_fname
    else:
        with fits.open(img_fname) as hdu_list:
            if 'DATA' not in hdu_list:
                hdu_list[0].header = hdr
            else:
                hdu_list['DATA'].header = hdr
            hdu_list['ERR'].header.update(wcs_keys)
            hdu_list.writeto(output, overwrite=True)

    msg.append(" [OUTPUT] - Saving WCS calibrated image: %s" % output)
    msg.append("")
    output_msg = "\n".join(msg)
    return output_msg
