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

import sys, time
import concurrent.futures
import numpy as np
import stack_maps as stack
from astropy.io import fits
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

def aperture(stack, px_size, center, angle):
    """
    Create a boolean mask for the region to work on the stack
    """
    yy, xx = np.indices(stack.shape)
    radius = (angle/60.)/px_size
    mask = np.sqrt(np.abs(xx-center[0])**2 + np.abs(yy-center[1])**2) <= radius

    return mask


def stack_flux(freq, stacked_data, px_size, center, angle):
    """
    Compute integrated flux in Mjy/sr of stacked image given an aperture angle
    """
    conversion = {100: 244.1, 143: 371.74, 217: 483.690, 353: 287.450, 545: 1., 857: 1.}    # Conversion coefficient in Mjy/sr/K
    mask = aperture(stacked_data, px_size, center, angle)
    region = stacked_data[mask]
    background = np.mean(stacked_data[~mask])
    flux = np.sum(region - background)*conversion[freq]

    return flux


def bootstrap_flux(data_sample, wcs_sample, freq, angle, n_it, savename=None):
    """
    Implement the 'bootstrap' technic to compute flux and error given a sample.
    """
    n_obj = data_sample.shape[0]
    flux = np.zeros(n_it)
    for i in range(n_it):
        ind_it = np.random.randint(0, n_obj, n_obj)
        data_sample_it = data_sample[ind_it]
        wcs_sample_it = wcs_sample[ind_it]

        stack_it, px_size, center = stack.stack_sample(data_sample_it, wcs_sample_it)
        flux[i] = stack_flux(freq, stack_it, px_size, center, angle)

    if savename != None:
        np.save("../data/flux_{0:s}_{1:d}.npy".format(savename,freq),np.array([flux.mean(), flux.std()]))

    return flux.mean(), flux.std()


def parallel_bootstrap(data_sample, wcs_sample, freq, angle, n_it, savename=None):
    """
    Implement the 'bootstrap' technic with multithreading for multiple frequencies
    """
    nfreq = len(freq)
    exe = concurrent.futures.ThreadPoolExecutor()
    future_flux = []
    flux_sample = np.zeros((nfreq, 2))
    for i in range(nfreq):
        future_flux.append(exe.submit(bootstrap_flux, data_sample[i], wcs_sample[i], freq[i], angle, n_it))
    for i in range(nfreq):
        flux_sample[i] = future_flux[i].result()

    if savename != None:
        np.save("../data/flux_{0:s}.npy".format(savename),flux_sample)

    return flux_sample


def s_SZ(nu, y_SZ):
    """
    SED law for thermal SZ effect, in Mjy/sr
    """
    h = 6.62607015e-34  #SI units
    k_b = 1.380649e-23
    T_cmb = 2.72548
    x = h/(k_b*T_cmb)*nu   #reduced frequency

    g = x*(np.exp(x)+1.)/(np.exp(x)-1.)-4.
    return y_SZ*g*x**4*np.exp(x)/(np.exp(x)-1.)**2


def s_d(nu, A_d, beta_d, T_d):
    """
    SED law for thermal dust, in Mjy/sr (appear in IR)
    """
    h = 6.62607015e-34  #SI units
    k_b = 1.380649e-23
    nu_0 = 545e9
    gamma = h/(k_b*T_d)
    return A_d*(nu/nu_0)**(beta_d+1.)*(np.exp(gamma*nu_0)-1.)/(np.exp(gamma*nu)-1.)


def flux_fit(nu, y_SZ, A_d, beta_d, T_d):
    """
    SED law containing thermal SZ effect and IR contamination by thermal dust
    """
    return s_SZ(nu, y_SZ)+s_d(nu, A_d, beta_d, T_d)


def plot_SED(freq, flux, samplename, save=False):
    """
    Nicely plot SED and contributions from SZ and dust for a given sample
    """
    freq_lin = np.linspace(freq.min(), freq.max(), 300)
    popt, _ = curve_fit(flux_fit, xdata=freq*1e9, ydata=flux[:,0], p0=[1., 10., 2.3, 16.])
    chisq = np.sum(((flux[:,0]-flux_fit(freq,*popt))/flux[:,1])**2)

    plt.figure(figsize=(10,5))
    plt.errorbar(freq,flux[:,0], yerr=flux[:,1], fmt='ok', label=samplename)
    plt.plot(freq_lin,s_SZ(freq_lin*1e9,popt[0]), 'b--', label=r"$S_{SZ}$ model fit with $y_{SZ}$ = "+"{0:.5f}".format(popt[0]))
    plt.plot(freq_lin,s_d(freq_lin*1e9,*popt[1:]), 'r--', label=r"$S_{d}$ model fit with $A_{d}$ = "+"{0:.3e}".format(popt[1])+r"; $\beta_{d}$ = "+"{0:.5f}".format(popt[2])+r"; $T_{d}$ = "+"{0:.5f}\n".format(popt[3]))
    plt.plot(freq_lin,flux_fit(freq_lin*1e9, *popt), 'g--', label=r"$S_{SZ}+S_{d}$ model fit with $\chi^2$/ndf = "+"{0:.5f}".format(chisq/(len(freq)-len(popt))))
    plt.plot(freq,[0.]*len(freq), 'k--')

    plt.title("SED of the SZ effect for '{0:s}' sample.".format(samplename))
    plt.xlabel("Frequency [GHz]")
    plt.ylabel("Flux density [Mjy/sr]")
    plt.legend()
    if save:
        plt.savefig("../plots/SED_{0:s}.png".format(samplename), bbox_inches='tight')
    plt.show()

    return 0


def compute_SED(sample, savename, ROI=[0.5, 128], angle=None, n_it=500, plot=False):
    """
    Automatically compute and save SED
    """
    freq = np.array([100, 143, 217, 353, 545, 857])

    if angle == None:
        rimo = fits.getdata("/home/abeelen/Planck/products_2015/HFI_RIMO_R2.00.fits", 2)
        true_FWHM = rimo[:]['FWHM']
        angle = 1.*true_FWHM[freq == 143]

    print("Getting sample data for {0:s}...".format(savename))
    t0 = time.time()
    wcs_sample, data_sample = stack.get_data_sample(freq, sample, ROI=ROI, savename=savename)
    t1 = time.time()
    print("Elapsed time : {0:f}sec".format(t1-t0))

    stack_freq = np.zeros((len(freq), data_sample.shape[2], data_sample.shape[3]))
    for i in range(len(freq)):
        stack_freq[i] = stack.stack_sample(data_sample[i], wcs_sample[i])[0]
    np.save("../data/stack_{0:s}.npy".format(savename),stack_freq)

    print("Computing flux values...")
    t2 = time.time()
    flux_sample = parallel_bootstrap(data_sample, wcs_sample, freq, angle, n_it, savename=savename)
    t3 = time.time()
    print("Elapsed time : {0:f}sec".format(t3-t2))

    if plot:
        plot_SED(freq, flux_sample, savename, save=True)

    return 0


def main():
    # General information on the Planck data channels
    rimo = fits.getdata("/home/abeelen/Planck/products_2015/HFI_RIMO_R2.00.fits", 2)
    true_FWHM = rimo[:]['FWHM']

    frequencies = np.array([100, 143, 217, 353, 545, 857])

    # Get objects sample from PSZ catalog
    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)]
#    redshift_sup0_4 = PSZ_catalog[PSZ_catalog['redshift'] > 0.4]
#    redshift_inf0_4 = PSZ_catalog[np.abs(PSZ_catalog['redshift']) < 0.4]
#    QN_inf0_4 = PSZ_catalog[PSZ_catalog['q_neural'] < 0.4]
    QN_sup0_4 = PSZ_catalog[PSZ_catalog['q_neural'] > 0.4]

    # Compute SED and plot
    angle = 1.*true_FWHM[frequencies == 143]

    #compute_SED(clusters, "clusters", ROI=[0.5, 128], angle=angle, n_it=500, plot=True)
    #compute_SED(candidates, "candidates", ROI=[0.5, 128], angle=angle, n_it=500, plot=True)
    #compute_SED(QN_sup0_4, "QN_sup0_4", ROI=[0.5, 128], angle=angle, n_it=500, plot=True)
    flux_clusters = np.load("../data/flux_clusters.npy")
    flux_candidates = np.load("../data/flux_candidates.npy")
    flux_QN_sup0_4 = np.load("../data/flux_QN_sup0_4.npy")

    plot_SED(frequencies, flux_clusters, "Known clusters", save=True)
    plot_SED(frequencies, flux_candidates, "Candidate clusters", save=True)
    plot_SED(frequencies, flux_QN_sup0_4, "High-reliability clusters", save=True)

    stack_clusters = np.load("../data/stack_clusters.npy")
    stack_candidates = np.load("../data/stack_candidates.npy")
    stack_QN_sup0_4 = np.load("../data/stack_QN_sup0_4.npy")

    stack.plot_stacks(np.array([stack_clusters, stack_candidates, stack_QN_sup0_4]), ["Known clusters", "Candidate clusters", "High-reliability clusters"], save=True)

    return 0


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