# Below taken from the GUI profile plot methods in Refl1d - modified to be stand alone
# This is to allow for the layer labels to be easily added to the SLD plots

# Originally written for use in a jupyter notebook
# Adapted to act as a prototype for the profile plotting in the new GUI for refl1d

from typing import List, TypedDict

import numpy as np
from numpy import inf

from refl1d.experiment import Experiment

from .colors import COLORS as MULTI_PLOT_COLORS

SINGLE_PLOT_COLORS = ("black", "blue", "green", "gold")

SINGLE_PLOT_STYLES = ["solid"] * 4
MULTI_PLOT_STYLES = ["solid", "dash", "dashdot", "dot"]

# === Sample information ===


class FindLayers:
    def __init__(self, experiment, axes=None, x_offset=0):
        self.experiment = experiment
        self.axes = axes
        self.x_offset = x_offset
        self._show_labels = True
        self._show_boundaries = True

        self._find_layer_boundaries()

    def find(self, z):
        return self.experiment.sample.find(z - self.x_offset)

    def _find_layer_boundaries(self):
        offset = self.x_offset
        boundary = [-inf, offset]
        if hasattr(self.experiment, "sample"):
            for L in self.experiment.sample[1:-1]:
                dx = L.thickness.value
                offset += dx
                boundary.append(offset)
        boundary.append(inf)
        self.boundary = np.asarray(boundary)

    def sample_labels(self):
        return [L.name for L in self.experiment.sample]

    def sample_layer(self, n):
        return self.experiment.sample[n]

    def label_offsets(self):
        z = self.boundary[1:-1]
        left = -20
        middle = [(a + b) / 2.0 for a, b in zip(z[:-1], z[1:])]
        right = z[-1] + 20
        return [left] + middle + [right]

    def reset_markers_plotly(self):
        """
        Reset all markers, for plotly plots
        """
        import plotly.graph_objs as go

        # self.clear_markers()
        if not isinstance(self.axes, go.Figure):
            raise ValueError("reset_markers_plotly can only be used with type: axes=plotly.graph_objs.Figure")

        fig = self.axes

        # Add bars
        fittable = [self.experiment.sample[idx].thickness.fittable for idx, _ in enumerate(self.boundary[1:-1])]
        fittable[0] = False  # First interface is not fittable

        for f, z_pos in zip(fittable, self.boundary[1:-1]):
            if not f:
                line_dict = dict(dash="dash")
            else:
                line_dict = dict(dash="solid")

            fig.add_vline(
                x=z_pos,
                opacity=0.75,
                line=line_dict,
            )

        # for f,m in zip(fittable,self.markers):
        #     if not f: m.set(linestyle='--', linewidth=1.25)
        # self.connect_markers(m for f,m in zip(fittable,self.markers) if f)

        # Add labels
        offsets = self.label_offsets()
        labels = self.sample_labels()
        for label_pos, label in zip(offsets, labels):
            fig.add_annotation(
                text=label, textangle=-30, x=label_pos, yref="paper", y=1.0, yanchor="bottom", showarrow=False
            )


def generate_best_profile(model: Experiment):
    if model.ismagnetic:
        best = model.magnetic_smooth_profile(dz=0.5)
    else:
        best = model.smooth_profile()
    return best


# ============================================================================= #
# Plotting script below
# ============================================================================= #


class ModelSpec(TypedDict):
    model_index: int
    sample_index: int


class PlotItem(TypedDict):
    model: Experiment
    spec: ModelSpec
    color_index: int


def plot_multiple_sld_profiles(plot_items: List[PlotItem]):
    import plotly.graph_objs as go

    fig = go.Figure()
    nplots = len(plot_items)
    multiplot = nplots > 1
    for plot_item in plot_items:
        model = plot_item["model"]
        spec = plot_item["spec"]
        color_index = plot_item["color_index"]
        model_index = spec["model_index"]
        sample_index = spec["sample_index"]
        if multiplot:
            colors = [MULTI_PLOT_COLORS[color_index % len(MULTI_PLOT_COLORS)]] * 4
        else:
            colors = SINGLE_PLOT_COLORS
        line_styles = MULTI_PLOT_STYLES if multiplot else SINGLE_PLOT_STYLES
        legendgroup = f"{model_index}:{sample_index} ({model.name})" if multiplot else None
        plot_sld_profile_plotly(model, fig, colors=colors, line_styles=line_styles, legendgroup=legendgroup)
    return fig


def plot_sld_profile_plotly(model, fig, colors=SINGLE_PLOT_COLORS, line_styles=SINGLE_PLOT_STYLES, legendgroup=None):
    if model.ismagnetic:
        z_best, rho_best, irho_best, rhoM_best, thetaM_best = generate_best_profile(model)
        yaxis_title = "SLD: ρ, ρ<sub>i</sub>, ρ<sub>M</sub> / 10<sup>-6</sup> Å<sup>-2</sup>"
    else:
        z_best, rho_best, irho_best = generate_best_profile(model)
        yaxis_title = "SLD: ρ, ρ<sub>i</sub> / 10<sup>-6</sup> Å<sup>-2</sup>"

    hovertemplate = (
        f"(%{{x}}, %{{y}})<br>{legendgroup if legendgroup is not None else ''} %{{data.hovertext}}<extra></extra>"
    )

    fig.add_scatter(
        x=z_best,
        y=rho_best,
        name="ρ",
        legendgroup=legendgroup,
        legendgrouptitle_text=legendgroup,
        hovertemplate=hovertemplate,
        hovertext="SLD",
        line={
            "color": colors[0],
            "dash": line_styles[0],
        },
    )

    fig.add_scatter(
        x=z_best,
        y=irho_best,
        name="ρ<sub>i</sub>",
        legendgroup=legendgroup,
        hovertemplate=hovertemplate,
        hovertext="Im. SLD",
        line={
            "color": colors[1],
            "dash": line_styles[1],
        },
    )

    if model.ismagnetic:
        fig.add_scatter(
            x=z_best,
            y=rhoM_best,
            name="ρ<sub>M</sub>",
            legendgroup=legendgroup,
            hovertemplate=hovertemplate,
            hovertext="Mag. SLD",
            line={
                "color": colors[2],
                "dash": line_styles[2],
            },
        )

        fig.add_scatter(
            x=z_best,
            y=thetaM_best,
            name="θ<sub>M</sub>",
            yaxis="y2",
            legendgroup=legendgroup,
            hovertemplate=hovertemplate,
            hovertext="Theta M",
            line={
                "color": colors[3],
                "dash": line_styles[3],
            },
        )

        # TODO: need to make axis scaling for thetaM dependent on if thetaM exceeds 0-360
        fig.update_layout(
            yaxis2={
                "title": {"text": "Magnetic Angle θ<sub>M</sub> / °"},
                "type": "linear",
                "autorange": False,
                "range": [0, 360],
                "anchor": "x",
                "overlaying": "y",
                "side": "right",
                "ticks": "inside",
                # 'ticklen': 20,
            }
        )

    fig.update_layout(uirevision=1, plot_bgcolor="white")
    fig.update_layout(
        xaxis={
            "title": {"text": "depth (Å)"},
            "type": "linear",
            "autorange": True,
            # 'gridcolor': "Grey",
            "ticks": "inside",
            # 'ticklen': 20,
            "showline": True,
            "linewidth": 1,
            "linecolor": "black",
            "mirror": "ticks",
            "side": "bottom",
        }
    )

    fig.update_layout(
        yaxis={
            "title": {"text": yaxis_title},
            "exponentformat": "e",
            "showexponent": "all",
            "type": "linear",
            "autorange": True,
            # 'gridcolor': "Grey",
            "ticks": "inside",
            # 'ticklen': 20,
            "showline": True,
            "linewidth": 1,
            "linecolor": "black",
            "mirror": True,
        }
    )

    fig.update_layout(margin={"l": 75, "r": 50, "t": 50, "b": 75, "pad": 4})

    fig.update_layout(
        legend={
            "x": -0.1,
            "bgcolor": "rgba(255,215,0,0.15)",
            # "traceorder": "reversed"
        }
    )

    marker_positions = FindLayers(model, axes=fig)
    marker_positions.reset_markers_plotly()

    # fig.show(renderer='firefox')
    return fig


# def plot_contours(model, title, savepath=None, nsamples=1000, show_contours=True, show_mag=True, savetitle=None,
#                   store=None, ultranest=False, dream=False, align=0):
#     if show_contours:
#         # data, columns = generate_contour_data(model)
#         if ultranest:
#             points_array = model.post_samples
#             sub_array = points_array[np.sort(np.random.randint(points_array.shape[0], size=nsamples)), :]
#             errors = calc_errors(model.problem, sub_array)
#             data, columns = generate_profile_data(errors)
#         # if this comment is removed the above code will not overwrite the previous reference to the errors object
#         # and it will act as a bad memory leak - using more ram everytime the function is called
#         if dream:
#             state, points = load_dream_state_notebook(problem=model.problem, store=store)
#             errors = calc_errors(model.problem, points[-nsamples:-1])
#             data, columns = generate_profile_data(errors, align)
#
#     # get best profile manually as we are not using DREAM state
#     if show_mag:
#         z_best, rho_best, irho_best, rhoM_best, thetaM_best = generate_best_profile(model)
#     else:
#         z_best, rho_best, irho_best = generate_best_profile(model)
#
#     fig, ax = plt.subplots(1, 1, figsize=(8, 6))
#
#     # rho
#     if show_contours:
#         z, best, sig2lo, sig2hi, siglo, sighi = data[0]
#         ax.plot(z, best, label=r"$\rho$ SLD", color="orange")
#         ax.fill_between(z, siglo, sighi, alpha=0.5, color="orange", label=r"$\rho$ SLD $1\sigma$")
#         ax.fill_between(z, sig2lo, sig2hi, alpha=0.25, color="orange", label=r"$\rho$ SLD $2\sigma$")
#     else:
#         ax.plot(z_best, rho_best, label=r"$\rho$ SLD", color="orange")
#
#     # irho
#     # ax.plot(z_best, irho_best, label=r"$\rho_{i}$ Im SLD")
#     # if show_contours:
#     #     z, best, sig2lo, sig2hi, siglo, sighi = data[1]
#     #     ax.fill_between(z, siglo, sighi, alpha=0.5)
#     #     ax.fill_between(z, sig2lo, sig2hi, alpha=0.25)
#     if show_mag:
#         if show_contours:
#             z, best, sig2lo, sig2hi, siglo, sighi = data[2]
#             ax.plot(z, best, label=r"$\rho_{M}$ mSLD", color="b")
#             ax.fill_between(z, siglo, sighi, alpha=0.5, color="b", label=r"$\rho_{M}$ mSLD $1\sigma$")
#             ax.fill_between(z, sig2lo, sig2hi, alpha=0.25, color="b", label=r"$\rho_{M}$ mSLD $2\sigma$")
#         else:
#             ax.plot(z_best, rhoM_best, label=r"$\rho_{M}$ mSLD", color="b")
#
#     # ax.legend(fontsize=14, loc='upper right', framealpha=0.5, ncol=2)
#     ax.legend(fontsize=16, framealpha=0.5, ncol=3)
#
#     ax.grid(True)
#     ax.set_xlabel(r"z $\left(\AA\right)$", fontsize=18)
#     ax.set_ylabel(r"$\rho, \rho_{i} \left(10^{-6}\AA^{-2}\right)$", fontsize=18)
#     if show_mag:
#         ax.set_ylabel(r"$\rho, \rho_{M} \left(10^{-6}\AA^{-2}\right)$", fontsize=18)
#     ax.tick_params(axis='y', labelsize=16)
#     ax.tick_params(axis='x', labelsize=16)
#
#     experiment = model.determine_problem_fitness()
#
#     layer_markers = FindLayers(experiment, axes=ax)
#     layer_markers.reset_markers()
#     fig.suptitle(title)
#     # if np.max(rho_best) > 75:
#     #     top = np.max(rho_best)*1.15
#     # else:
#     #     top =75
#     # ax.set_ylim(top=top)
#     if savepath:
#         if savetitle:
#             title = savetitle
#         plt.savefig(f"{savepath}/{title}.png")
