import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from ftio.freq.helper import MyConsole
from ftio.plot.print_html import PrintHtml
from ftio.plot.helper import format_plot_and_ticks
from ftio.api.metric_proxy.phasemode import PhaseMode


CONSOLE = MyConsole()
CONSOLE.set(True)


def phases_and_timeseries(metrics, data, argv=[]):
    phasemode_list, t = classify_waves(data, True)

    if argv and '-n' in argv:
        n = argv[argv.index('-n') + 1]
        plot_waves_and_timeseries(argv,metrics, phasemode_list, t, n)
    else:
        plot_waves_and_timeseries(argv,metrics, phasemode_list, t)


def phases(data, argv=[]):
    phasemode_list, t = classify_waves(data)

    if argv and '-n' in argv:
        n = argv[argv.index('-n') + 1]
        plot_waves(argv,phasemode_list, t, n)
    else:
        plot_waves(argv,phasemode_list, t)


def classify_waves(data, normed=True):
    phasemode_list = []
    phasemode_list.append(
        PhaseMode(
            'Network',
            [
                'network',
                'mpi_bcast',
                'sock',
                'send',
                'listen',
                'recieve',
                'gather',
                'reduce',
                'cast',
            ],
        )
    )
    phasemode_list.append(PhaseMode('write', ['write']))
    phasemode_list.append(PhaseMode('read', ['read']))
    phasemode_list.append(
        PhaseMode(
            'I/O',
            [
                'io',
                'disk',
                'stat',
                'cntl',
                'read',
                'write',
                'openat',
                'mmap',
                'seek',
                'close',
            ],
        )
    )
    phasemode_list.append(PhaseMode('Temperature', ['temp']))
    phasemode_list.append(PhaseMode('CPU', ['cpu']))
    phasemode_list.append(PhaseMode('Wait', ['poll', 'barrier']))
    seen = []
    for mode in phasemode_list:
        seen.extend(mode.matches)
    other = [x['metric'] for x in data if all(y not in x['metric'] for y in seen)] 
    phasemode_list.append(PhaseMode('Other', other))

    sampling_freq = np.NaN
    t_s = np.inf
    t_e = 0
    for d in data:
        if len(d['dominant_freq']) > 0 and len(d['conf']) > 0:
            if np.isnan(sampling_freq):
                sampling_freq = d['freq']
                t_s = min(d['t_start'],t_s)
                t_e = max(d['t_end'],t_e)

            add_metric(phasemode_list, d)

    if not np.isnan(sampling_freq):
        t = np.arange(t_s, t_e, 1 / sampling_freq)
        add_time(phasemode_list, t)
    else:
        CONSOLE.print('[red] No metrics classified. Try increasing the sampling frequency[/]')    
        exit(0)

    text = print_len(phasemode_list, data)
    CONSOLE.print(f'[blue]{text}[/]')
    aggregate_wave_for_all_modes(phasemode_list, t, normed)

    return phasemode_list, t


def plot_waves(argv:list, arr: list[PhaseMode], t, n=None):
    plot_mode = False
    if argv and '-e' in argv:
        plot_mode = "mat" in argv[argv.index('-e') + 1]

    if plot_mode:
        # Create a figure and axis using variables
        fig, ax = plt.subplots()

        # Loop through each mode in arr and plot
        for mode in arr:
            ax.plot(t, mode.wave, marker='o', label=mode.name)  # Plot with markers


        # Set axis titles
        ax.set_xlabel('Time (s)')
        ax.set_ylabel('Normed metrics')

        # Set the legend title if n is provided
        if n:
            ax.legend(title=f'{n} Frequencies')
        else:
            ax.legend()

        plt.show()
    else:
        fig = go.Figure()
        for mode in arr:
            fig.add_trace(
                go.Scatter(x=t, y=mode.wave, mode='lines+markers', name=mode.name)
            )

        fig.update_layout(
            xaxis_title='Time (s)',
            yaxis_title='Normed metrics',
        )
        if n:
            fig.update_layout(legend_title_text=f'{n} Frequencies')
        fig = format_plot_and_ticks(fig)
        fig.show()
        # fig.write_image("waves.png")


def plot_waves_and_timeseries(argv:list,  metrics: dict, arr: list[PhaseMode], t, n=None):
    names = [] #get_names(arr)

    out = PrintHtml('./', names, outdir='phase_plots')
    out.generate_html_start()
    # t = []
    for mode in arr:
        f = plot_mode(mode,metrics,t,n,True)
        # for fig in f:
        #     fig.show
        out.generate_html_core(mode.name+'.html', f)

    out.generate_html_end()


def plot_mode(mode,metrics,t,n,subfig=False)-> list[go.Figure]:

    f = []
    spec=[{},{}]
    if subfig:
        f.append(make_subplots(rows=2,cols=1))
        spec=[{"row":1,"col":1},{"row":2,"col":1}]
    else:
        f.append(go.Figure())

    f[-1].add_trace(
        go.Scatter(x=t, y=mode.wave, mode='lines+markers', name=mode.name),**spec[0])

    if not subfig:
        f.append(go.Figure())

    for metric, arrays in metrics.items():
        if len(arrays[0]) > 1 and any(n in metric for n in mode.matches):
            f[-1].add_trace(
                go.Scatter(
                    x=arrays[1],
                    y=arrays[0],
                    mode='lines+markers',
                    name=metric,
                    hovertemplate='<i>Time </i>: %{x} s'
                    + '<br><b>Metric</b>: %{y}<br>',
                    legendgroup=metric,
                ),**spec[1],
            )
            wave,name = mode.get_wave(metric)
            f[-1].add_trace(
                go.Scatter(
                    x=mode.get("t"),
                    y=wave,
                    mode='lines+markers',
                    name=name,
                    hovertemplate='<i>Time </i>: %{x} s'
                    + '<br><b>Metric</b>: %{y}<br>',
                    legendgroup=metric,
                ),**spec[0],
            )

    for fig in f:
        fig = format_plot_and_ticks(fig, font=False)
        fig.update_layout(
            xaxis_title='Time (s)',
            yaxis_title=f'{mode.name.capitalize()} Metrics',
            width=1400,
            height=800,
            showlegend=True,
        ) 
    if subfig:
        fig.update_xaxes(title_text='Time (s)', **spec[1])
        fig.update_yaxes(title_text=f'{mode.name.capitalize()} Metrics',**spec[1])

        if f.index(fig) == 0:
            if n:
                fig.update_layout(title=f'{mode.name.capitalize()}: Using {n} Frequencies per metric')
            else:
                fig.update_layout(title=f'{mode.name.capitalize()}: Using Dominant Frequencies')
    return f




def add_metric(arr: list[PhaseMode], d):
    for mode in arr:
        if mode.match(d):
            mode.add(d)
            break

def add_time(arr: list[PhaseMode], t):
    for mode in arr:
        mode.set_time(t)


def get_names(arr:list[PhaseMode])-> list[str]:
    names = []
    for mode in arr:
        names.append(mode.name)
    return names 


def aggregate_wave_for_all_modes(arr: list[PhaseMode], t, normed=True):
    for mode in arr:
        mode.aggregates_waves(t, normed)


def print_len(arr: list[PhaseMode], data) -> str:
    text = '\n'
    total_metrics = 0
    for mode in arr:
        text += f'{len(mode.data):3} {mode.name} metrics\n'
        total_metrics += len(mode.data)

    text += f'---------------------\n{total_metrics:3}/{len(data):3} total metrics with dominant frequencies\n'

    return text
