# (c) Kevin Dunn, 2010-2021. MIT License. Based on own private work over the years.

# Built-in libraries
from typing import Dict

import plotly.graph_objects as go
from pydantic import BaseModel


def plot_pre_checks(model, pc_horiz, pc_vert) -> bool:
    assert (
        pc_horiz <= model.A
    ), f"The model has {model.A} components. Ensure that pc_horiz<={model.A}."
    assert (
        pc_vert <= model.A
    ), f"The model has {model.A} components. Ensure that pc_vert<={model.A}."
    assert pc_horiz != pc_vert, "Specify two different components: pc_horiz != pc_vert"
    return True


def score_plot(
    model, pc_horiz=1, pc_vert=2, settings: Dict = None, fig=None
) -> go.Figure:
    """Generates a 2-dimensional score plot for the given latent variable model.

    Parameters
    ----------
    model : MVmodel object (PCA, or PLS)
        A latent variable model generated by this library.
    pc_horiz : int, optional
        Which component to plot on the horizontal axis, by default 1 (the first component)
    pc_vert : int, optional
        Which component to plot on the vertical axis, by default 2 (the second component)
    settings : dict
        Default settings are = {
            "show_ellipse": True [bool],
                Should the Hotelling's T2 ellipse be added

            "ellipse_conf_level": 0.95 [float]
                If the ellipse is added, which confidence level is used. A number < 1.00.

            "title": f"Score plot of component {pc_horiz} vs component {pc_vert}"
                Overall plot title

            "show_observation_labels": False,
                Adds a label for each observation. Labels are always available in the hover.

            "show_legend": True,
                Shows a clickable legend (allows to turn the ellipse(s) on/off)

            "html_image_height": 500,
                in pixels

            "html_aspect_ratio_w_over_h": 16/9,
                sets the image width, as a ratio of the height

        }
    """
    plot_pre_checks(model, pc_horiz, pc_vert)
    margin_dict: Dict = dict(l=10, r=10, b=5, t=80)  # Defaults: l=80, r=80, t=100, b=80

    class Settings(BaseModel):
        show_ellipse: bool = True
        ellipse_conf_level: float = 0.95  # TODO: check constraint
        title: str = f"Score plot of component {pc_horiz} vs component {pc_vert}"
        show_observation_labels: bool = False  # TODO
        show_legend: bool = True
        html_image_height: float = 500.0
        html_aspect_ratio_w_over_h: float = 16 / 9.0

    if settings:
        setdict = Settings(settings).dict()
    else:
        setdict = Settings().dict()
    if fig is None:
        fig = go.Figure()

    fig = model.x_scores.plot.scatter(x=pc_horiz, y=pc_vert, text=model.x_scores.index)
    fig.update_traces(textposition="top center")
    ellipse = model.ellipse_coordinates(
        score_horiz=pc_horiz,
        score_vert=pc_vert,
        T2_limit_conf_level=setdict["ellipse_conf_level"],
    )
    fig.add_trace(go.Scatter(x=ellipse[0], y=ellipse[1], name="Hotelling's T^2 [95%]"))
    fig.update_layout(
        xaxis_title_text=f"PC {pc_horiz}", yaxis_title_text=f"PC {pc_vert}"
    )
    fig.add_hline(y=0, line_color="black")
    fig.add_vline(x=0, line_color="black")
    fig.update_layout(
        title_text=setdict["title"],
        margin=margin_dict,
        hovermode="closest",
        showlegend=setdict["show_legend"],
        legend=dict(
            orientation="h",
            traceorder="normal",
            font=dict(family="sans-serif", size=12, color="#000"),
            bordercolor="#DDDDDD",
            borderwidth=1,
        ),
        autosize=False,
        xaxis=dict(
            gridwidth=1,
            mirror=True,  # ticks are mirror at the top of the frame also
            showspikes=True,
            visible=True,
        ),
        yaxis=dict(
            gridwidth=2,
            type="linear",
            autorange=True,
            showspikes=True,
            visible=True,
            showline=True,  # show a separating line
            side="left",  # show on the RHS
        ),
        width=setdict["html_aspect_ratio_w_over_h"] * setdict["html_image_height"],
        height=setdict["html_image_height"],
    )
    return fig


def loadings_plot(
    model, loadings_type="p", pc_horiz=1, pc_vert=2, settings: Dict = None, fig=None
) -> go.Figure:
    """Generates a 2-dimensional loadings for the given latent variable model.

    Parameters
    ----------
    model : MVmodel object (PCA, or PLS)
        A latent variable model generated by this library.

    loadings_type : str, optional
        A choice of the following:
            'p' : (default for PCA) : the P (projection) loadings: only option possible for PCA
            'w' : the W loadings: Suitable for PLS
            'w*' : (default for PLS) the W* (or R) loadings: Suitable for PLS
            'w*c' : the W* (from X-space) with C loadings from the Y-space: Suitable for PLS
            'c' : the C loadings from the Y-space: Suitable for PLS

        For PCA model any other choice besides 'p' will be ignored.

    pc_horiz : int, optional
        Which component to plot on the horizontal axis, by default 1 (the first component)
    pc_vert : int, optional
        Which component to plot on the vertical axis, by default 2 (the second component)
    settings : dict
        Default settings are = {
            "title": f"Loadings plot of component {pc_horiz} vs component {pc_vert}"
                Overall plot title

            "show_column_labels": True,
                Adds a label for each column. Labels are always available in the hover.

            "html_image_height": 500,
                in pixels

            "html_aspect_ratio_w_over_h": 16/9,
                sets the image width, as a ratio of the height

        }
    """
    plot_pre_checks(model, pc_horiz, pc_vert)
    margin_dict: Dict = dict(l=10, r=10, b=5, t=80)  # Defaults: l=80, r=80, t=100, b=80

    class Settings(BaseModel):
        title: str = (
            f"Loadings plot [{loadings_type.upper()}] of component {pc_horiz} vs "
            f"component {pc_vert}"
        )
        show_column_labels: bool = True  # TODO
        html_image_height: float = 500.0
        html_aspect_ratio_w_over_h: float = 16 / 9.0

    if settings:
        setdict = Settings(settings).dict()
    else:
        setdict = Settings().dict()
    if fig is None:
        fig = go.Figure()

    what = model.x_loadings  # PCA default
    if hasattr(model, "direct_weights"):
        what = model.direct_weights  # PLS default
    extra = None
    if loadings_type.lower() == "p":
        what = model.x_loadings
    if loadings_type.lower() == "w":
        what = model.x_weights
    elif loadings_type.lower() == "w*":
        what = model.direct_weights
    elif loadings_type.lower() == "w*c":
        loadings_type = loadings_type[0:-1]
        what = model.direct_weights
        extra = model.y_loadings
    elif loadings_type.lower() == "c":
        what = model.y_loadings

    fig.add_trace(
        go.Scatter(
            x=what.loc[:, pc_horiz],
            y=what.loc[:, pc_vert],
            name="X-space loadings W*",
            mode="markers+text",
            marker=dict(
                color="darkblue",
                symbol="circle",
            ),
            marker_size=7,
            text=what.index,
            textposition="top center",
        )
    )
    add_legend = False

    # Note, we have cut off the 'c' from loadings_type
    add_legend = False
    if loadings_type.lower() == "w*" and extra is not None:
        add_legend = True
        fig.add_trace(
            go.Scatter(
                x=extra.loc[:, pc_horiz],
                y=extra.loc[:, pc_vert],
                name="Y-space loadings C",
                mode="markers+text",
                marker=dict(
                    color="purple",
                    symbol="star",
                ),
                marker_size=7,
                text=extra.index,
                textposition="bottom center",
            )
        )

    fig.update_layout(
        xaxis_title_text=f"PC {pc_horiz}", yaxis_title_text=f"PC {pc_vert}"
    )
    fig.add_hline(y=0, line_color="black")
    fig.add_vline(x=0, line_color="black")
    fig.update_layout(
        title_text=setdict["title"],
        margin=margin_dict,
        hovermode="closest",
        showlegend=add_legend,
        autosize=False,
        xaxis=dict(
            gridwidth=1,
            mirror=True,  # ticks are mirror at the top of the frame also
            showspikes=True,
            visible=True,
        ),
        yaxis=dict(
            gridwidth=2,
            type="linear",
            autorange=True,
            showspikes=True,
            visible=True,
            showline=True,  # show a separating line
            side="left",  # show on the RHS
        ),
        width=setdict["html_aspect_ratio_w_over_h"] * setdict["html_image_height"],
        height=setdict["html_image_height"],
    )
    return fig


def spe_plot(model, with_a=-1, settings: Dict = None, fig=None) -> go.Figure:
    """Generates a squared-prediction error (SPE) plot for the given latent variable model using
    `with_a` number of latent variables. The default will use the total number of latent variables
    which have already been fitted.

    Parameters
    ----------
    model : MVmodel object (PCA, or PLS)
        A latent variable model generated by this library.
    with_a : int, optional
        Uses this many number of latent variables, and therefore shows the SPE after this number of
        model components. By default the total number of components fitted will be used.
    settings : dict
        Default settings are = {
            "show_limit": True [bool],
                Should the SPE limit be plotted.

            "conf_level": 0.95 [float]
                If the limit line is added, which confidence level is used. Number < 1.00.

            "title": f"Squared prediction error plot after fitting {with_a} components,
                       with the {conf_level*100}% confidence limit"
                Overall plot title

            "show_observation_labels": False,
                Adds a label for each observation. Labels are always available in the hover.

            "show_legend": True,
                Shows a clickable legend (allows to turn the limit on/off)

            "html_image_height": 500,
                Image height, in pixels.

            "html_aspect_ratio_w_over_h": 16/9,
                Sets the image width, as a ratio of the height.

        }
    """
    # TO CONSIDER: allow a setting `as_line`: which connects the points with line segments
    margin_dict: Dict = dict(l=10, r=10, b=5, t=80)  # Defaults: l=80, r=80, t=100, b=80

    if with_a < 0:
        with_a = model.squared_prediction_error.columns[with_a]

    # TODO: check `with_a`: what should it plot if `with_a` is zero, or > A?

    class Settings(BaseModel):
        show_limit: bool = True
        conf_level: float = 0.95  # TODO: check constraint < 1
        title: str = (
            "Squared prediction error plot after "
            f"fitting {with_a} component{'s' if with_a > 1 else ''}"
            f", with the {conf_level*100}% confidence limit"
        )

        show_observation_labels: bool = False  # TODO
        show_legend: bool = True
        html_image_height: float = 500.0
        html_aspect_ratio_w_over_h: float = 16 / 9.0

    if settings:
        setdict = Settings(settings).dict()
    else:
        setdict = Settings().dict()
    if fig is None:
        fig = go.Figure()

    fig = model.squared_prediction_error.loc[:, [with_a]].plot.scatter(
        x=model.squared_prediction_error.index,
        y=model.squared_prediction_error.columns[with_a - 1],
    )
    limit_SPE_conf_level = model.SPE_limit(conf_level=setdict["conf_level"])

    fig.add_hline(
        y=limit_SPE_conf_level,
        line_color="red",
        annotation_text=f'{setdict["conf_level"]*100:.3g}% limit',
        annotation_position="bottom right",
    )
    fig.add_hline(y=0, line_color="black")
    fig.update_layout(
        title_text=setdict["title"],
        margin=margin_dict,
        hovermode="closest",
        showlegend=setdict["show_legend"],
        legend=dict(
            orientation="h",
            traceorder="normal",
            font=dict(family="sans-serif", size=12, color="#000"),
            bordercolor="#DDDDDD",
            borderwidth=1,
        ),
        autosize=False,
        xaxis=dict(
            gridwidth=1,
            mirror=True,  # ticks are mirror at the top of the frame also
            showspikes=True,
            visible=True,
        ),
        yaxis=dict(
            title=f"SPE values after fitting {with_a} component{'s' if with_a > 1 else ''}",
            gridwidth=2,
            type="linear",
            autorange=True,
            showspikes=True,
            visible=True,
            showline=True,  # show a separating line
            side="left",  # show on the RHS
        ),
        width=setdict["html_aspect_ratio_w_over_h"] * setdict["html_image_height"],
        height=setdict["html_image_height"],
    )
    return fig


def t2_plot(model, with_a=-1, settings: Dict = None, fig=None) -> go.Figure:
    """Generates a Hotelling's T2 (T^2) plot for the given latent variable model using
    `with_a` number of latent variables. The default will use the total number of latent variables
    which have already been fitted.

    Parameters
    ----------
    model : MVmodel object (PCA, or PLS)
        A latent variable model generated by this library.
    with_a : int, optional
        Uses this many number of latent variables, and therefore shows the SPE after this number of
        model components. By default the total number of components fitted will be used.
    settings : dict
        Default settings are = {
            "show_limit": True [bool],
                Should the T2 limit be plotted.

            "conf_level": 0.95 [float]
                If the limit line is added, which confidence level is used. Number < 1.00.

            "title": f"Hotelling's T2 plot after fitting {with_a} components,
                       with the {conf_level*100}% confidence limit""
                Overall plot title

            "show_observation_labels": False,
                Adds a label for each observation. Labels are always available in the hover.

            "show_legend": True,
                Shows a clickable legend (allows to turn the limit on/off)

            "html_image_height": 500,
                Image height, in pixels.

            "html_aspect_ratio_w_over_h": 16/9,
                Sets the image width, as a ratio of the height.
        }
    """
    # TO CONSIDER: allow a setting `as_line`: which connects the points with line segments
    margin_dict: Dict = dict(l=10, r=10, b=5, t=80)  # Defaults: l=80, r=80, t=100, b=80

    if with_a < 0:
        with_a = model.Hotellings_T2.columns[with_a]

    # TODO: check `with_a`: what should it plot if `with_a` is zero, or > A?

    class Settings(BaseModel):
        show_limit: bool = True
        conf_level: float = 0.95  # TODO: check constraint < 1
        title: str = (
            f"Hotelling's T2 plot after fitting {with_a} component{'s' if with_a > 1 else ''}"
            f", with the {conf_level*100}% confidence limit"
        )
        show_observation_labels: bool = False  # TODO
        show_legend: bool = True
        html_image_height: float = 500.0
        html_aspect_ratio_w_over_h: float = 16 / 9.0

    if settings:
        setdict = Settings(settings).dict()
    else:
        setdict = Settings().dict()
    if fig is None:
        fig = go.Figure()

    fig = model.Hotellings_T2.loc[:, [with_a]].plot.scatter(
        x=model.Hotellings_T2.index,
        y=model.Hotellings_T2.columns[with_a - 1],
    )
    limit_HT2_conf_level = model.T2_limit(conf_level=setdict["conf_level"])

    fig.add_hline(
        y=limit_HT2_conf_level,
        line_color="red",
        annotation_text=f'{setdict["conf_level"]*100:.3g}% limit',
        annotation_position="bottom right",
    )
    fig.add_hline(y=0, line_color="black")
    fig.update_layout(
        title_text=f"T2 values after fitting {with_a} component{'s' if with_a > 1 else ''}",
        margin=margin_dict,
        hovermode="closest",
        showlegend=setdict["show_legend"],
        legend=dict(
            orientation="h",
            traceorder="normal",
            font=dict(family="sans-serif", size=12, color="#000"),
            bordercolor="#DDDDDD",
            borderwidth=1,
        ),
        autosize=False,
        xaxis=dict(
            gridwidth=1,
            mirror=True,  # ticks are mirror at the top of the frame also
            showspikes=True,
            visible=True,
        ),
        yaxis=dict(
            title_text=setdict["title"],
            gridwidth=2,
            type="linear",
            autorange=True,
            showspikes=True,
            visible=True,
            showline=True,  # show a separating line
            side="left",  # show on the RHS
        ),
        width=setdict["html_aspect_ratio_w_over_h"] * setdict["html_image_height"],
        height=setdict["html_image_height"],
    )
    return fig
