#!/usr/bin/ipython3
#-*- coding:utf-8 -*-

import sys
import matplotlib.pyplot as plt
from astropy.io import fits
from astropy import wcs
import numpy as np
import healpy as hp


def create_WCS(sky_pos, ROI=[0.5, 128], radec=False):
    """
    Create a WCS object centered at given sky position.
    """
    w = wcs.WCS(naxis=2)
    center = np.array([int(ROI[1]/2.), int(ROI[1]/2.)])
    w.wcs.crpix = center
    w.wcs.crval = sky_pos
    w.wcs.cdelt = [-ROI[0]/60., ROI[0]/60.]

    if radec:
        w.wcs.ctype = ['RA--TAN','DEC--TAN']
    else:
        w.wcs.ctype = ['GLON-TAN','GLAT-TAN']

    return w


def create_ROI(hpmap, sky_pos, ROI=[0.5, 128], radec=False):
    """
    Create a projected map given a Healpix map and a region of interest.
    """
    nside = hp.get_nside(hpmap)
    w = create_WCS(sky_pos, ROI, radec)

    data = np.zeros((ROI[1], ROI[1]))
    yy, xx = np.indices(data.shape)
    theta, phi = w.wcs_pix2world(xx, yy, 0)
    ipx = hp.ang2pix(nside, theta, phi, lonlat=True)
    data = hpmap[ipx]

    return w, data


def plot_ROI(w, data, name=None, freq=None, pos=None, saveplot=False):
    """
    Display a nice plot of the given data.
    """
    fig = plt.figure(figsize=(10,5))
    ax = plt.subplot(projection=w, label='overlays', figure=fig)

    ax.imshow(data, origin='lower')

    ax.coords.grid(True, color='white', ls='solid')
    ax.coords[0].set_axislabel('Galactic Longitude')
    ax.coords[1].set_axislabel('Galactic Latitude')

    overlay = ax.get_coords_overlay('fk5')
    overlay.grid(color='white', ls='dotted')
    overlay[0].set_axislabel('Right Ascension (J2000)')
    overlay[1].set_axislabel('Declination (J2000)')

    if name != None:
        ax.set_title("Object {0:s} at {1:d}GHz".format(name, freq), y=1.0, pad=+40)
        if saveplot:
            plt.savefig("../plots/{0:s}_G{1:.2f}±{2:.2f}_{2:d}GHz.png".format(name, *pos, freq), bbox_inches='tight')
    plt.show()

    return 0


def save_file(map_name,sky_pos,header,data):
    """
    Save given projected map to fits file with appropriate header.
    """
    hdu = fits.PrimaryHDU(data,header)
    filename = '../data/' + str(map_name) + "_G{0:.2f}±{1:.2f}".format(*sky_pos) + '.fits'

    hdu.writeto(filename, overwrite=True)

    return 0


def get_data_sample(freq, sample, ROI=[0.5, 128], savename=None):
    """
    Construct on array (nobj, npix, npix) at a given frequency for a given sample and ROI parameters. Allow local saving.
    """
    if type(freq) == int:
        freq = np.array([freq])
    nfreq = len(freq)
    nobj = len(sample)

    # Get Planck data
    dirf = "../data/Planck2015/"
    fits_basename = "HFI_SkyMap_{0:d}_2048_R2.00_full.fits"

    data_sample = np.zeros((nfreq, nobj, ROI[1], ROI[1]))
    wcs_sample = np.zeros((nfreq, nobj), dtype=object)
    for i in range(nfreq):
        hpmap = hp.read_map(dirf+fits_basename.format(freq[i]), verbose=False)
        for j in range(nobj):
            pos_obj = [sample['GLON'][j], sample['GLAT'][j]]
            wcs_obj, data_obj = create_ROI(hpmap, pos_obj, ROI, radec=False)
            data_sample[i,j] = data_obj
            wcs_sample[i,j] = wcs_obj

    if nfreq == 1:
        wcs_sample, data_sample = wcs_sample[0], data_sample[0]

    if savename != None:
        np.save("../data/wcs_{0:s}.npy".format(savename), wcs_sample)
        np.save("../data/data_{0:s}.npy".format(savename), data_sample)

    return wcs_sample, data_sample


def save_sample(sample, sample_name, freq, ROI=[0.5, 128]): #sample from PSZ2v1.fits, freq from planck full sky map
    # Get Planck data
    dirf = "../data/Planck2015/"
    fits_basename = "HFI_SkyMap_{0:d}_2048_R2.00_full.fits"
    hpmap = hp.read_map(dirf+fits_basename.format(freq), verbose=False)

    hdul = fits.HDUList([])
    primary_hdu = fits.PrimaryHDU()
    hdul.append(primary_hdu)
    hdr = hdul[0].header
    hdr['sample'] = (sample_name, 'Same of the sample')
    hdr['freq'] = (freq, 'Detection frequency in GHz')
    hdr['nobj'] = len(sample)
    hdr['ROI_size'] = (ROI[1], 'Size in pixels of the ROI (square)')
    hdr['px_size'] = (ROI[0], 'Size of a pixel in arcmin')

    for i in range(len(sample)):
        sample_pos = [sample['GLON'][i], sample['GLAT'][i]]
        sample_w, sample_data = create_ROI(hpmap, sample_pos, ROI, radec=False)

        sample_header = sample_w.to_header()
        sample_header['targname'] = sample['NAME'][i]

        hdu_sample = fits.ImageHDU(data=sample_data,header=sample_header)
        hdul.append(hdu_sample)

    filename = '../data/' + str(sample_name) + "_{0:d}GHz".format(freq) + '.fits'
    hdul.writeto(filename, overwrite=True)

    return 0


def data_sample_from_fits(sample):
    """
    Construct an array (ncluster, npix, npix) of 'sample'.
    Input 'sample' must be a valid path to a fits file.
    """
    with fits.open(sample) as hdul:
        header = hdul[0].header

        nobj = header['nobj']
        ROI_size = header['ROI_size']
        data_sample = np.zeros((nobj, ROI_size, ROI_size))
        wcs_sample = np.zeros(nobj,dtype=object)

        for i in range(nobj):
            data_sample[i] = hdul[i+1].data
            wcs_sample[i] = wcs.WCS(hdul[i+1].header)

    return wcs_sample, data_sample


def stack_sample(data_sample, wcs_sample=None):
    """
    Compute the stack of a given sample.
    Input 'sample' can either be an array (ncluster, npixel, npixel) or a path to fits file containing the sample.
    """
    if str(type(data_sample)) != "<class 'numpy.ndarray'>":
        wcs_sample, data_sample = data_sample_from_fits(sample)

    data_stack = data_sample.mean(axis=0)
    px_size = np.array([np.abs(wcs_sample[i].wcs.cdelt) for i in range(wcs_sample.shape[0])]).mean()
    center = np.array([wcs_sample[i].wcs.crpix for i in range(wcs_sample.shape[0])]).mean(axis=0)

    return data_stack, px_size, center


def plot_stacks(stack_list, stack_names, save=False):
    """
    Plot together stacks at each frequencies for different samples.
    """
    freq = np.array([100, 143, 217, 353, 545, 857])
    nfreq, nsample = len(freq), len(stack_list)

    grid = plt.GridSpec(nsample, nfreq, hspace=0, wspace=0)
    fig = plt.figure(figsize=(16,9))
    for i in range(nsample):
        for j in range(nfreq):
            ax = fig.add_subplot(grid[i,j], xticks=[], yticks=[])
            ax.imshow(stack_list[i,j], origin='lower')
            if i == nsample-1:
                ax.set_xlabel("{0:d}GHz".format(freq[j]))
            if j == 0:
                ax.set_ylabel(stack_names[i])
    if save:
        filename = "../plots/Stacks"
        for name in stack_names:
            filename += "_"+name
        plt.savefig(filename+".png", bbox_inches='tight')
    plt.show()
    return 0


def main():
    # Get Planck data
    dirf = "../data/Planck2015/"
    ext = ".fits"
    fits_basename = "HFI_SkyMap_{0:d}_2048_R2.00_{1:s}"
    mission = "full"    #Can be 'full', 'halfmission-1', 'halfmission-2'
    frequencies = np.array([100, 143, 217, 353, 545, 857])

    hfi_maps = np.array([fits_basename.format(freq,mission) for freq in frequencies])

    ## Get known clusters and candidates
    with fits.open("../data/PSZ2v1.fits") as hdu:
        PSZ_catalog = hdu[1].data

    clusters = PSZ_catalog[PSZ_catalog['redshift']>0]
    candidates = PSZ_catalog[(PSZ_catalog['redshift']<0) & (PSZ_catalog['validation']<0)]

#    wcs_clusters, data_clusters = get_data_sample(frequencies, clusters, ROI=[0.5, 128], savename="clusters")
#    wcs_candidates, data_candidates = get_data_sample(frequencies, candidates, ROI=[0.5, 128], savename="candidates")
    wcs_clusters, data_clusters = np.load("../data/wcs_clusters.npy"), np.load("../data/data_clusters.npy")
    wcs_candidates, data_candidates = np.load("../data/wcs_candidates.npy"), np.load("../data/data_candidates.npy")

    stacks_clusters = np.zeros((len(frequencies), data_clusters.shape[2], data_clusters.shape[3]))
    stacks_candidates = np.zeros((len(frequencies), data_candidates.shape[2], data_candidates.shape[3]))

    for i in range(len(frequencies)):
        stacks_clusters[i] = stack_sample(data_clusters[i], wcs_clusters[i])[0]
        stacks_candidates[i] = stack_sample(data_candidates[i], wcs_candidates[i])[0]

    plot_stacks(np.array([stacks_clusters, stacks_candidates]), ['clusters', 'candidates'], save=True)

    return 0


if __name__ == "__main__":
    sys.exit(main())
