
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(1)

"""
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
End of explanation
"""


from astropy.io import fits
from astropy.table import Table

"""
Explanation: Read in CIGALE SEDs
End of explanation
"""


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

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

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


onp.array2string(np.arange(0.001,1,0.05),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


from astropy.cosmology import WMAP9 as cosmo

SEDs['age']=cosmo.age(SEDs['universe.redshift'])

SEDs

plt.plot(SEDs['age'],SEDs['sfh.age_main']/1000.0,'o')
plt.xlabel('age')
plt.ylabel('age_main')

SEDs=SEDs[SEDs['age']>SEDs['sfh.age_main']/1000.0]

from astropy.table import vstack

# define a range of scales
scale=np.arange(8,14,0.1)
#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'])
dataset['log10_dust.luminosity']=onp.log10(dataset['dust.luminosity'])
# 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)]

dataset


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


print(len(input_cols),inner_batch_size,onp.round(len(train)/inner_batch_size),120*25457,len(train))

0.75*len(dataset)

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_dust.luminosity','agn.fracAGN','universe.redshift','sfh.age_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)

print(train)

# Use stax to set up network initialization and evaluation functions
net_init, net_apply = stax.serial(
    Dense(128), LeakyRelu,
    Dense(128), LeakyRelu,
    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_20210420_log10sfr_uniformAGN_z',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
"""


