#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul 23 17:31:30 2020

@author: jonty/sascha
"""

import numpy as np
import matplotlib.pyplot as plt
from astropy.io import ascii
from extinction import apply, fitzpatrick99
from scipy import interpolate
import math
import glob
from astropy.io import fits
from spectres import spectres

# constants
pc  = 3.086e16 # m
rsol = 696340e3 # m

#Define stellar spectrum

#
# Here we want a generalised function to read in a stellar spectrum from the subfolder.
#
#B0V T = 30000 K, log g = 4, Rstar = 10.0 Rsol
#B5V T = 15200 K, log g = 4, Rstar = 3.9 Rsol
#B9V T = 10600 K, log g = 4, Rstar = 2.7 Rsol
#B2V T = 20600 K, log g = 4, Rstar = 3.85 Rsol
#O8V T = 35000 K, log g = 4, Rstar = 8.75 Rsol


def read_stellar_spectrum(dstar=1000,spectrum='./stellar_spectra/BT-NEXTGEN_Photosphere_Teff_15000_logg_4.0_Rstar_3.9.spec'):
    spec_trim = spectrum[:-5]
    spec_split = spec_trim.split('_')
    #print(spec_split)
    tstar = float(spec_split[4])
    logg  = float(spec_split[6])
    rstar = float(spec_split[8])

    data = ascii.read(spectrum,comment='#',names=['Wave','Flux'])

    wavelengths =  data['Wave'].data #Angstroms
    model_spect =  data['Flux'].data #Ergs/cm**2/s/A
    if spec_split[0] == 'BOSZ':
        model_spect *= np.pi #factor of pi *only* as already scaled by 4x

    wav_um = wavelengths / 1e4
    flx_mjy = model_spect *  1e3 * wavelengths**2  * 3.33564095E+04 * ((rstar*rsol)**2 / (dstar*pc)**2) #coversion to mJy

    return wav_um,flx_mjy


# get list of stellar spectra filenames from subfolder

spectra_list = glob.glob("./stellar_spectra/*.spec")

# print list of filenames

for i in range(len(spectra_list)):
    print(i,spectra_list[i])

# ask user for desired spectrum, distance

spec_index  = int(input("Please enter index of desired stellar spectrum:"))
spec = spectra_list[spec_index]

dstar = float(input("Please enter the distance (in pc) of the star:"))

# read in stellar spectrum
wav_um, flx_mjy = read_stellar_spectrum(dstar,spec)
wav_aa = wav_um*1e4 # wavelengths in angstroms for the extinction calculation

spectrum_total_int = 0

# Set the depletion of silicon
depletion_si = float(input("Please enter the depletion of silicon:"))

# Set the depletion of carbon
depletion_carbon = float(input("Please enter the depletion of carbon:"))

# get list of stellar spectra filenames from subfolder

ice_list = glob.glob("./water_ice/*.q")
print(ice_list)
# print list of filenames
for i in range(len(ice_list)):
    print(i,ice_list[i])

water_ice_yesno = float(input("Please enter 0 for no water ice and 1 for including water ice (10 ppm abundance):"))
if water_ice_yesno==1:
    ice_index = input("Please enter index of desired ice component:")

    # Set the abundance of ices in the diffuse ISM

    ice_species=[]
    ice_species_string = ""
    ice_species.append(ice_list[int(ice_index)])
    ice_species_string = ice_species[0]#[12:]+"_"


# Set the fraction of small silicates
small_silicates_fraction = float(input("Please enter the fraction of silicates in nano silicates (recommended 0.03-0.1):"))

#
# Here we want a function to select the dust component(s) we want to add to the extincted stellar model
#
def read_dust_qext(species='mg05fe05sio3_am_cde01.q'):
    # read in the Q values
    data_dust=np.loadtxt(species)

    wavelengths_q=data_dust[:,0]
    Qext=data_dust[:,1]

    return wavelengths_q, Qext

# Here we read the properties, e.g. density and molmass of the dust species

def read_dust_prop(species='mg05fe05sio3_am_cde01.q'):
    # read in the Q values
    data_dust=np.loadtxt(species)
    sample = open('./dust_qext/silicate_sample.txt','r')
    linessample=sample.readlines()
    sample_column_number = 0
    resultsample=[]
    for x in linessample:
        resultsample.append(x.split()[sample_column_number])
    sample.close()
    #print(resultsample)

    index = resultsample.index(species)
    name,density,molmass=linessample[index].split()

    return density, molmass

def read_ice_prop(ice_species='water_ice_cde01.q'):
    # read in the Q values
    data_dust=np.loadtxt(ice_species)
    ice_sample = open('./water_ice/ice_sample.txt','r')
    linessample=ice_sample.readlines()
    sample_column_number = 0
    resultsample=[]
    for x in linessample:
        resultsample.append(x.split()[sample_column_number])
    ice_sample.close()
    #print(resultsample)

    index = resultsample.index(ice_species)
    name,ice_density, ice_grain_radius =linessample[index].split()

    return ice_density, ice_grain_radius

#
# Calculate number of atoms along line of sight based on extinction
#
def calculate_Nd(Av,rho,gmol,fraction, depletion_si):
    NH=1.9e21*Av
    si_abundance=3.4e-5
    # total amount of silicon atoms available
    si_los=NH*si_abundance*depletion_si*(1.0-small_silicates_fraction)
    si_los_specie=si_los*fraction
    # calculate the silicon atoms per particle
    #rho=3.2 # check per composition
    rho=rho
    V=4./3.*math.pi*(0.1*1e-4)**3
    grams=V*rho
    gmol=gmol
    #gmol=100. # check per composition
    mol=grams/gmol
    Avogadro=6.0221409e+23
    si_atoms_per_particle=mol*Avogadro
    Nd=si_los/si_atoms_per_particle

    # this needs to be adapted, because it will be different for different species
    # also has the assumption that all the silicon is locked up in dust

    print('Nd',Nd)
    return Nd

def water_ice(Av, ice_grain_radius, rho_ice):#grain radius can be where this code breaks! not yet added
    NH=1.9e21*Av
    water_ice_abundance=10e-6  # water molecules per
    ice_los=NH*water_ice_abundance
    # 0.1 micron grains
    rho_ice=float(rho_ice) # check per composition
    V=4./3.*math.pi*(float(ice_grain_radius)*1e-4)**3
    grams=V*rho_ice
    gmol=18.015281
    mol=grams/gmol
    Avogadro=6.0221409e+23
    water_molecules_per_particle=mol*Avogadro
    N_waterice=ice_los/water_molecules_per_particle
    return N_waterice

def calculate_hydrocarbon_spectrum(Av, depletion_carbon):
    NH=1.9e21*Av
    carbon_abundance=3.91e-4 # from tielens book
    # total amount of silicon atoms available
    carbon_dust_col=carbon_abundance*depletion_carbon*NH
    abs_strength=[24.3, 15.2, 23.7, 14.8, 2.58,0.275,0.275,0.984629]         # from paper Jean Chiar
    # fractions for hydrocarbon
    fractions=np.loadtxt('./hydrocarbon/fractions_chx2.txt')
    wavelength_carbon=np.arange(2,40.0,0.001)
    # read fits file with all the gaussians
    hdulist = fits.open('./hydrocarbon/spectrum_34_micron2.fits')
    hdu = hdulist[0]
    data_gauss34 = fits.getdata('./hydrocarbon/spectrum_34_micron2.fits')
    print(data_gauss34.shape)
    spectrum=np.empty(data_gauss34.shape)

    for i in range(len(fractions)):

        spectrum[:,i]=data_gauss34[:,i]*abs_strength[i]*1e-18*fractions[i]*carbon_dust_col

    tau_total=np.sum(spectrum, axis=1)

    # this needs to be adapted, because it will be different for different species
    # also has the assumption that all the silicon is locked up in dust

    #plt.plot(wavelength_carbon,tau_total)
    return wavelength_carbon, tau_total

def small_silicates(Av,depletion_si,small_silicates_fraction):
    # this returns the optical depth from small silicates
    # the silicates are a mix between olivine and pyroxene grains, we will assume it is 50/50 for now
    NH=1.9e21*Av
    si_abundance=3.9e-5
    si_los=NH*si_abundance*depletion_si*small_silicates_fraction

    si_atoms_per_unit_py=271

    Nd_py=(si_los)/si_atoms_per_unit_py

    data_smallsil_py=np.loadtxt('./nanosil_cross_sec/pyro_0_cross.txt')

    #tau_small_sil_ol=data_smallsil_ol[:,1]*Nd_ol
    tau_small_sil_py=data_smallsil_py[:,1]*Nd_py

    tau_total_sms=tau_small_sil_py

    return data_smallsil_py[:,0], tau_total_sms

#generate noise for a spectrum
def spec_noise(spectrum,noise_para=[0.0,1.0],noise_type='Poisson'):

    if noise_type == 'Poisson':
        noise = np.zeros((len(spectrum)))

        for k in range(len(noise)):
            noise[k] = 1e-3*np.sqrt(1e3*spectrum[k])*np.random.uniform(-1,1)

    if noise_type == 'Gaussian':
        mu=noise_para[0]
        sigma = noise_para[1]
        noise = np.zeros((len(spectrum)))

        for k in range(len(noise)):
           noise[k] = spectrum[k]*np.random.normal(mu,sigma)

    return noise

# get list of stellar spectra filenames from subfolder

dust_list = glob.glob("./dust_qext/*.q")
print(dust_list)
# print list of filenames
for i in range(len(dust_list)):
    print(i,dust_list[i])

# ask user for desired spectrum, distance

dust_index  = input("Please enter index(s) of desired dust component(s), separated by commas (no spaces e.g. '1,2,3'):")

dust_split = dust_index.split(',')

species=[]
species_string = ""
for i in range(len(dust_split)):
    species.append(dust_list[int(dust_split[i])])
    species_string += species[i][12:]+"_"
Nspecies = len(species)

print("Using ",Nspecies," dust species to generate features on reddened stellar spectrum.")

##################################
dust_fraction  = input("Please enter the fractions(s) of the desired dust component(s), separated by commas (no spaces e.g. '0.2,0.2,0.6', must add up to 1):")

dust_split_fraction = dust_fraction.split(',')
fraction=np.empty(len(dust_split_fraction))

for i in range(len(dust_split_fraction)):
    fraction[i]=float(dust_split_fraction[i])

# Here we define the extinction properties of the ISM
#Avs = [1.0,3.0,5.0,6.5,8.0,10.0,15.0] # in mags
#Avs = [4.0,5.5,7.0] # in mags
Avs = [7.0] # in mags


Rv = 3.1 # for MW

# create list of filenames based on material composition and extinction
filename_array = []
for i in range(len(Avs)):
    filename_array.append("extinction_"+species_string+"av"+str(Avs[i])+"_"+str(int(dstar/1000))+"kpc.txt")

#filename_array=["extinction_enstatite_av1_1kpc.txt","extinction_enstatite_av3_1kpc.txt","extinction_enstatite_av5_1kpc.txt","extinction_enstatite_av8_1kpc.txt","extinction_enstatite_av10_1kpc.txt","extinction_enstatite_av15_1kpc.txt"]

#Apply extinction models
for i in range(len(Avs)):
    Av = Avs[i]
    spectrum = apply(fitzpatrick99(wav_aa,Av,Rv),flx_mjy)

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlabel(r'Wavelength [$\mu$m]',fontsize=16)
    ax.set_ylabel(r'Flux density [Jy]',fontsize=16)
    ax.set_ylim(1e-5,1e3)
    ax.set_xlim(1e-2,1e3)
    ax.loglog(wav_um,spectrum,color='blue')
    ax.loglog(wav_um,flx_mjy,color='red')
    plt.show()
    plt.close()
    print(np.min(wav_um),np.max(wav_um))
    # wavelength for interpolation
    wavelength_int=np.arange(0.1,40.0,0.001)
    spectrum3_int = spectres(wavelength_int,wav_um,spectrum)
    #interpfunc_spectrum=interpolate.interp1d(wav_um,spectrum)
    #spectrum3_int=interpfunc_spectrum(wavelength_int)

    #interpfunc_model=interpolate.interp1d(wav_um,flx_mjy)
    #model_int=interpfunc_model(wavelength_int)

    model_int=spectres(wavelength_int,wav_um,flx_mjy)

    # optical depth array

    tau_cont=np.log(spectrum3_int/model_int)*-1.

    wavelength_int_q=np.arange(2,40.0,0.001)
    tau_int=spectres(wavelength_int_q,wavelength_int,tau_cont)
    model_int2=spectres(wavelength_int_q,wav_um,flx_mjy)

    #tau_cont=np.log(spectrum3_int/model_int)*-1.
    #wavelength_int_q=np.arange(2,40.0,0.001)
    #tau_int_func=interpolate.interp1d(wavelength_int,tau_cont)
    #tau_int=tau_int_func(wavelength_int_q)
    #model_int2=interpfunc_model(wavelength_int_q)

    tau_lab=np.empty([len(species),len(wavelength_int_q)])

    teller=0

    for specie in species:

        wavelengths_q, Qext = read_dust_qext(specie)

        rho, molmass = read_dust_prop(specie)
        rho=float(rho)
        molmass=float(molmass)
        # interpolating lab data
        interpfunc_Qext=interpolate.interp1d(wavelengths_q,Qext)
        Qext_int=interpfunc_Qext(wavelength_int_q)

        # How many dust particles do we have along the line of sight for a certain Av value?
        #Nd = calculate_Nd(Av)
        Nd=calculate_Nd(Av, rho, molmass, fraction, depletion_si)

        # Now we need the lab spectrum part of the intensity
        tau_lab[teller,:]=Qext_int*math.pi*(0.1*1e-4)**2*Nd*fraction[teller]

        teller=teller+1

    tau_lab=np.sum(tau_lab, axis=0)

    # Here we should add the carbon features
    # how did we integrate things ... , perhaps just expand 2-3 to -40 micron
    wavelength_carbon, tau_carbon=calculate_hydrocarbon_spectrum(Av,depletion_carbon)

    # water ice
    # here we add tau from water ice
    # THIS NEEDS THE OPTIONS TO CHOOSE DIFFERENT COMPOSITIONS!!!!!
    if water_ice_yesno==1:
        data_waterice = np.loadtxt(ice_species[0])
        wavelengths_q_waterice=data_waterice[:,0]
        Qext_waterice=data_waterice[:,1]
        # interpolating lab data
        interpfunc_Qext_waterice=interpolate.interp1d(wavelengths_q_waterice,Qext_waterice)
        Qext_int_waterice=interpfunc_Qext_waterice(wavelength_int_q)

        rho_ice, ice_grain_radius = read_ice_prop(ice_species[0])
        N_waterice=water_ice(Av, ice_grain_radius,rho_ice)
        tau_waterice=Qext_int_waterice*math.pi*(float(ice_grain_radius)*1.0e-4)**2*N_waterice
    elif water_ice_yesno==0:
        tau_waterice=0

    # small silicates
    wavelength_small_silicates, tau_small_silicates=small_silicates(Av,depletion_si,small_silicates_fraction)

    tau_smallsil = spectres(wavelength_int_q,wavelength_small_silicates,tau_small_silicates)
    tau_smallsil[np.isnan(tau_smallsil)] = 0

    # interpolating lab data

# now we need to bring back the original stellar spectrum
    tau=tau_int+tau_lab+tau_carbon+tau_waterice+tau_smallsil
    intensity=model_int2*np.exp(-1*(tau))
    wavelength_int_where=np.where(wavelength_int<2)
    wavelength_int_again=wavelength_int[wavelength_int_where]
    spectrum_again=spectrum3_int[wavelength_int_where]

    spectrum_total=np.concatenate((spectrum_again,intensity))
    wavelength_total=np.concatenate((wavelength_int_again,wavelength_int_q))

    spectrum_total_func=interpolate.interp1d(wavelength_total,spectrum_total,fill_value="extrapolate")
    spectrum_total_int = spectrum_total_func(wavelength_int)
    # Adding noise to the spectrum
    noise = spec_noise(spectrum_total_int,noise_para=None,noise_type='Poisson') #Noise = sqrt(Signal)
    #noise = spec_noise(spectrum_total_int,noise_para=[0.0,0.10],noise_type='Gaussian') #Noise = 10% Gaussian uncertainty

    spectrum_total_int_noisy = spectrum_total_int + noise

    # we now also need to consider binning and the
    # rebinning the spectrum
    #JWST MIRSPEC resolutions by filter:
    #https://jwst-docs.stsci.edu/mid-infrared-instrument/miri-observing-modes/miri-medium-resolution-spectroscopy#MIRIMediumResolutionSpectroscopy-wavelengthMRSwavelengthcoverage

    #lam_min, lam_max, avg. spec. resln.
    spectral_resolution = {'nirspec_prism' :[0.6,5.0,100],
                           'nirspec_lores' :[0.6,5.0,1000],
                           'nirspec_hires' :[0.6,5.0,2700],
                           'mirispec_ch1_s':[4.88,5.75,3515],
                           'mirispec_ch1_m':[5.63,6.63,3470],
                           'mirispec_ch1_l':[6.41,7.52,3355],
                           'mirispec_ch2_s':[7.48,8.76,3050],
                           'mirispec_ch2_m':[8.71 ,10.23,2960],
                           'mirispec_ch2_l':[10.02,11.75,3080],
                           'mirispec_ch3_s':[11.52,13.49,2705],
                           'mirispec_ch3_m':[13.36,15.65,2215],
                           'mirispec_ch3_l':[15.43,18.08,2385],
                           'mirispec_ch4_s':[17.65,20.94,1695],
                           'mirispec_ch4_m':[20.41,24.22,1725],
                           'mirispec_ch4_l':[23.88,28.34,1495]}

    #empty arrays for the final spectrum
    jwst_spec_wave = np.asarray([])
    jwst_spec_flux = np.asarray([])

    for key in spectral_resolution:
        channel = spectral_resolution[key]

        new_waves = np.arange(channel[0],channel[1],
                                (channel[1]-0.5*(channel[1]-channel[0]))/channel[2])
        new_waves = np.append(new_waves,channel[1])
        new_fluxs = spectres(new_waves,wavelength_int,spectrum_total_int_noisy)

        jwst_spec_wave = np.append(jwst_spec_wave,new_waves)
        jwst_spec_flux = np.append(jwst_spec_flux,new_fluxs)

    #plot figures

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlabel(r'Wavelength [$\mu$m]',fontsize=16)
    ax.set_ylabel(r'Flux density [mJy]',fontsize=16)
    ax.set_ylim(1e-1,1e3)
    ax.set_xlim(1,40)
    ax.loglog(wavelength_int,spectrum_total_int,color='blue')
    ax.loglog(wav_um,flx_mjy,color='red')
    ax.text(20,5e2,r'$A_{V}$ = '+str(Av),fontsize=16)

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlabel(r'Wavelength [$\mu$m]',fontsize=16)
    ax.set_ylabel(r'Relative intenstity (compared to photosphere)',fontsize=16)
    ax.set_ylim(1e-1,1.1e0)
    ax.set_xlim(1,40)
    ax.loglog(wavelength_int,spectrum_total_int_noisy/spectrum3_int,color='green')
    ax.loglog(wavelength_int,spectrum_total_int/spectrum3_int,color='blue')
    ax.loglog(wavelength_int,model_int/model_int,color='red')
    ax.text(20,0.7,r'$A_{V}$ = '+str(Av),fontsize=16)
    plt.show()
    plt.close()

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlabel(r'Wavelength [$\mu$m]',fontsize=16)
    ax.set_ylabel(r'Flux density [mJy]',fontsize=16)
    ax.set_ylim(1e-1,1e3)
    ax.set_xlim(1,40)
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.loglog(wavelength_int,spectrum3_int,color='black')
    cval = np.arange(0,1,1/len(spectral_resolution))
    nk = 0
    for keys in spectral_resolution:
        #cval = (nk+1)/len(spectral_resolution)
        lam_lo = spectral_resolution[keys][0]
        lam_hi = spectral_resolution[keys][1]
        filt = np.where((jwst_spec_wave >= lam_lo) & (jwst_spec_wave <= lam_hi))
        ax.scatter(jwst_spec_wave[filt],jwst_spec_flux[filt], cmap='rainbow',linestyle='-',marker='.',linewidth=0.5)
        nk += 1
    ax.text(20,5e2,r'$A_{V}$ = '+str(Av),fontsize=16)

    plt.show()
    plt.close()


    np.savetxt('./results/'+filename_array[i], np.c_[wavelength_int,spectrum_total_int/spectrum3_int], delimiter=' ')