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

from pyoptim.optimizers import GradientDescent

FIG_WIDTH = 7


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"] = 10
    rcParams["axes.labelsize"] = 10
    rcParams["axes.titlesize"] = 10
    rcParams["legend.fontsize"] = 8
    rcParams["xtick.labelsize"] = 8
    rcParams["ytick.labelsize"] = 8
    # 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, full_width: bool = False
) -> 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.
    """
    if full_width:
        return (FIG_WIDTH, FIG_WIDTH * rel_height)
    return (FIG_WIDTH * 0.5, FIG_WIDTH * 0.5 * 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"])

    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))
    return fig, ax


def scatter_dataset(
    x_train: np.ndarray, x_test: np.ndarray, y_train: np.ndarray, y_test: np.ndarray
) -> tuple[plt.Figure, plt.Axes]:
    fig, ax = plt.subplots(figsize=get_figsize(0.5))
    ax.scatter(x_train, y_train, label="Train Data", s=1)
    ax.scatter(x_test, y_test, label="Test Data", s=1)
    ax.set_xlabel("$x$")
    ax.set_ylabel("$y$")
    ax.legend()
    return fig, ax


def get_cost_history(
    optimizer: type[GradientDescent], X: np.ndarray, y: np.ndarray, **kwargs
) -> tuple[np.ndarray, np.ndarray]:
    opt = optimizer(**kwargs)
    opt.fit(X, y)
    iterations = opt.get_epochs()
    return iterations, opt.cost_history


def plot_optimizers(
    ax: plt.Axes,
    optimizer_list: list[tuple[type[GradientDescent], dict]],
    X: np.ndarray,
    y: np.ndarray,
    ylabel="Cost",
    labels: list | None = None,
) -> None:
    for i, (opt_class, params) in enumerate(optimizer_list):
        iter, history = get_cost_history(opt_class, X, y, **params)
        if labels:
            label = labels[i]
        else:
            label = f"{opt_class.__name__}"
        ax.plot(iter, history, label=label)
    ax.set_xlabel("Iteration")
    ax.set_ylabel(ylabel)
    ax.legend()


def optimization_performance_evaluation(
    optimizers: list[GradientDescent],
    labels: list[str],
    X_tr: np.ndarray,
    y_tr: np.ndarray,
):
    fig, ax = plt.subplots(figsize=get_figsize(0.6))

    for optimizer, label in zip(optimizers, labels):
        start_time = time.time()
        optimizer.fit(X_tr, y_tr)
        time_taken = time.time() - start_time
        cost_history = optimizer._cost_history
        if "Stochastic" in label:
            b_p_e = getattr(optimizer, "batches_per_epoch", 1)
            epochs = (
                np.arange(1, optimizer.num_iterations + 1, dtype=np.float64) / b_p_e
            )
            ls = "-"
        else:
            epochs = np.arange(1, optimizer.num_iterations + 1, dtype=np.float64)
            ls = "--"
        c = "C0" if "Ridge" not in label else "C1"
        ax.plot(
            epochs,
            cost_history,
            label=f"{label} ({time_taken:.2f} s)",
            linestyle=ls,
            color=c,
        )
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Cost")
    ax.legend()
    return fig, ax
