import os.path
from typing import Literal, Tuple, Union

import pandas as pd
from matplotlib import pyplot as plt

from pyafs.graphics.utils import export_figure, set_axis_ticks


def plot_norm_spec(
    spec_df: pd.DataFrame,
    key_prefix: Literal["primitive", "final"],
) -> Tuple[plt.Figure, list[plt.Axes]]:
    """Plot the spectrum in both scaled and normalised."""
    colour_dict = {"primitive": "tab:red", "final": "tab:green"}

    fig, (scaled_ax, norm_ax) = plt.subplots(
        2,
        1,
        figsize=(10, 7),
        dpi=300,
        height_ratios=[4, 3],
        gridspec_kw=dict(hspace=0.1),
    )

    if spec_df["is_outlier"].any():
        scaled_ax.plot(
            spec_df["wvl"],
            spec_df["scaled_intensity"],
            "-",
            c="grey",
            lw=1,
            alpha=0.8,
            label=f"source {key_prefix} spec.",
        )
        norm_ax.plot(
            spec_df["wvl"],
            spec_df[f"{key_prefix}_norm_intensity"],
            "-",
            c="grey",
            lw=1,
            alpha=0.8,
            label=f"source {key_prefix} spec.",
        )

        cleanup_spec_df = spec_df[~spec_df["is_outlier"]]
        scaled_ax.plot(
            cleanup_spec_df["wvl"],
            cleanup_spec_df["scaled_intensity"],
            "-",
            c="k",
            lw=1,
            alpha=0.8,
            label=f"non-outlier {key_prefix} spec.",
        )
        norm_ax.plot(
            cleanup_spec_df["wvl"],
            cleanup_spec_df[f"{key_prefix}_norm_intensity"],
            "-",
            c="k",
            lw=1,
            alpha=0.8,
            label=f"non-outlier {key_prefix} spec.",
        )
    else:
        scaled_ax.plot(
            spec_df["wvl"],
            spec_df["scaled_intensity"],
            "k-",
            lw=1,
            alpha=0.8,
            label=f"{key_prefix} spec.",
        )
        norm_ax.plot(
            spec_df["wvl"],
            spec_df[f"{key_prefix}_norm_intensity"],
            "k-",
            lw=1,
            alpha=0.8,
            label=f"{key_prefix} norm spec.",
        )

    scaled_ax.plot(
        spec_df["wvl"],
        spec_df[f"{key_prefix}_blaze"],
        "-",
        c=colour_dict[key_prefix],
        lw=1,
        alpha=0.8,
        label=f"{key_prefix} blaze",
    )

    for axis in [scaled_ax, norm_ax]:
        set_axis_ticks(axis)
        axis.tick_params(labelsize="large")

    norm_ax.axhline(1, ls=":", c="k", lw=1, alpha=0.8)
    norm_ax.set_ylim(0, 1.2)
    norm_ax.yaxis.set_major_locator(plt.MultipleLocator(0.3, offset=0.1))

    scaled_ax.set_ylabel("scaled intensity", fontsize="x-large")
    scaled_ax.set_ylim(None, int(scaled_ax.get_ylim()[1]) + 1)

    norm_ax.set_ylabel("normalised intensity", fontsize="x-large")
    norm_ax.set_xlabel("wavelength", fontsize="x-large")

    return fig, [scaled_ax, norm_ax]


def plot_primitive_norm_spec(
    spec_df: pd.DataFrame,
    debug: Union[bool, str] = False,
) -> None:
    """Plot the primitive normalised spectrum."""
    fig, (scaled_ax, norm_ax) = plot_norm_spec(spec_df, "primitive")

    scaled_ax.plot(
        spec_df["wvl"],
        spec_df["tilde_AS_alpha"],
        "-",
        lw=1,
        c="tab:cyan",
        alpha=0.8,
        label="$\\widetilde{AS}_{\\alpha}$",
    )
    scaled_ax.plot(
        spec_df["wvl"],
        spec_df["tilde_AS_alpha"] - spec_df["primitive_blaze"],
        "-",
        lw=1,
        c="tab:purple",
        alpha=0.8,
        label="$\\widetilde{AS}_{\\beta}$ - primitive blaze",
    )
    scaled_ax.axhline(0, ls=":", c="k", lw=1, alpha=0.8)
    scaled_ax.plot(
        spec_df[spec_df["is_intersect_with_tilde_AS_alpha"]]["wvl"],
        spec_df[spec_df["is_intersect_with_tilde_AS_alpha"]]["scaled_intensity"],
        "o",
        ms=6,
        mew=1,
        mec="tab:orange",
        mfc="None",
        alpha=0.8,
        label="$\\widetilde{AS}_{\\alpha}$ $\cap$ spec.",
    )

    # different legend layouts for DataFrame with and without outliers
    legend_n_col = 3 if spec_df["is_outlier"].any() else 5

    for axis in [scaled_ax, norm_ax]:
        axis.legend(
            fontsize="medium",
            ncol=legend_n_col,
            handlelength=1.4,
            columnspacing=1.6,
            loc="upper center" if axis == scaled_ax else "lower center",
        )

    fig.suptitle("Primitive Normalised Spectrum", fontsize="xx-large", y=0.92)

    if isinstance(debug, str):
        export_figure(fig, filename=os.path.join(debug, "primitive_norm.png"))
    elif debug:
        export_figure(fig)


def plot_afs_final_norm_spec(
    spec_df: pd.DataFrame, debug: Union[bool, str] = False
) -> None:
    """Plot the final normalised spectrum."""
    fig, (scaled_ax, norm_ax) = plot_norm_spec(spec_df, "final")

    scaled_ax.plot(
        spec_df[spec_df["is_above_quantile"]]["wvl"],
        spec_df[spec_df["is_above_quantile"]]["scaled_intensity"],
        "x",
        ms=6,
        mew=1,
        mec="tab:green",
        mfc="None",
        alpha=0.8,
        label=f"selected pixels",
    )
    scaled_ax.plot(
        spec_df[spec_df["is_intersect_with_tilde_AS_alpha"]]["wvl"],
        spec_df[spec_df["is_intersect_with_tilde_AS_alpha"]]["scaled_intensity"],
        "o",
        ms=6,
        mew=1,
        mec="tab:orange",
        mfc="None",
        alpha=0.8,
        label="$\widetilde{\mathrm{AS}}_{\\alpha}$ $\cap$ spec.",
    )
    scaled_ax.plot(
        spec_df["wvl"],
        spec_df["primitive_blaze"],
        "-",
        c="tab:red",
        lw=1,
        alpha=0.8,
        label="primitive blaze",
    )

    # different legend layouts for DataFrame with and without outliers
    legend_n_col = 3 if spec_df["is_outlier"].any() else 5

    for axis in [scaled_ax, norm_ax]:
        axis.legend(
            fontsize="medium",
            ncol=legend_n_col,
            handlelength=1.4,
            columnspacing=1.6,
            loc="upper center" if axis == scaled_ax else "lower center",
        )

    fig.suptitle("Final Normalised Spectrum", fontsize="xx-large", y=0.92)

    if isinstance(debug, str):
        export_figure(fig, filename=os.path.join(debug, "final_norm.png"))
    elif debug:
        export_figure(fig)
