from tardis import run_tardis
import tardis
import plotly.graph_objects as go
from collections import defaultdict

marker_colors = ["#958aff","#ff8b85","#5cff74"]
marker_line_colors = ['#27006b', '#800000', '#00801c']
marker_colors = ['#636EFA', '#EF553B', '#00CC96']

luminosities = ["Emitted", "Absorbed", "Requested"]
X = list(range(1,21))


fig = go.FigureWidget().set_subplots(
        2,1, 
        shared_xaxes=True, 
        vertical_spacing = 0.1,
        row_heights=[0.75, 0.25]
    )

for luminosity, marker_color, marker_line_color in zip(luminosities, marker_colors, marker_line_colors):
    fig.add_scatter(
        name = luminosity + "<br>Luminosity", 
        mode='lines+markers', 
        row = 1, 
        col = 1, 
        marker_color=marker_color, 
        marker_line_color=marker_line_color,
        legendgroup = luminosity, 
        marker_line_width=1.5, 
        opacity=0.6
    )
fig.add_scatter(
    name = "Residual<br>Luminosity",
    row = 2, 
    col = 1, 
    hovertext="text",
    marker_color='rgb(158,202,225)', 
    marker_line_color='rgb(8,48,107)',
    marker_line_width=1.5, 
    mode='lines+markers',
    opacity=0.7

)


fig = fig.update_layout(
    xaxis = dict(
        range = [0, 21],
        dtick = 2,
        
    ),
    xaxis2 = dict(
        matches = "x",
        title = r"$\mbox{Iteration Number}$", 
        dtick = 2,
    ),
    
    yaxis = dict(
        title = r"$\mbox{Luminosity}~(erg~sec^{-1})$",
        automargin = True,
        side = "top",
        exponentformat = "e",
    ),
    yaxis2 = dict(
        exponentformat = "e",
        title = r"$~~\mbox{Residual}\\\mbox{Luminosity(%)}$", 
        title_font_size = 13,
        automargin = True
    ),
    legend_tracegroupgap = 0,
    hidesources = True,
    height = 600,
    hoverlabel_align = 'right',
    legend_title_text = "Luminosity",
)

   

fig

traces = defaultdict(list)

def update_traces(sim):

    global traces

    emitted_luminosity = sim.runner.calculate_emitted_luminosity(
        sim.luminosity_nu_start, sim.luminosity_nu_end
    )
    reabsorbed_luminosity = sim.runner.calculate_reabsorbed_luminosity(
        sim.luminosity_nu_start, sim.luminosity_nu_end
    )

    traces["Emitted"].append(emitted_luminosity.value)
    traces["Absorbed"].append(reabsorbed_luminosity.value)
    traces["Requested"].append(sim.luminosity_requested.value)
    

    with fig.batch_update():
        for index, luminosity in zip(range(3), luminosities):
            fig.data[index].x = X
            fig.data[index].y = traces[luminosity]
            fig.data[index].hovertemplate = '<b>%{y:.2g}</b>'+'<br>at X = %{x}<br>'
    
        y = [((emitted-requested)*100)/requested for emitted, requested in zip(traces["Emitted"],traces["Requested"])]
        fig.data[-1].x = X
        fig.data[-1].y = y
        fig.data[-1].hovertemplate='Residual Luminosity: %{y:.2f}% at X = %{x:,.0f}<extra></extra>'