import matplotlib.pyplot as plt
import h5py
import json
import matplotlib.gridspec as gridspec


NoSpectra = json.load(open("Setting.json"))[0]["Spectra"];
for spectr in ["Halpha_0.h5", "IR_0.h5", "VIS_0.h5", "UV_0.h5", "IRVISUV_0.h5"]:

    Spectra = h5py.File(spectr, "r")
    ShotNo=open("ShotNo","r").read()
    
    fig = plt.figure()
    gs = fig.add_gridspec(NoSpectra, hspace=0)
    axs = gs.subplots(sharex=True, sharey=True)
    fig.set_figheight(10);fig.set_figwidth(20)
    plt.suptitle("Golem spectum #" + str(ShotNo) +str(spectr[0:-5]))
    
    for i in range(0, NoSpectra):
          axs[i].plot(Spectra['Wavelengths'][0:], Spectra['Spectra'][i:][0], label=i)
    plt.savefig('ScreenShotAll.png')

NoSpectra = json.load(open("Setting.json"))[0]["Spectra"]
ShotNo = open("ShotNo", "r").read()

spectr_list = ["Halpha_0.h5", "IR_0.h5", "VIS_0.h5", "UV_0.h5", "IRVISUV_0.h5"]
num_files = len(spectr_list)

fig, axs = plt.subplots(NoSpectra, num_files, sharex=True, figsize=(20, 10))
fig.subplots_adjust(hspace=0, wspace=0)

for col_idx, spectr in enumerate(spectr_list):
    Spectra = h5py.File(spectr, "r")
    for row_idx in range(NoSpectra):
        if NoSpectra > 1:
            ax = axs[row_idx, col_idx]
        else:
            ax = axs[col_idx]
        ax.plot(Spectra['Wavelengths'][:], Spectra['Spectra'][row_idx], label=f"Spectrum {row_idx}")
        
        if row_idx != NoSpectra - 1:
            ax.tick_params(labelbottom=False)
        else:
            ax.set_xlabel('Wavelength')
        
        if col_idx != 0:
            ax.tick_params(labelleft=False)

        
        if row_idx == 0:
            ax.set_title(spectr[:-5])

fig.suptitle(f"Golem Spectrum #{ShotNo}", fontsize=16)

plt.savefig('ScreenShotAll_horizontal_stack.png')
plt.show()

NoSpectra = json.load(open("Setting.json"))[0]["Spectra"]
spectr_list = ["Halpha_0.h5", "IR_0.h5", "VIS_0.h5", "UV_0.h5", "IRVISUV_0.h5"]
number_of_files = len(spectr_list)
ShotNo = open("ShotNo", "r").read()

rows_per_file = NoSpectra + 1
total_rows = number_of_files * rows_per_file

fig = plt.figure(figsize=(20, total_rows * 0.75))
gs = gridspec.GridSpec(total_rows, 1, hspace=0)

current_row = 0 
for spectr in spectr_list:
    ax_title = fig.add_subplot(gs[current_row, 0])
    ax_title.text(0.5, 0.5, f"{spectr[:-5]}", ha='center', va='center', fontsize=16)
    ax_title.axis('off')  # Hide the axis
    current_row += 1  # Move to the next row

    with h5py.File(spectr, "r") as Spectra:
        wavelengths = Spectra['Wavelengths'][:]
        spectra_data = Spectra['Spectra']

        for i in range(NoSpectra):
            # Create a subplot for each spectrum
            ax = fig.add_subplot(gs[current_row, 0], sharex=ax if current_row > 1 else None)
            ax.plot(wavelengths, spectra_data[i], label=f"Spectrum {i}")


            current_row += 1  # Move to the next row

fig.suptitle(f"Golem Spectrum #{ShotNo}", fontsize=20)

plt.tight_layout(rect=[0, 0, 1, 0.98])  # Leave space for the super title
plt.savefig('ScreenShotAll_vertical_stack.png', bbox_inches='tight')
plt.show()

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import h5py
import json

# load basic data
NoSpectra = json.load(open("Setting.json"))[0]["Spectra"]
ShotNo = open("ShotNo", "r").read()

# store names of spectroscopes
spectr_list = ["Halpha_0.h5", "IR_0.h5", "VIS_0.h5", "UV_0.h5", "IRVISUV_0.h5"]

# define colors for each spectra file
color_list = ['red', 'green', 'blue', 'orange', 'purple']
spectr_colors = dict(zip(spectr_list, color_list))

# subplot with shared axes
fig = make_subplots(rows=NoSpectra, cols=1, shared_xaxes=False, vertical_spacing=0.02)

# initialize annotations list
annotations = []

# loop over each spectrum index (corresponding to different times)
for ti in range(NoSpectra):
    # loop over each spectroscopes
    for spectr in spectr_list:
        with h5py.File(spectr, "r") as Spectra:
            wavelengths = Spectra['Wavelengths'][:]
            spectrum_data = Spectra['Spectra'][ti]
            # create a trace with labels
            fig.add_trace(
                go.Scatter(
                    x=wavelengths,
                    y=spectrum_data,
                    mode='lines',
                    name=f'{spectr[:-5]}',
                    showlegend=False,  # Disable the global legend
                    line=dict(color=spectr_colors[spectr]),
                    hovertemplate=f'{spectr[:-5]}<br>Wavelength: %{{x}}<br>Intensity: %{{y}}<extra></extra>'
                ),
                row=ti + 1,
                col=1,
            )
    # update y-axis label for each subplot
    fig.update_yaxes(title_text=f'Spectrum {ti}', row=ti + 1, col=1)

    # determine yref for annotations
    yref = 'y' + (str(ti+1) if ti > 0 else '') + ' domain'

    # ddd legend-like annotations to each subplot (showing both legend and datapoint descriptions)
    for idx, spectr in enumerate(spectr_list):
        annotations.append(
            dict(
                x=1.01,  # Slightly outside the plot area to the right
                y=1 - (idx * 0.08),  # Adjust y position for each legend entry
                xref='x domain',
                yref=yref,
                xanchor='left',
                yanchor='middle',
                text=f'<span style="color:{spectr_colors[spectr]};">{spectr[:-5]}</span>',
                showarrow=False,
                font=dict(size=10),
            )
        )

# update the figure layout to include annotations from above
fig.update_layout(
    annotations=annotations,
    height=300 * NoSpectra,
    width=1000,
    title_text=f"Golem Spectrum #{ShotNo}",
    margin=dict(r=150),  # Increase right margin to make space for annotations
)

# Update x-axis label
fig.update_xaxes(showticklabels=True, title_text='Wavelength', row=NoSpectra, col=1)

# Save to HTML
fig.write_html('interactive_plot.html', auto_open=False)