"""
Build trace plots (plotly)
"""

from typing import Union

__all__ = ['plot_trace']

import numpy as np
from .colors import COLORS

# TODO: when minimum python version is 3.10, can use | to combine types 
def plot_trace(x: Union[list, np.ndarray], ys: Union[list, np.ndarray], label='', alpha=0.4):

    import plotly.graph_objs as go

    traces = []
    color_idx = 0
    for y in ys:
        color = COLORS[color_idx % len(COLORS)]
        traces.append(go.Scatter(x=x,
                                y=y,
                                mode='lines',
                                showlegend=False,
                                line=dict(color=color),
                                opacity=alpha))
        color_idx += 1

    fig = go.Figure(data=traces)

    fig.update_layout(
        template = 'plotly_white',
        xaxis_title=dict(text='Generation number'),
        yaxis_title=dict(text=label),
        yaxis=dict(exponentformat='e'),
        showlegend=False
    )

    return fig