from matplotlib.axes import Axes
from matplotlib.figure import Figure, SubFigure
import numpy as np

from shnitsel.analyze.populations import PopulationStatistics
from shnitsel.filtering.state_selection import StateSelection

from .common import figaxs_defaults, centertext
import xarray as xr
from ....units.conversion import convert_time
from ....units.definitions import time


def plot_time_interstate_error(
    data: xr.Dataset, state_selection: StateSelection, ax: Axes
) -> Axes:
    """Function to plot the error bars/confidence intervals on time series interstate plots.

    Uses the `upper`, `lower` and `mean` variables of `data`.

    Parameters
    ----------
    data : Dataset
            Interstate Dataset containing the `upper`, `lower` and `mean` variables generated by a confidence interval calculation.
        Data must have the appropriate unit as conversion is not possible without knowing the target dimension or unit.
        You can set the units in the `unit` attribute of the data.
        Similarly, you can specify the y-axis label with the `tex` attribute of the data.
        The time will, however, be converted to fs if not yet converted.
    state_selection : StateSelection
            State selection object to limit the states included in plotting and to provide state names.
    ax : Axes
            Axes object to plot the graph to.

    Returns
    -------
    Axes
        The axes that have been plotted to
    """
    # TODO: FIXME: Support more states.
    # vas = {
    #     '$S_2 - S_0$': 'bottom',
    #     '$S_2 - S_1$': 'bottom',
    #     '$S_1 - S_0$': 'top',
    # }
    label_modes = ["top", "bottom"]
    for i, (sc, scdata) in enumerate(data.groupby('statecomb')):
        if not state_selection.has_state_combination(sc):
            continue

        data_filtered_na = scdata.dropna(dim='time', how='all')
        if data_filtered_na.time.size == 0:
            # Skip combinations without data
            continue

        # c = scdata['_color'].item()
        color = state_selection.get_state_combination_color(sc)
        scdata = scdata.squeeze('statecomb')

        # state_name_from = str(data['state_names'][sc[0] - 1].item())
        # state_name_to = str(data['state_names'][sc[1] - 1].item())

        time_in_fs = convert_time(scdata['time'], time.femto_seconds)

        # ax.fill_between(
        #     time_in_fs, scdata['upper'], scdata['lower'], color=color, alpha=0.3
        # )
        ax.plot(time_in_fs, scdata['mean'], c=color, lw=0.8)
        # va = vas.get(sc, 'baseline')
        # va = 'baseline'

        label_va = label_modes[i % len(label_modes)]

        label_x: float
        label_y: float
        if label_va == "bottom":
            try:
                max_index = scdata['upper'].dropna(dim='time', how='all').argmax()
            except ValueError:
                max_index = scdata['mean'].dropna(dim='time', how='all').argmax()
            label_x = float(scdata['time'][max_index])
            label_y = float(scdata['mean'][max_index])
        elif label_va == "top":
            try:
                min_index = scdata['lower'].dropna(dim='time', how='all').argmin()
            except ValueError:
                min_index = scdata['mean'].dropna(dim='time', how='all').argmin()
            label_x = float(scdata['time'][min_index])
            label_y = float(scdata['mean'][min_index])
        else:
            label_x = float(scdata['time'][-1])
            label_y = float(scdata['mean'][-1])

        label_y = max(label_y, 0.0)
        combination_label = state_selection.get_state_combination_tex_label(sc)
        ax.text(
            label_x,
            label_y,
            f"${combination_label}$",  # f"${state_name_from} \\to {state_name_to}$",
            c=color,
            va=label_va,
            ha='right',
        )

    ylabel = data.attrs['tex']
    if u := data.attrs.get('units'):
        ylabel += f" / {u}"
    ax.set_ylabel(ylabel)

    return ax


def plot_populations_graph(
    pops: PopulationStatistics, state_selection: StateSelection, ax: Axes
) -> Axes:
    """Plot a graph of populations to the provided axes.

    Parameters
    ----------
    pops : PopulationStatistics
            The per-state population data. We expect the `time` data to be in fs. If it can be converted, it will be converted to fs.
    state_selection : StateSelection
            State selection object to limit the states included in plotting and to provide state names.
    ax : Axes
            The Axes object to plot the population graph into

    Returns
    -------
    Axes
        The resulting Axes object after the plot
    """
    num_states_with_pops = 0
    for state, sdata in pops.relative.groupby('state'):
        if not state_selection.has_state(state):
            continue

        if sdata.max() < 1e-9:
            # No population data, skip
            continue
        num_states_with_pops += 1

        color = state_selection.get_state_color(state)
        # c = sdata['_color'].item()
        state_label = state_selection.get_state_tex_label(
            state
        )  # str(sdata['state_names'].item())

        time_in_fs = convert_time(sdata['time'], time.femto_seconds)

        ax.plot(time_in_fs, sdata, c=color, lw=0.5)

        max_index = sdata.argmax()

        # ax.text(float(time_in_fs[-1]), float(sdata[-1]), f"${state_label}$", c=c)
        ax.text(
            float(time_in_fs[max_index]),
            float(sdata[max_index]) + 0.5e-1,
            f"${state_label}$",
            c=color,
        )
    if num_states_with_pops == 0:
        centertext(r"No pop data", ax)
        # axs['ntd'].tick_params(axis='y', labelleft=False)
        ax.get_yaxis().set_visible(False)
        ax.get_xaxis().set_visible(False)

    ax.set_ylabel('Population')
    return ax


@figaxs_defaults(mosaic=[['pop'], ['de'], ['ft']], scale_factors=(1 / 3, 1 / 2))
def plot_timeplots(
    pops: PopulationStatistics,
    delta_E: xr.Dataset,
    fig: Figure | SubFigure | None,
    state_selection: StateSelection,
    fosc_time: xr.Dataset | None = None,
    axs: dict[str, Axes] | None = None,
) -> dict[str, Axes]:
    """Function to generate all time plots, involving population plots and state-transition plots.

    Parameters
    ----------
    pops : PopulationStatistics
            object containing state population data with a time coordinate
    delta_E : xr.Dataset
            Dataset containing the energy delta per state combination.
    fig : Figure| SubFigure
            A figure, consumed by the autmatic axes generation. Not used by the function itself.
    state_selection : StateSelection
            State selection object to limit the states included in plotting and to provide state names.
    fosc_time : xr.DataArray, optional
            The oscillation frequency in a DataArray groupable by `time`
    axs : dict[str, Axes], optional
            The dictionary of named Axes objects to plot the inidividual graphs into. Defaults to None.

    Returns
    -------
    dict[str, Axes]
        The axes dictionary after plotting.
    """
    assert axs is not None, "Could not generate axes for time plots."
    plot_populations_graph(
        pops,
        state_selection,
        axs['pop'],
    )
    non_degenerate_selection = state_selection.non_degenerate()
    plot_time_interstate_error(delta_E, non_degenerate_selection, axs['de'])
    axs['de'].set_ylim(0.0, None)
    if fosc_time is not None:
        plot_time_interstate_error(fosc_time, non_degenerate_selection, axs['ft'])
        axs['ft'].set_ylim(0.0, None)
        lowest_ax = axs['ft']
        higher_axnames = ['de', 'pop']
    else:
        centertext(r"No $\mathbf{\mu}_{ij}$ data", ax=axs['ft'])
        axs['ft'].get_yaxis().set_visible(False)
        axs['ft'].get_xaxis().set_visible(False)
        lowest_ax = axs['de']
        higher_axnames = ['pop']

    lowest_ax.set_xlabel(r'$t$ / fs')  # TODO: Make configurable or ensure conversion.
    lowest_ax.minorticks_on()

    for axn in higher_axnames:
        axs[axn].sharex(lowest_ax)
        axs[axn].tick_params(axis='x', labelbottom=False)

    return axs
