import matplotlib.pyplot as plt
from matplotlib.colors import SymLogNorm
import numpy as np

FIG_WIDTH = 6


def get_rc_params():
    colors = ["FF220C", "70D6FF", "8AAA79", "666370", "1C1F33"]
    rcParams = plt.rcParams
    # Use LaTeX for rendering
    # Setup fonts
    rcParams["text.usetex"] = True
    rcParams["font.family"] = "serif"
    rcParams["font.size"] = 12
    rcParams["axes.labelsize"] = 12
    rcParams["axes.titlesize"] = 12
    rcParams["legend.fontsize"] = 10
    rcParams["xtick.labelsize"] = 10
    rcParams["ytick.labelsize"] = 10
    # Figure size and resolution
    rcParams["figure.figsize"] = (FIG_WIDTH, 4)
    rcParams["figure.dpi"] = 300
    # Use colors from the palette
    rcParams["axes.prop_cycle"] = plt.cycler(color=[f"#{color}" for color in colors])
    # Grid
    rcParams["axes.grid"] = True
    rcParams["grid.alpha"] = 0.5
    rcParams["grid.linestyle"] = "--"
    # Point ticks to the inside of the axes
    rcParams["xtick.direction"] = "in"
    rcParams["ytick.direction"] = "in"
    rcParams["xtick.top"] = True
    rcParams["ytick.right"] = True
    return rcParams


def set_rc_params():
    plt.rcParams.update(get_rc_params())


set_rc_params()


def get_figsize(rel_height: float = 2 / 3) -> tuple[float, float]:
    """Returns a figure size tuple based on the global FIG_WIDTH and a relative height.
    Args:
        rel_height: The relative height of the figure compared to FIG_WIDTH.
    Returns:
        A tuple (width, height) for the figure size.
    """
    return (FIG_WIDTH, FIG_WIDTH * rel_height)


def mse_r2_plot(
    polynomial_degrees: np.ndarray,
    train_mse_list: list[float],
    mse_list: list[float],
    train_r2_list: list[float],
    r2_list: list[float],
    labels: dict[str, str] | None = None,
) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes]]:
    std_labels = {
        "set1": "Train Set",
        "set2": "Test Set",
        "xlabel": "Polynomial Degree",
        "ylabel1": "Mean Squared Error",
        "ylabel2": "$R^2$ Score",
    }

    if labels is None:
        labels = std_labels
    else:
        for key in std_labels:
            if key not in labels:
                labels[key] = std_labels[key]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=get_figsize(0.5))
    ax1.plot(polynomial_degrees, train_mse_list, marker="o", label=labels["set1"])
    ax1.plot(polynomial_degrees, mse_list, marker="o", label=labels["set2"])
    ax1.legend()
    ax1.set_xlabel(labels["xlabel"])
    ax1.set_ylabel(labels["ylabel1"])

    ax2.plot(polynomial_degrees, train_r2_list, marker="o", label=labels["set1"])
    ax2.plot(polynomial_degrees, r2_list, marker="o", label=labels["set2"])
    ax2.legend()
    ax2.set_xlabel(labels["xlabel"])
    ax2.set_ylabel(labels["ylabel2"])

    fig.tight_layout()
    return fig, (ax1, ax2)


def parameter_plot(
    polynomial_degrees: np.ndarray,
    beta_OLS_list: list[np.ndarray],
    labels: dict[str, str] | None = None,
) -> tuple[plt.Figure, plt.Axes]:
    std_labels = {
        "xlabel": "Polynomial Degree of Fit",
        "ylabel": "Coefficient Index $i$",
        "cbar": "Coefficient Value",
    }

    if labels is None:
        labels = std_labels
    else:
        for key in std_labels:
            if key not in labels:
                labels[key] = std_labels[key]

    fig, ax = plt.subplots(figsize=get_figsize(0.5))
    beta_OLS = np.zeros((beta_OLS_list[-1].shape[0], len(polynomial_degrees)))
    beta_OLS[:] = np.nan
    for i in range(len(beta_OLS_list)):
        for j in range(len(beta_OLS_list[i])):
            beta_OLS[j, i] = beta_OLS_list[i][j]
    cmap = plt.get_cmap("coolwarm")
    norm = SymLogNorm(
        linthresh=1e-3, vmin=np.nanmin(beta_OLS), vmax=np.nanmax(beta_OLS)
    )
    im = ax.imshow(beta_OLS, aspect="auto", cmap=cmap, origin="lower", norm=norm)
    cbar = fig.colorbar(im, ax=ax)
    cbar.set_label(labels["cbar"])
    ax.set_xlabel(labels["xlabel"])
    ax.set_ylabel(labels["ylabel"])
    ax.set_xticks(np.arange(0, len(polynomial_degrees), step=3))
    ax.set_xticklabels(polynomial_degrees[::3])
    ax.set_yticks(np.arange(0, beta_OLS.shape[0], 5))
    ax.set_yticklabels(np.arange(1, beta_OLS.shape[0] + 1, step=5))
    fig.tight_layout()
    return fig, ax