from IPython.core.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))

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

cmap = cm.get_cmap('jet', 20)
colors = []
for i in range(cmap.N):
    rgb = cmap(i)[:3]
    colors.append(matplotlib.colors.rgb2hex(rgb))

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))

import plotly.graph_objects as go
fig = go.FigureWidget().set_subplots(
    rows = 3,
    cols = 2,
    specs = [
        [{}, {}],
        [{"colspan": 2}, None],
        [{"colspan": 2}, None]
    ],
    row_heights=[0.55, 0.4,0.1],
    vertical_spacing = 0.07
)

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))


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 = 2,
        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 = 3, 
    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.add_scatter(row = 1, col = 1)
fig.add_scatter(row = 1, col = 2)

fig.update_layout(
    height = 1000,
    yaxis = {
        "range": [9000, 14000],
        "tickformat": "g",
        "title": r"$T_{rad}\ [K]$"
    },
    yaxis2 = {
        "title": r"$W$"
    },
    yaxis3 = {
        "title": r"$\mbox{Luminosity}~(erg~sec^{-1})$",
        "exponentformat": "e",
    },
    yaxis4 = {
        "title": r"$~~\mbox{Residual}\\\mbox{Luminosity(%)}$", 
        "exponentformat": "e",
        "title_font_size": 13,
    },
    
    xaxis = {
        "title": r"$Shell~~Velocity$",
        "tickformat": "g"
    },
    xaxis2 = {
        "title": r"$Shell~~Velocity$",
        "tickformat": "g",
        "matches": "x"
    },
    xaxis3 = {
        "dtick": 2,
        "range": [0, 21],
    },
    xaxis4 = {
        "title": r"$\mbox{Iteration Number}$", 
        "matches": "x3",
        "dtick": 2,
        "range": [0, 21],
    },
    
    
    
)