import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import itertools

from astropy import units as u
import astropy.constants as const
from literature_attenuation import *


def find_nearest(array,value):
    idx = (np.abs(np.array(array)-value)).argmin()
    return idx




data = np.load('ext_dmsf_mw_tseris.npz')

#first set up sizes plot
palette = itertools.cycle(sns.color_palette('rocket'))


fig = plt.figure()
ax = fig.add_subplot(111)

yr_list = ['700Myr','1000Myr','1300Myr','1600Myr']
year_labels = ['700 Myr','1000 Myr','1.4 Gyr','1.6 Gyr']

for counter,yr in enumerate(yr_list):
    ax.loglog(data['a'],data['a3N_'+yr],lw=4,color=next(palette),label=year_labels[counter])
    #sns.lineplot(data['a'],data['a3N_'+yr],lw=2,color=next(palette))
ax.set_xlabel(r'Size ($\mu$m)',fontsize=14)
ax.set_ylabel(r'4$\pi \rho a^3/3\mathrm{M_d} \times \partial(n/\partial_\mathrm{ log} a \ (\mathrm{M_\odot})$',fontsize=14)

plt.legend(loc=4,fontsize=16)
fig.savefig('sizes.png',dpi=300)



#second, extinction plot
#reset the palette
palette = itertools.cycle(sns.color_palette('rocket'))

wave = np.linspace(0.1,1,1000) #micron
x = 1./wave

v_idx = find_nearest(wave,0.5500) #5500 angstrom
b_idx = find_nearest(wave,0.445) #4450 angstrom

fig = plt.figure()
ax = fig.add_subplot(111)



Rv_array = np.linspace(2,5,500)
norm = matplotlib.colors.Normalize(
    vmin = np.min(Rv_array),
    vmax = np.max(Rv_array))
c_m = matplotlib.cm.viridis_r
s_m  = matplotlib.cm.ScalarMappable(cmap = c_m,norm=norm)
s_m.set_array([])

for idx,Rv in enumerate(Rv_array):
    tau = cardelli(wave*1.e4) #required input is in angstrom
    ax.plot(x,cardelli(wave*1.e4,R_v = Rv),alpha=0.075,color='grey')#s_m.to_rgba(Rv))
#cb = fig.colorbar(s_m,orientation='vertical')
#cb.set_label(r'R$_\mathrm{V}$ for Cardelli et al. (1989) MW Parameterization',fontsize=8)
#cb.ax.tick_params(labelsize=8)


#just plot average now for cardelli
Rv = 3.1
tau = cardelli(wave*1.e4)
#ax.plot(x,cardelli(wave*1.e4,R_v = Rv),lw=3,color='black',label=r'Galactic Average Extinction Curve: R$_\mathrm{V} = 3.1$')


#now plot the simulations extinction curve
for counter,yr in enumerate(yr_list):
    if counter == len(yr_list)-1:
        ax.plot(data['wlen_inverse'],data['A_'+yr],lw=1,color=next(palette),alpha=0.75,label='Pilot Study')
    else:
        ax.plot(data['wlen_inverse'],data['A_'+yr],lw=1,color=next(palette),alpha=0.75)
        
ax.set_xlabel(r' 1/$\mu$m',fontsize=14)
ax.set_ylabel(r'A/A_\mathrm{V}',fontsize=14)


tau_smc = smc(wave*1.e4)
tau_lmc = lmc(wave*1.e4)

ax.plot(x,tau_smc,color='dodgerblue',lw=4,label='SMC observed')
ax.plot(x,tau_lmc,color='seagreen',lw=4,label='LMC observed')

ax.set_xlabel(r'x $\equiv 1/\lambda (\mu$m)',fontsize=14)
ax.set_ylabel(r'$\tau$',fontsize=14)

#dummy shaded region plot for label
ax.plot([-1,-1],lw=5,color='grey',label='MW observed range')


ax.set_xlim([1,10])
ax.set_ylim([1,8])
plt.legend(loc=2,fontsize=16)
fig.savefig('mw_only.png',dpi=300)




#----------------------------------------
#final plot: slope-z
#----------------------------------------

data = np.load('slope_z.npz')
fig = plt.figure()
ax = fig.add_subplot(111)

#ax.scatter(data['logZ'],data['slope'])
sns.kdeplot(data['logZ']+0.1,data['slope'],fill=True,palette='rocket')
ax.set_xlabel(r'log$_\mathrm{10}(Z)$',fontsize=14)
ax.set_ylabel(r'Extinction Law Slope',fontsize=14)
fig.savefig('slope_Z.png',dpi=300)