from astropy.io import ascii, fits
import pylab as plt
%matplotlib inline
from astropy import wcs


import numpy as np
import xidplus
from xidplus import moc_routines
import pickle
import seaborn as sns
import aplpy

"""
Explanation: XID+ Example Run Script
(This is based on a Jupyter notebook, available in the XID+ package and can be interactively run and edited)
XID+ is a probababilistic deblender for confusion dominated maps. It is designed to:

Use a MCMC based approach to get FULL posterior probability distribution on flux
Provide a natural framework to introduce additional prior information
Allows more representative estimation of source flux density uncertainties
Provides a platform for doing science with the maps (e.g XID+ Hierarchical stacking, Luminosity function from the map etc)

Cross-identification tends to be done with catalogues, then science with the matched catalogues.
XID+ takes a different philosophy. Catalogues are a form of data compression. OK in some cases, not so much in others, i.e. confused images: catalogue compression loses correlation information. Ideally, science should be done without compression.
XID+ provides a framework to cross identify galaxies we know about in different maps, with the idea that it can be extended to do science with the maps!!
Philosophy: 

build a probabilistic generative model for the SPIRE maps
Infer model on SPIRE maps

Bayes Theorem
$p(\mathbf{f}|\mathbf{d}) \propto p(\mathbf{d}|\mathbf{f}) \times p(\mathbf{f})$
In order to carry out Bayesian inference, we need a model to carry out inference on.
For the SPIRE maps, our model is quite simple, with likelihood defined as:
    $L = p(\mathbf{d}|\mathbf{f}) \propto |\mathbf{N_d}|^{-1/2} \exp\big{ -\frac{1}{2}(\mathbf{d}-\mathbf{Af})^T\mathbf{N_d}^{-1}(\mathbf{d}-\mathbf{Af})\big}$
where:
    $\mathbf{N_{d,ii}} =\sigma_{inst.,ii}^2+\sigma_{conf.}^2$
Simplest model for XID+ assumes following:

All sources are known and have positive flux (fi)
A global background (B) contributes to all pixels 
PRF is fixed and known
Confusion noise is constant and not correlated across pixels


Because we are getting the joint probability distribution, our model is generative i.e. given parameters, we generate data and vica-versa
Compared to discriminative model (i.e. neural network), which only obtains conditional probability distribution i.e. Neural network, give inputs, get output. Can't go other way'
Generative model is full probabilistic model. Allows more complex relationships between observed and target variables
XID+ SPIRE
XID+ applied to GALFORM simulation of COSMOS field

SAM simulation (with dust) ran through SMAP pipeline_ similar depth and size as COSMOS
Use galaxies with an observed 100 micron flux of gt. $50\mathbf{\mu Jy}$. Gives 64823 sources
Uninformative prior: uniform $0 - 10{^3} \mathbf{mJy}$

Import required modules
End of explanation
"""


xidplus.__path__[0]

#Folder containing maps
imfolder=xidplus.__path__[0]+'/../test_files/'

pswfits=imfolder+'cosmos_itermap_lacey_07012015_simulated_observation_w_noise_PSW_hipe.fits.gz'#SPIRE 250 map
pmwfits=imfolder+'cosmos_itermap_lacey_07012015_simulated_observation_w_noise_PMW_hipe.fits.gz'#SPIRE 350 map
plwfits=imfolder+'cosmos_itermap_lacey_07012015_simulated_observation_w_noise_PLW_hipe.fits.gz'#SPIRE 500 map


#Folder containing prior input catalogue
catfolder=xidplus.__path__[0]+'/../test_files/'
#prior catalogue
prior_cat='lacey_07012015_MillGas.ALLVOLS_cat_PSW_COSMOS_test.fits'


#output folder
output_folder='./'

"""
Explanation: Set image and catalogue filenames
End of explanation
"""


#-----250-------------
hdulist = fits.open(pswfits)
im250phdu=hdulist[0].header
im250hdu=hdulist[1].header

im250=hdulist[1].data*1.0E3 #convert to mJy
nim250=hdulist[2].data*1.0E3 #convert to mJy
w_250 = wcs.WCS(hdulist[1].header)
pixsize250=3600.0*w_250.wcs.cd[1,1] #pixel size (in arcseconds)
hdulist.close()
#-----350-------------
hdulist = fits.open(pmwfits)
im350phdu=hdulist[0].header
im350hdu=hdulist[1].header

im350=hdulist[1].data*1.0E3 #convert to mJy
nim350=hdulist[2].data*1.0E3 #convert to mJy
w_350 = wcs.WCS(hdulist[1].header)
pixsize350=3600.0*w_350.wcs.cd[1,1] #pixel size (in arcseconds)
hdulist.close()
#-----500-------------
hdulist = fits.open(plwfits)
im500phdu=hdulist[0].header
im500hdu=hdulist[1].header 
im500=hdulist[1].data*1.0E3 #convert to mJy
nim500=hdulist[2].data*1.0E3 #convert to mJy
w_500 = wcs.WCS(hdulist[1].header)
pixsize500=3600.0*w_500.wcs.cd[1,1] #pixel size (in arcseconds)
hdulist.close()

"""
Explanation: Load in images, noise maps, header info and WCS information
End of explanation
"""


hdulist = fits.open(catfolder+prior_cat)
fcat=hdulist[1].data
hdulist.close()
inra=fcat['RA']
indec=fcat['DEC']
# select only sources with 100micron flux greater than 50 microJy
sgood=fcat['S100']>0.050
inra=inra[sgood]
indec=indec[sgood]

"""
Explanation: Load in catalogue you want to fit (and make any cuts)
End of explanation
"""


from astropy.coordinates import SkyCoord
from astropy import units as u
c = SkyCoord(ra=[150.74]*u.degree, dec=[2.03]*u.degree)  
import pymoc
moc=pymoc.util.catalog.catalog_to_moc(c,100,15)

"""
Explanation: XID+ uses Multi Order Coverage (MOC) maps for cutting down maps and catalogues so they cover the same area. It can also take in MOCs as selection functions to carry out additional cuts. Lets use the python module pymoc to create a MOC, centered on a specific position we are interested in. We will use a HEALPix order of 15 (the resolution: higher order means higher resolution), have a radius of 100 arcseconds centered around an R.A. of 150.74 degrees and Declination of 2.03 degrees.
End of explanation
"""


#---prior250--------
prior250=xidplus.prior(im250,nim250,im250phdu,im250hdu, moc=moc)#Initialise with map, uncertianty map, wcs info and primary header
prior250.prior_cat(inra,indec,prior_cat)#Set input catalogue
prior250.prior_bkg(-5.0,5)#Set prior on background (assumes Gaussian pdf with mu and sigma)
#---prior350--------
prior350=xidplus.prior(im350,nim350,im350phdu,im350hdu, moc=moc)
prior350.prior_cat(inra,indec,prior_cat)
prior350.prior_bkg(-5.0,5)

#---prior500--------
prior500=xidplus.prior(im500,nim500,im500phdu,im500hdu, moc=moc)
prior500.prior_cat(inra,indec,prior_cat)
prior500.prior_bkg(-5.0,5)

"""
Explanation: XID+ is built around two python classes. A prior and posterior class. There should be a prior class for each map being fitted. It is initiated with a map, noise map, primary header and map header and can be set with a MOC. It also requires an input prior catalogue and point spread function.
End of explanation
"""


#pixsize array (size of pixels in arcseconds)
pixsize=np.array([pixsize250,pixsize350,pixsize500])
#point response function for the three bands
prfsize=np.array([18.15,25.15,36.3])
#use Gaussian2DKernel to create prf (requires stddev rather than fwhm hence pfwhm/2.355)
from astropy.convolution import Gaussian2DKernel

##---------fit using Gaussian beam-----------------------
prf250=Gaussian2DKernel(prfsize[0]/2.355,x_size=101,y_size=101)
prf250.normalize(mode='peak')
prf350=Gaussian2DKernel(prfsize[1]/2.355,x_size=101,y_size=101)
prf350.normalize(mode='peak')
prf500=Gaussian2DKernel(prfsize[2]/2.355,x_size=101,y_size=101)
prf500.normalize(mode='peak')

pind250=np.arange(0,101,1)*1.0/pixsize[0] #get 250 scale in terms of pixel scale of map
pind350=np.arange(0,101,1)*1.0/pixsize[1] #get 350 scale in terms of pixel scale of map
pind500=np.arange(0,101,1)*1.0/pixsize[2] #get 500 scale in terms of pixel scale of map

prior250.set_prf(prf250.array,pind250,pind250)#requires PRF as 2d grid, and x and y bins for grid (in pixel scale)
prior350.set_prf(prf350.array,pind350,pind350)
prior500.set_prf(prf500.array,pind500,pind500)

print('fitting '+ str(prior250.nsrc)+' sources \n')
print('using ' +  str(prior250.snpix)+', '+ str(prior250.snpix)+' and '+ str(prior500.snpix)+' pixels')


"""
Explanation: Set PRF. For SPIRE, the PRF can be assumed to be Gaussian with a FWHM of 18.15, 25.15, 36.3 '' for 250, 350 and 500 $\mathrm{\mu m}$ respectively. Lets use the astropy module to construct a Gaussian PRF and assign it to the three XID+ prior classes.
End of explanation
"""


prior250.get_pointing_matrix()
prior350.get_pointing_matrix()
prior500.get_pointing_matrix()


"""
Explanation: Before fitting, the prior classes need to take the PRF and calculate how much each source contributes to each pixel. This process provides what we call a pointing matrix. Lets calculate the pointing matrix for each prior class
End of explanation
"""


prior250.upper_lim_map()
prior350.upper_lim_map()
prior500.upper_lim_map()

"""
Explanation: Default prior on flux is a uniform distribution, with a minimum and maximum of 0.00 and 1000.0 $\mathrm{mJy}$ respectively for each source. running the function upper_lim _map resets the upper limit to the maximum flux value (plus a 5 sigma Background value) found in the map in which the source makes a contribution to.
End of explanation
"""


%%time 
from xidplus.numpyro_fit import SED_prior
fit_numpyro=SED_prior.all_bands([prior250,prior350,prior500],'../../../../tests/GB_emulator_20210106.pkl')

posterior_numpyro=xidplus.posterior_numpyro_sed(fit_numpyro,[prior250,prior350,prior500],'../../../../tests/GB_emulator_20210106.pkl')
xidplus.save([prior250,prior350,prior500],posterior_numpyro,'test_numpyro_sed')

"""
Explanation: You can fit with the numpyro backend.
End of explanation
"""


plt.hist(posterior_numpyro.samples['src_f'][:,2,:].flatten(),normed=True,alpha=0.5);
plt.hist(posterior_numpyro.samples['src_f'][posterior_numpyro.divergences,2,:].flatten(),normed=True,alpha=0.5);
plt.xlabel('Flux')

fig,axes=plt.subplots(1,3,figsize=(10,5))
for i in range(0,3):
    axes[i].plot(posterior_numpyro.samples['sigma_conf'][:,i],posterior_numpyro.samples['bkg'][:,i],'o')
    axes[i].plot(posterior_numpyro.samples['sigma_conf'][posterior_numpyro.divergences,i],
             posterior_numpyro.samples['bkg'][posterior_numpyro.divergences,i],'o')

def inverse_transform_parameters(param):
    LIR_sd=(15.0-8)/np.sqrt(12.0)
    LIR_mean=0.5*(15+8.0)
    
    red_sd=(6-0.01)/np.sqrt(12.0)
    red_mean=0.5*(0.01+6)
    
    temp_sd=(70-10.0)/np.sqrt(12.0)
    temp_mean=0.5*(70+10.0)
    
    LIR=param[0]*LIR_sd+LIR_mean
    red=param[1]*red_sd+red_mean
    temp=param[2]*temp_sd+temp_mean
    return np.vstack((LIR,red,temp)).T

import pandas as pd
params=posterior_numpyro.samples['params'][::100,:,:]

df=pd.DataFrame(np.asarray(params.reshape(-1, params.shape[-1])),columns=['LIR','z','T'])
g=sns.PairGrid(df)
g.map_diag(plt.hist)
#g.map_lower(sns.kdeplot, cmap="Reds",alpha=0.8,n_levels=10,normed=True, shade=True,shade_lowest=False)
g.map_lower(plt.scatter,color='red',alpha=0.5)
params=posterior_numpyro.samples['params'][posterior_numpyro.divergences,:,:]

df=pd.DataFrame(np.asarray(params.reshape(-1, params.shape[-1])),columns=['LIR','z','T'])
g.data=df
g.map_diag(plt.hist)
#g.map_lower(sns.kdeplot,cmap="Blues",alpha=0.5,n_levels=10,normed=True, shade=True,shade_lowest=False)
g.map_lower(plt.scatter,color='blue',alpha=0.5)


inverse_transform_parameters(params.reshape(-1, params.shape[-1]).T)

import pandas as pd
#params=inverse_transform_parameters(posterior_numpyro.samples['params'][:,:,:].T).reshape(2000,51,3)[::100,:,:]
params=posterior_numpyro.samples['params'][::100,:,:]


df=pd.DataFrame(inverse_transform_parameters(params.reshape(-1, params.shape[-1]).T),columns=['LIR','z','T'])
g=sns.PairGrid(df)
g.map_diag(plt.hist)
#g.map_lower(sns.kdeplot, cmap="Reds",alpha=0.8,n_levels=10,normed=True, shade=True,shade_lowest=False)
g.map_lower(plt.scatter,color='red',alpha=0.5)

params=posterior_numpyro.samples['params'][posterior_numpyro.divergences,:,:]

df=pd.DataFrame(inverse_transform_parameters(params.reshape(-1, params.shape[-1]).T),columns=['LIR','z','T'])
g.data=df
g.map_diag(plt.hist)
#g.map_lower(sns.kdeplot,cmap="Blues",alpha=0.5,n_levels=10,normed=True, shade=True,shade_lowest=False)
g.map_lower(plt.scatter,color='blue',alpha=0.5)


inverse_transform_parameters(params).shape

posterior_numpyro.samples['params'].shape

"""
Explanation: Where are the divergences
End of explanation
"""


priors=[prior250,prior350,prior500]
import copy
priors_prior_pred=copy.deepcopy(priors)
for i in range(0,3):
    priors_prior_pred[i].sim=None
fit_numpyro_prior=SED_prior.all_bands(priors_prior_pred,'../../../../tests/GB_emulator_20210106.pkl')

posterior_numpyro_prior=xidplus.posterior_numpyro_sed(fit_numpyro,[prior250,prior350,prior500],'../../../../tests/GB_emulator_20210106.pkl')
xidplus.save([prior250,prior350,prior500],posterior_numpyro,'test_numpyro_sed_prior_pred')

from numpyro.infer import Predictive
from jax import random, vmap

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)

post_predictive=Predictive(SED_prior.spire_model,fit_numpyro.get_samples())

post_pred=post_predictive(random.PRNGKey(0),priors,SED_prior.load_emulator('../../../../tests/GB_emulator_20210106.pkl'))

prior_predictive=Predictive(SED_prior.spire_model,posterior_samples = {}, num_samples = 100,)

prior_pred=prior_predictive(random.PRNGKey(0),priors_prior_pred,SED_prior.load_emulator('../../../../tests/GB_emulator_20210106.pkl'))

plt.hist(prior_pred['obs_psw'].flatten(),alpha=0.5,normed=True);
plt.hist(post_pred['obs_psw'].flatten(),alpha=0.5,normed=True);

prior_pred['obs_psw'][0,:]

plt.hist(prior_pred['bkg'][:,0],normed=True,alpha=0.5);
plt.hist(posterior_numpyro.samples['bkg'][:,0],normed=True,alpha=0.5);

plt.hist(np.log10(prior_pred['obs_psw'][0,:]))

"""
Explanation: Generate Prior Predictive distribution
End of explanation
"""


hdulists=[xidplus.posterior_maps.make_fits_image(prior250,prior_pred['obs_psw'][1,:]),
         xidplus.posterior_maps.make_fits_image(prior350,prior_pred['obs_pmw'][1,:]),
         xidplus.posterior_maps.make_fits_image(prior500,prior_pred['obs_plw'][1,:])]
cmap=sns.cubehelix_palette(8, start=.5, rot=-.75,as_cmap=True)

         
fig = plt.figure(figsize=(10*len(priors),10))
figs=[]
for i in range(0,len(priors)):
    figs.append(aplpy.FITSFigure(hdulists[i][1],figure=fig,subplot=(1,len(priors),i+1)))

for i in range(0,len(priors)):
    vmin=np.min(priors[i].sim)
    vmax=np.max(priors[i].sim)
    figs[i].show_colorscale(vmin=vmin,vmax=vmax,cmap=cmap)
    figs[i].show_markers(priors[i].sra, priors[i].sdec, edgecolor='black', facecolor='black',
            marker='o', s=20, alpha=0.5)
    figs[i].tick_labels.set_xformat('dd.dd')
    figs[i].tick_labels.set_yformat('dd.dd')
    figs[i].add_colorbar()
    figs[i].colorbar.set_location('top')

hdulists=[xidplus.posterior_maps.make_fits_image(prior250,post_pred['obs_psw'][0,:]),
         xidplus.posterior_maps.make_fits_image(prior350,post_pred['obs_pmw'][0,:]),
         xidplus.posterior_maps.make_fits_image(prior500,post_pred['obs_plw'][0,:])]
cmap=sns.cubehelix_palette(8, start=.5, rot=-.75,as_cmap=True)

         
fig = plt.figure(figsize=(10*len(priors),10))
figs=[]
for i in range(0,len(priors)):
    figs.append(aplpy.FITSFigure(hdulists[i][1],figure=fig,subplot=(1,len(priors),i+1)))

for i in range(0,len(priors)):
    vmin=np.min(priors[i].sim)
    vmax=np.max(priors[i].sim)
    figs[i].show_colorscale(vmin=vmin,vmax=vmax,cmap=cmap)
    figs[i].show_markers(priors[i].sra, priors[i].sdec, edgecolor='black', facecolor='black',
            marker='o', s=20, alpha=0.5)
    figs[i].tick_labels.set_xformat('dd.dd')
    figs[i].tick_labels.set_yformat('dd.dd')
    figs[i].add_colorbar()
    figs[i].colorbar.set_location('top')

"""
Explanation: Make replicated map
End of explanation
"""


