
from astropy.cosmology import WMAP9 as cosmo
import jax
import numpy as onp
import pylab as plt
import astropy.units as u
import scipy.integrate as integrate
%matplotlib inline

import jax.numpy as np
from jax import grad, jit, vmap, value_and_grad
from jax import random
from jax import vmap # for auto-vectorizing functions
from functools import partial # for use with vmap
from jax import jit # for compiling functions for speedup
from jax.experimental import stax # neural network library
from jax.experimental.stax import Conv, Dense, MaxPool, Relu, Flatten, LogSoftmax, LeakyRelu # neural network layers

from jax.experimental import optimizers
from jax.tree_util import tree_multimap  # Element-wise manipulation of collections of numpy arrays 
import matplotlib.pyplot as plt # visualization
# Generate key which is used to generate random numbers
key = random.PRNGKey(2)
from xidplus import cigale

onp.random.seed(2)

"""
Explanation: The JAX emulator: CIGALE prototype
In this notebook, I will prototype my idea for emulating radiative transfer codes with a Deepnet in order for it to be used inside xidplus. As numpyro uses JAX, the Deepnet wil ideally be trained with a JAX network. I will use CIGALE
Advice from Kasia
Use the following modules:
* Dale 2014 dust module with one parameter ($\alpha$) however, $\alpha$ can only take certian values in Cigale
    * 0.0625, 0.1250, 0.1875, 0.2500,0.3125, 0.3750, 0.4375, 0.5000, 0.5625, 0.6250, 0.6875, 0.7500,0.8125, 0.8750, 0.9375, 1.0000, 1.0625, 1.1250, 1.1875, 1.2500,1.3125, 1.3750, 1.4375, 1.5000, 1.5625, 1.6250, 1.6875, 1.7500, 1.8125, 1.8750, 1.9375, 2.0000, 2.0625, 2.1250, 2.1875, 2.2500,2.3125, 2.3750, 2.4375, 2.5000, 2.5625, 2.6250, 2.6875, 2.7500,2.8125, 2.8750, 2.9375, 3.0000, 3.0625, 3.1250, 3.1875, 3.2500, 3.3125, 3.3750, 3.4375, 3.5000, 3.5625, 3.6250, 3.6875, 3.7500, 3.8125, 3.8750, 3.9375, 4.0000
* sfhdelayed starforamtion history module. Has parameters $\tau$ (500-6500) ($age$ can be calculated from redshift). $f_{burst}$ is set to 0
* bc03stellar population synthesis module (don't change parameters)
* dustatt_2powerlaws 
    * set $Av_BC$ the V band attenuation in the birth clouds to between 0 - 4
    * set BC_to_ISM_factor to 0.7
Final parameters: $alpha$, $AV_BC$,$\tau$,$z$,$SFR$,$AGN$
Ideally, I would generate values from prior. I can do that for $AV_BC$,$\tau$,$z$,$SFR$,$AGN$ but not $\alpha$ given that there are fixed values.
End of explanation
"""


from astropy.io import fits
from astropy.table import Table
import scipy.stats as stats

alpha=onp.array([0.0625, 0.1250, 0.1875, 0.2500,0.3125, 0.3750, 0.4375, 0.5000, 0.5625, 0.6250, 0.6875, 0.7500,0.8125, 0.8750, 0.9375, 1.0000, 1.0625, 1.1250, 1.1875, 1.2500,1.3125, 1.3750, 1.4375, 1.5000, 1.5625, 1.6250, 1.6875, 1.7500, 1.8125, 1.8750, 1.9375, 2.0000, 2.0625, 2.1250, 2.1875, 2.2500,2.3125, 2.3750, 2.4375, 2.5000, 2.5625, 2.6250, 2.6875, 2.7500,2.8125, 2.8750, 2.9375, 3.0000, 3.0625, 3.1250, 3.1875, 3.2500, 3.3125, 3.3750, 3.4375, 3.5000, 3.5625, 3.6250, 3.6875, 3.7500, 3.8125, 3.8750, 3.9375, 4.0000])
alpha_rv = stats.randint(0, len(alpha))
av_bc_rv=stats.uniform(0.1,4.0)
tau_rv=stats.randint(500,6500)
z_rv=stats.uniform(0.01,6)
sfr_rv=stats.loguniform(0.01,30000)
agn_frac_rv=stats.beta(1,3)

from astropy.cosmology import Planck13

z=z_rv.rvs(1)[0]
onp.int(Planck13.age(z).value*1000)
alpha[alpha_rv.rvs(1)[0]]

nsamp=1
from astropy.constants import L_sun, M_sun
from astropy.table import vstack
col_scale=['spire_250','spire_350','spire_500','dust.luminosity','sfh.sfr','stellar.m_star']
parameter_names=onp.array(['tau_main','age_main','Av_BC','alpha','fracAGN','redshift'])
all_SEDs=[]
for i in range(0,nsamp):
    z=z_rv.rvs(1)[0]
    
    parameters={'tau_main':[tau_rv.rvs(1)[0]],'age_main':[onp.int(Planck13.age(z).value*1000)],
                'Av_BC':[av_bc_rv.rvs(1)[0]],'alpha':[alpha[alpha_rv.rvs(1)[0]]],'fracAGN':[agn_frac_rv.rvs(1)[0]],'redshift':[z]}
    path_to_cigale='/Volumes/pdh_storage/cigale/'
    path_to_ini_file='pcigale_kasia_nn.ini'
    SEDs=cigale.generate_SEDs(parameter_names, parameters, path_to_cigale, path_to_ini_file, filename = 'tmp_single')
    #set more appropriate units for dust

    SEDs['dust.luminosity']=SEDs['dust.luminosity']/L_sun.value
    scale=1.0/SEDs['sfh.sfr']
    for c in col_scale:
        SEDs[c]=SEDs[c]*scale*sfr_rv.rvs(1)[0]
    all_SEDs.append(SEDs)
    if i and i % 100 == 0:
        tmp_SEDs=vstack(all_SEDs)
        tmp_SEDs.write('kasia_gen_SEDs_{}.fits'.format(i),overwrite=True)
        all_SEDs=[]

all_SEDs[0]

"""
Explanation: Generate CIGALE SEDs
End of explanation
"""


onp.array2string(10.0**np.arange(-2.5,0.77,0.1), separator=',',formatter={'float_kind':lambda x: "%.4f" % x}).replace('\n','')

onp.array2string(np.arange(0.1,4,0.3),separator=',',formatter={'float_kind':lambda x: "%.4f" % x}).replace('\n','')

"""
Explanation: Generate values for CIGALE
Redshift
End of explanation
"""


onp.array2string(np.arange(0.001,1,0.075),separator=',',formatter={'float_kind':lambda x: "%.3f" % x}).replace('\n','')

SEDs=Table.read('/Volumes/pdh_storage/cigale/out/models-block-0.fits')
#set more appropriate units for dust
from astropy.constants import L_sun, M_sun
SEDs['dust.luminosity']=SEDs['dust.luminosity']/L_sun.value


SEDs=SEDs[onp.isfinite(SEDs['spire_250'])]

SEDs

from astropy.table import vstack

(1.0/dataset['sfh.sfr'])*dataset['sfh.sfr']*10.0**scale_table

# define a range of scales
scale=np.arange(8,14,0.25)
#repeat the SED table by the number of scale steps
dataset=vstack([SEDs for i in range(0,scale.size)])
#repeat the scale range by the number of entries in table (so I can easily multiply each column)
scale_table=np.repeat(scale,len(SEDs))
#parameters to scale
col_scale=['spire_250','spire_350','spire_500','dust.luminosity','sfh.sfr','stellar.m_star']
for c in col_scale:
    dataset[c]=dataset[c]*10.0**scale_table
dataset['log10_sfh.sfr']=onp.log10(dataset['sfh.sfr'])
dataset['log10_universe.redshift']=onp.log10(dataset['universe.redshift'])
# transform AGN fraction to logit scale
dataset['logit_agnfrac']=onp.log(dataset['agn.fracAGN']/(1-dataset['agn.fracAGN']))
#shuffle dataset
dataset=dataset[onp.random.choice(len(dataset), len(dataset), replace=False)]

plt.hist(dataset['log10_sfh.sfr'],bins=(np.arange(0,14)));


dataset

"""
Explanation: AGN frac
End of explanation
"""


dataset=dataset[0:18000000]

len(dataset)/1200

split=0.75
inner_batch_size=1200

train_ind=onp.round(0.75*len(dataset)).astype(int)
train=dataset[0:train_ind]
validation=dataset[train_ind:]
input_cols=['log10_sfh.sfr','agn.fracAGN','universe.redshift', 'attenuation.Av_BC','dust.alpha','sfh.tau_main']
output_cols=['spire_250','spire_350','spire_500']
train_batch_X=np.asarray([i.data for i in train[input_cols].values()]).reshape(len(input_cols)
                                                                               ,inner_batch_size,onp.round(len(train)/inner_batch_size).astype(int)).T.astype(float)
train_batch_Y=np.asarray([np.log(i.data) for i in train[output_cols].values()]).reshape(len(output_cols),
                                                                                          inner_batch_size,onp.round(len(train)/inner_batch_size).astype(int)).T.astype(float)
validation_batch_X=np.asarray([i.data for i in validation[input_cols].values()]).reshape(len(input_cols)
                                                                               ,inner_batch_size,onp.round(len(validation)/inner_batch_size).astype(int)).T.astype(float)
validation_batch_Y=np.asarray([np.log(i.data) for i in validation[output_cols].values()]).reshape(len(output_cols),
                                                                                          inner_batch_size,onp.round(len(validation)/inner_batch_size).astype(int)).T.astype(float)

# Use stax to set up network initialization and evaluation functions
net_init, net_apply = stax.serial(
    Dense(128), LeakyRelu,
    Dense(128), LeakyRelu,
    Dense(128), LeakyRelu,
    Dense(128), Relu,
    Dense(len(output_cols))
)
in_shape = (-1, len(input_cols),)
out_shape, net_params = net_init(key,in_shape)

def loss(params, inputs, targets):
    # Computes average loss for the batch
    predictions = net_apply(params, inputs)
    return np.mean((targets - predictions)**2)

def batch_loss(p,x_b,y_b):
    loss_b=vmap(partial(loss,p))(x_b,y_b)
    return np.mean(loss_b)

opt_init, opt_update, get_params= optimizers.adam(step_size=5e-4)
out_shape, net_params = net_init(key,in_shape)
opt_state = opt_init(net_params)

@jit
def step(i, opt_state, x1, y1):
    p = get_params(opt_state)
    g = grad(batch_loss)(p, x1, y1)
    loss_tmp=batch_loss(p,x1,y1)
    return opt_update(i, g, opt_state),loss_tmp

np_batched_loss_1 = []
valid_loss=[]
for i in range(10000):
    opt_state, l = step(i, opt_state, train_batch_X, train_batch_Y)
    p = get_params(opt_state)
    valid_loss.append(batch_loss(p,validation_batch_X,validation_batch_Y))
    np_batched_loss_1.append(l)
    if i % 100 == 0:
        print(i)
net_params = get_params(opt_state)

for i in range(2000):
    opt_state, l = step(i, opt_state, train_batch_X, train_batch_Y)
    p = get_params(opt_state)
    valid_loss.append(batch_loss(p,validation_batch_X,validation_batch_Y))
    np_batched_loss_1.append(l)
    if i % 100 == 0:
        print(i)
net_params = get_params(opt_state)

plt.figure(figsize=(20,10))
plt.semilogy(np_batched_loss_1,label='Training loss')
plt.semilogy(valid_loss,label='Validation loss')
plt.xlabel('Iteration')
plt.ylabel('Loss (MSE)')
plt.legend()


"""
Explanation: DeepNet building
I will build a multi input, multi output deepnet model as my emulator, with parameters as input and the observed flux as outputs. I will train on log10 flux to make the model easier to train, and have already standarised the input parameters. I wilkl be using stax which can be thought of as the Keras equivalent for JAX. This blog was useful starting point.
I will use batches to help train the network
End of explanation
"""


net_params = get_params(opt_state)
predictions = net_apply(net_params,validation_batch_X)

validation_batch_X.shape

validation_batch_X[0,:,:].shape

res=((np.exp(predictions)-np.exp(validation_batch_Y))/(np.exp(validation_batch_Y)))
fig,axes=plt.subplots(1,len(output_cols),figsize=(50,len(output_cols)))
for i in range(0,len(output_cols)):
    axes[i].hist(res[:,:,i].flatten()*100.0,np.arange(-10,10,0.1))
    axes[i].set_title(output_cols[i])
    axes[i].set_xlabel(r'$\frac{f_{pred} - f_{True}}{f_{True}} \ \%$ error')
plt.subplots_adjust(wspace=0.5)

"""
Explanation: Investigate performance of each band of emulator
To visulise performance of the trainied emulator, I will show the difference between real and emulated for each band.
End of explanation
"""


import cloudpickle


with open('CIGALE_emulator_20210330_log10sfr_uniformAGN_z.pkl', 'wb') as f:
    cloudpickle.dump({'net_init':net_init,'net_apply': net_apply,'params':net_params}, f)
    
net_init, net_apply

"""
Explanation: Save network
Having trained and validated network, I need to save the network and relevant functions
End of explanation
"""


wave=np.array([250,350,500])

plt.loglog(wave,np.exp(net_apply(net_params,np.array([2.95, 0.801, 0.1]))),'o')
#plt.loglog(wave,10.0**net_apply(net_params,np.array([3.0,0.0,0.0])),'o')
plt.loglog(wave,dataset[(dataset['universe.redshift']==0.1) & (dataset['agn.fracAGN'] == 0.801) & (dataset['sfh.sfr']>900) & (dataset['sfh.sfr']<1100)][output_cols].values())

dataset[(dataset['universe.redshift']==0.1) & (dataset['agn.fracAGN'] == 0.801) & (dataset['sfh.sfr']>900) & (dataset['sfh.sfr']<1100)]

import xidplus

from xidplus.numpyro_fit.misc import load_emulator

obj=load_emulator('CIGALE_emulator_20210330_log10sfr_uniformAGN_z.pkl')

type(obj['params'])

import json

import numpy as np

np.savez('CIGALE_emulator_20210610_kasia',obj['params'],allow_pickle=True)

ls

x=np.load('params_save.npz',allow_pickle=True)

x['arr_0'].tolist()

obj['params']

"""
Explanation: Does SED look right?
End of explanation
"""


