# 
# Code to create a slider graph that shows the effect of varying f_sed and Kzz on spectra
# 
# 

import matplotlib
matplotlib.use('qt5Agg') 
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.signal import savgol_filter
from matplotlib.widgets import Button, Slider
from matplotlib.ticker import FormatStrFormatter

from pyodide.http import open_url

url = "https://www.star.bris.ac.uk/matt_lodge/pyscript_test/spectra" # add url for folder where data is stored on webserver

# ---- SET PARAMETERS -----

# list of d_f values
d_f_list=np.linspace(1.2,2.8,5)
width=2 # set linewidth


# ------------- MAKE SURE TO COPY AND PASTE THESE FROM THE "looper.py" SCRIPT -------------
# set array of Kzz values
Kzz_values = np.linspace(9.0, 10.0, 51)
# set array of fsed values
fsed_values = np.linspace(0.30, 1.00, 51)
# -----------------------------------------------------------------------------------------


# set default values of Kzz and fsed (just picking the middle value of each array)
initial_Kzz = Kzz_values[round(len(Kzz_values)/2)]
initial_fsed = fsed_values[round(len(fsed_values)/2)]

# choose a colormap (based on the number of shapes to analyse)
colors = plt.cm.cool(np.linspace(0,1,len(d_f_list)))

# calculate length of allowed value arrays
num_Kzz = len(Kzz_values)
num_fsed = len(fsed_values)

colnames=['wavelength', 'depth']

# ----- READ IN ALL SPECTRA DATASETS (do this at the beginning of the code (so that we don't have to read_csv every time the slider updates) -----

all_data=[] # top level list to hold all data for all fsed and Kzz

# loop over all fsed values
for i_fsed in range(num_fsed):

    # loop over all Kzz values
    for i_Kzz in range(num_Kzz):

        f_sed=fsed_values[i_fsed]
        Kzz=Kzz_values[i_Kzz]

        spectra_list=[] # initialise list to hold clear, spheres, and d_f spectra at this fsed and Kzz

        # load clear spectrum
        spectra_data= pd.read_csv(open_url(f"{url}/fsed_{f_sed:.4f}_Kzz_{Kzz:.4f}/clear.txt"), header=None, names=colnames, sep=",")
        spectra_list.append(savgol_filter(spectra_data['depth'], 25, 5)) # smooth the y-axis to make it look nicer for the plot and add it to the list

        # load sphere spectrum
        spectra_data= pd.read_csv(open_url(f"{url}/fsed_{f_sed:.4f}_Kzz_{Kzz:.4f}/spheres.txt"), header=None, names=colnames, sep=",") # load data
        spectra_list.append(savgol_filter(spectra_data['depth'], 25, 5)) # smooth the y-axis to make it look nicer for the plot and add it to the list

        # load spectra for d_f values
        for i in range(len(d_f_list)):
            spectra_data= pd.read_csv(open_url(f"{url}/fsed_{f_sed:.4f}_Kzz_{Kzz:.4f}/{d_f_list[i]:.1f}.txt"), header=None, names=colnames, sep=",")
            spectra_list.append(savgol_filter(spectra_data['depth'], 25, 5)) # smooth the y-axis to make it look nicer for the plot and add it to the list

        # add this list to ultimate list of all data
        all_data.append(spectra_list)

    print(f'Loading... {100.0*(i_fsed+1)/num_fsed:.1f} %')




# ----- PLOT DEFAULT SPECTRUM AT CHOSEN INITIAL VALUES OF Kzz and fsed -----

# find wavelength information for all plots by re-loading the clear spectrum for the default case
spectra_data= pd.read_csv(open_url(f"{url}/fsed_{f_sed:.4f}_Kzz_{Kzz:.4f}/clear.txt"), header=None, names=colnames, sep=",") # load data
wavelengths=spectra_data['wavelength']
    
fig,ax = plt.subplots()

# find the index of each of the initial values from our array
i_fsed=np.isclose(fsed_values, initial_fsed).nonzero()[0][0] # np.isclose finds the index of the closest value to "fsed" in the array "fsed_values". The .nonzero()[0][0] part is just because it returns an array, and we are only expecting one value in this array to match it perfectly (whereas this isn't always true for other programs)
i_Kzz=np.isclose(Kzz_values, initial_Kzz).nonzero()[0][0] # np.isclose finds the index of the closest value to "fsed" in the array "Kzz_values". The .nonzero()[0][0] part is just because it returns an array, and we are only expecting one value in this array to match it perfectly (whereas this isn't always true for other programs)

# translate this to the index of the "all_data" matrix to find the data for default values
index = i_fsed*num_Kzz + i_Kzz # (because the Kzz loop happens once per fsed value)

# plot clear spectrum
clear_line,=ax.plot(wavelengths.values, all_data[index][0], color = 'w', label= r'No condensate', linewidth=width) # the indices of "all_data" are [index from order of fsed/kzz loop][particle shape: clear=0, spheres=1, d_f=2-6]

# plot spheres
sphere_line,=ax.plot(wavelengths.values, all_data[index][1], color = 'yellow', label= r'spheres', linewidth=width)

# plot fractals
fractal_line_0,=ax.plot(wavelengths.values, all_data[index][2], color = colors[0], label= f'{d_f_list[0]:.1f}', linewidth=width) # plot d_f value 0
fractal_line_1,=ax.plot(wavelengths.values, all_data[index][3], color = colors[1], label= f'{d_f_list[1]:.1f}', linewidth=width) # plot d_f value 1
fractal_line_2,=ax.plot(wavelengths.values, all_data[index][4], color = colors[2], label= f'{d_f_list[2]:.1f}', linewidth=width) # plot d_f value 2
fractal_line_3,=ax.plot(wavelengths.values, all_data[index][5], color = colors[3], label= f'{d_f_list[3]:.1f}', linewidth=width) # plot d_f value 3
fractal_line_4,=ax.plot(wavelengths.values, all_data[index][6], color = colors[4], label= f'{d_f_list[4]:.1f}', linewidth=width) # plot d_f value 4

# set axes titles
ax.set_xlabel(r'$\lambda (\mu m)$', fontsize=23)
ax.set_ylabel(r'Transit depth (%)', fontsize=23)

ax.set_xscale('log') # log the wavelength axis

# set tick font size
ax.tick_params(axis='both', which='major', labelsize=20)
ax.tick_params(axis='both', which='minor', labelsize=20)
    
handles, labels = ax.get_legend_handles_labels()
ax.legend(reversed(handles), reversed(labels), title='$d_f$', title_fontsize=12, fontsize=12,loc='lower right', bbox_to_anchor=(0.7, 0))

# adjust the main plot to make room for the sliders
fig.subplots_adjust(bottom=0.3)



# ----- MAKE SLIDERS -----

# Make a horizontal slider to control Kzz
ax_Kzz = fig.add_axes([0.25, 0.15, 0.65, 0.03]) # axes position as: x0, y0, width, height
slider_Kzz = Slider(
    ax_Kzz, "$\log_{10}(K_{zz})$", min(Kzz_values),  max(Kzz_values), # title, start and end range of slider
    valinit=initial_Kzz, # initial default value
    valstep=Kzz_values, # allowed snapped values
    initcolor='none',  # remove the line marking the initial/default value position
    color="limegreen",
)

# Add ticks to slider
ax_Kzz.add_artist(ax_Kzz.xaxis)
Kzz_ticks = np.linspace(min(Kzz_values), max(Kzz_values),5) # set range and number of ticks
ax_Kzz.set_xticks(Kzz_ticks)
ax_Kzz.xaxis.set_major_formatter(FormatStrFormatter('%.1f')) # round ticks to nicer numbers




# Make a horizontal slider to control fsed
ax_fsed = fig.add_axes([0.25, 0.05, 0.65, 0.03]) # axes position as: x0, y0, width, height
slider_fsed = Slider(
    ax_fsed, "$f_{sed}$",  min(fsed_values),  max(fsed_values), # title, start and end range of slider
    valinit=initial_fsed, # initial default value
    valstep=fsed_values, # allowed snapped values
    initcolor='none',  # remove the line marking the initial/default value position
    color="tab:red"
)

# Add ticks to slider
ax_fsed.add_artist(ax_fsed.xaxis)
fsed_ticks = np.linspace(min(fsed_values), max(fsed_values),5) # set range and number of ticks
ax_fsed.set_xticks(fsed_ticks)
ax_fsed.xaxis.set_major_formatter(FormatStrFormatter('%.2f')) # round ticks to nicer numbers

'''# to find all attributes of the sliders, use this code to print all of the slider's properties!
from pprint import pprint
pprint(vars(slider_Kzz))'''



# ----- MAKE FUNCTION TO UPDATE FIGURE WHEN SLIDERS ARE CHANGED -----

# The function to be called anytime a slider's value changes
def update(val):
    fsed = slider_fsed.val
    Kzz = slider_Kzz.val

    # find the index of each of these values from our array
    i_fsed=np.isclose(fsed_values, fsed).nonzero()[0][0] # np.isclose finds the index of the closest value to "fsed" in the array "fsed_values". The .nonzero()[0][0] part is just because it returns an array, and we are only expecting one value in this array to match it perfectly (whereas this isn't always true for other programs)
    i_Kzz=np.isclose(Kzz_values, Kzz).nonzero()[0][0] # np.isclose finds the index of the closest value to "fsed" in the array "Kzz_values". The .nonzero()[0][0] part is just because it returns an array, and we are only expecting one value in this array to match it perfectly (whereas this isn't always true for other programs)

    # translate this to the index of the all_data matrix
    index = i_fsed*num_Kzz + i_Kzz # (because the Kzz loop happens once per fsed value

    #print(f'UPDATE: f_sed index {i_fsed} and Kzz index {i_Kzz} means using global index {index}')

    # update each line seperately (we need to this individually for them to keep the same colours and labels etc)
    clear_line.set_ydata(all_data[index][0])
    sphere_line.set_ydata(all_data[index][1])
    fractal_line_0.set_ydata(all_data[index][2])
    fractal_line_1.set_ydata(all_data[index][3]) # the indices of "all_data" are [index from order of fsed/kzz loop][particle shape: clear=0, spheres=1, d_f=2-6]
    fractal_line_2.set_ydata(all_data[index][4])
    fractal_line_3.set_ydata(all_data[index][5])
    fractal_line_4.set_ydata(all_data[index][6])

    fig.canvas.draw_idle()

# register the update function with each slider
slider_Kzz.on_changed(update)
slider_fsed.on_changed(update)

plt.show()