import pandas as pd
import sys
from os import listdir
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad
from scipy import interpolate
from scipy.stats import chisquare
from timeit import default_timer as timer
from scipy.optimize import curve_fit
import functions as fn
from model import Model, Spectral

s = Spectral(sys.argv[1])
s_norm = Spectral(sys.argv[1])
s_norm.cut_spectrum(end=17)
s.cut_spectrum(begin=5.24)
if sys.argv[1] =='DF-Tau':     #Removal of wavelengths with outlying fluxes
    s.cut_spectrum(begin=5.54)
    s_norm.cut_spectrum(begin=5.54)

MN_path = sys.argv[2]

try:
    print(sys.argv[3])
    """new data retrieval from Analyzer"""
    Inner_Rim = pd.DataFrame(np.array([fn.get_1sigma(sys.argv[2],'C1','short'),fn.get_1sigma(sys.argv[2],'Trim_min','short')
                            ,fn.get_1sigma(sys.argv[2],'qrim','short')]),columns=['Mean','Sigma','sh'])
    Inner_Rim = Inner_Rim['Mean']
    """--------------------------------"""
    print(Inner_Rim)
    n_params = 20
    s.cut_spectrum(17, 35)
    region = 'long'

except:
    n_params = 23
    s.cut_spectrum(end=17)
    region = 'short'
    print("Short wavelengths")
factor = s_norm.spitzer['flux'].max()/s.spitzer['flux'].max()
s.setunity()
s.rebin_starfit()
s.read_opacity()
s.restore_index()
MN_path = sys.argv[2]
MN_output = fn.get_all1sigma(sys.argv[2],n_params,region)
# MN_output = pd.read_csv(MN_path,sep='\s+',skiprows=2,nrows=n_params,index_col=False,usecols=[1,2])
# MN_output.columns = ['Mean','Sigma']
print(MN_output)
statistic = pd.read_csv(MN_path,sep='\s+',nrows=1)
loglike = round(float(statistic.columns[5]))
parameters_values = MN_output['Mean']

try:
    parameters_values= pd.concat([Inner_Rim,parameters_values],ignore_index=True)
except:
    print("shortwave")
print(parameters_values)

start = timer()
f = Model(s.spitzer['wavelength'],s.opacities)
f.define_parameters(parameters_values)
f.compute()
print(s.spitzer['wavelength'].max())
if s.spitzer['wavelength'].max()>=17:
    print(factor)
    f.longwave_rim(factor,s_norm.spitzer['wavelength'])

print(timer()-start)

"""---------------------------------------------"""
"""PLOTS"""
fig,axes = plt.subplots(2,1,figsize=(8,8),gridspec_kw={'height_ratios': [6, 2]})

axes[0].errorbar(s.spitzer['wavelength'],s.spitzer['flux'],s.spitzer['sigma'],label='Spitzer',color='black')

axes[0].plot(s.spitzer['wavelength'], s.stellar,label='star',color='green')
axes[0].plot(s.spitzer['wavelength'], f.components['midplane'],label='Midplane')
axes[0].plot(s.spitzer['wavelength'], f.components['inner_rim'],label='Inner Rim')
axes[0].plot(s.spitzer['wavelength'], f.components['atmosphere'],label = 'MODEL dust',color='orange')

axes[0].plot(s.spitzer['wavelength'], f.components['disk_model']+s.stellar,linewidth=2, label = 'MODEL',color='r')

axes[1].plot(s.spitzer['wavelength'],(s.spitzer['flux']-f.components['disk_model']-s.stellar)/s.spitzer['flux'],label = 'Residual',color='black')

axes[0].set_ylabel('$\\frac{F}{F_{max}}$',size=20)
axes[1].set_xlabel('$\lambda$ ($\mu$m)')
axes[1].set_ylabel('$\\frac{\Delta F}{F}$')
axes[0].legend()
axes[1].legend()
axes[0].set_title('Log-Likelihood: ' + str(loglike),fontsize=10 )
if n_params == 23:
    axes[0].set_title('5-17 $\mu m$' )
else:
    axes[0].set_title('17-35 $\mu m$' )
plt.subplots_adjust(hspace=0)
plt.show()
# fig.savefig(MN_path[:-14] + '/fitfig.png', dpi=fig.dpi)
# Model_data = pd.DataFrame()
# Model_data['wavelength'] = s.spitzer['wavelength']
# Model_data['flux'] = f.components['disk_model']
# Model_data.to_csv('/home/chris/Documents/Thesis/jupbooks/daniele/JulyRun/Model_spectra/' + sys.argv[1] + '.csv')

"""Dust Composition plot"""
kap = [0.1*3,2*3,5*3,0.1*2.8,2*2.8,5*2.8,0.1*3.2,2*3.2,0.1*2.8,2*2.8,0.1*2.2,2*2.2,5*2.2]
dust_parameters,dust_sigma_min, dust_sigma_plus = fn.dust_sigma(MN_output.iloc[-13:])
print(MN_output.iloc[-13:])
fig2,axes = plt.subplots(1,1,figsize=(8,8))
olivine = axes.bar(['.1 ','2 ','5 '],dust_parameters[0:3],yerr=[dust_sigma_min[0:3],dust_sigma_plus[0:3]]
            ,color='r',label='Olivine')
pyroxene = axes.bar([' .1 ',' 2 ',' 5 '],dust_parameters[3:6],yerr=[dust_sigma_min[3:6],dust_sigma_plus[3:6]]
            ,color='b',label='Pyroxene')
forsterite = axes.bar(['.1','2'],dust_parameters[6:8],yerr=[dust_sigma_min[6:8],dust_sigma_plus[6:8]]
            ,color='g',label='Forsterite')
enstatite = axes.bar([' .1',' 2'],dust_parameters[8:10],yerr=[dust_sigma_min[8:10],dust_sigma_plus[8:10]]
            ,color='y',label='Enstatite')
silica = axes.bar([' .1  ',' 2  ',' 5  '],dust_parameters[10:],yerr=[dust_sigma_min[10:],dust_sigma_plus[10:]]
            ,color='purple',label='Silica')

axes.set_ylim(bottom=0,top=1)
axes.set_ylabel('Mass fraction')
axes.set_xlabel('Grain size ($\mu$m)')
axes.set_title(sys.argv[1])
axes.set_title('Log-Likelihood: ' + str(loglike),fontsize=10)
axes.legend(fontsize=20)
plt.xticks(fontsize=15)
plt.show()
# fig2.savefig(MN_path[:-14] + '/dustspecies.png', dpi=fig.dpi)
