from __future__ import annotations

import warnings
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
from sklearn.exceptions import ConvergenceWarning
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ConstantKernel as C
from sklearn.gaussian_process.kernels import RBF


BASE_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = BASE_DIR / "outputs"

warnings.filterwarnings("ignore", category=ConvergenceWarning)
warnings.filterwarnings("ignore", category=RuntimeWarning)


def ensure_dir(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


def mod_10(x):
    return np.asarray(x) % 10


def mle_entropy(arr) -> float:
    data = np.asarray(arr)
    if data.ndim == 1:
        data = data.reshape(-1, 1)
    if data.ndim > 2:
        raise ValueError("Data must be 1D or 2D")

    _, counts = np.unique(data, axis=0, return_counts=True)
    p = counts / data.shape[0]
    return float(-np.sum(p * np.log(p)))


def fit_gp(x, y, n_restarts_optimizer: int = 10, random_state: int | None = None) -> GaussianProcessRegressor:
    kernel = C(1.0, (1e-3, 1e3)) * RBF(1.0, (1e-2, 1e2))
    gp = GaussianProcessRegressor(
        kernel=kernel,
        alpha=1e-2,
        n_restarts_optimizer=n_restarts_optimizer,
        random_state=random_state,
        normalize_y=False,
    )
    gp.fit(np.asarray(x).reshape(-1, 1), np.asarray(y).ravel())
    return gp


def shannon_surprise(data_old, data_new, gp=None, p_floor=1e-3) -> float:
    if gp is None:
        gp = fit_gp(data_old[:, 0], data_old[:, 1])

    x_new = np.asarray([data_new[0]], dtype=float).reshape(1, -1)
    y_new = float(data_new[1])
    mean, std = gp.predict(x_new, return_std=True)
    std = max(float(std[0]), 1e-12)
    p = (1.0 / (std * np.sqrt(2 * np.pi))) * np.exp(
        -0.5 * ((y_new - float(mean[0])) ** 2) / (std**2)
    )
    if p >= p_floor:
        return float(-np.log(p))
    return float("inf")


def bayesian_surprise(gp, x_new, y_new) -> float:
    x_new = np.asarray(x_new).reshape(1, -1)
    y_new = np.asarray([float(np.asarray(y_new).ravel()[0])])

    mean_prior, std_prior = gp.predict(x_new, return_std=True)
    var_prior = np.maximum(std_prior**2, 1e-12)

    x_updated = np.vstack([gp.X_train_, x_new])
    y_updated = np.hstack([gp.y_train_, y_new])
    gp_updated = GaussianProcessRegressor(
        kernel=gp.kernel_,
        alpha=gp.alpha,
        optimizer=None,
        normalize_y=gp.normalize_y,
    )
    gp_updated.fit(x_updated, y_updated)

    mean_post, std_post = gp_updated.predict(x_new, return_std=True)
    var_post = np.maximum(std_post**2, 1e-12)

    kl = 0.5 * (
        (var_post / var_prior)
        + ((mean_prior - mean_post) ** 2) / var_prior
        - 1
        + np.log(var_prior / var_post)
    )
    return float(np.ravel(kl)[0])


SCENARIOS = {
    "explore": {
        "seed": 1001,
        "xlabel": "Number of Explorations (m)",
        "init": lambda: (np.random.randint(0, 30, 100), None),
        "stream": lambda n: (np.random.randint(30, 60, n), None),
        "stream_y": lambda x: mod_10(x),
        "sb_file": "explore_SB.pdf",
        "mis_file": "explore_MIS.pdf",
        "sb_ylim": (0, 8),
        "bayes_ylim": (0, 20),
        "legend_loc": "best",
        "mis_legend_loc": [0.45, 0.7],
    },
    "over_exploitation": {
        "seed": 1002,
        "xlabel": "Number of Exploitations (m)",
        "init": lambda: (np.random.randint(0, 30, 100), None),
        "stream": lambda n: (np.full(n, 7), None),
        "stream_y": lambda x: mod_10(x),
        "sb_file": "over_exploitation_SB.pdf",
        "mis_file": "over_exploitation_MIS.pdf",
        "sb_ylim": None,
        "bayes_ylim": None,
        "legend_loc": "best",
        "mis_legend_loc": [0.45, 0.6],
    },
    "noisy_explore": {
        "seed": 1003,
        "xlabel": "Number of Explorations (m)",
        "init": lambda: (np.random.randint(0, 30, 100), None),
        "stream": lambda n: (np.random.randint(30, 60, n), np.random.randint(0, 10, n)),
        "stream_y": None,
        "sb_file": "noisy_explore_SB.pdf",
        "mis_file": "noisy_explore_MIS.pdf",
        "sb_ylim": (0, 8),
        "bayes_ylim": (0, 8),
        "legend_loc": [0.6, 0.1],
        "mis_legend_loc": [0.45, 0.7],
    },
    "pure_explore": {
        "seed": 1004,
        "xlabel": "Number of Explorations (m)",
        "init": lambda: (np.random.randint(0, 5, 100), None),
        "stream": lambda n: (np.random.choice(np.arange(5, 500), size=n, replace=False), None),
        "stream_y": lambda x: mod_10(x),
        "sb_file": "pure_explore_SB.pdf",
        "mis_file": "pure_explore_MIS.pdf",
        "sb_ylim": (0, 8),
        "bayes_ylim": (0, 20),
        "legend_loc": [0.14, 0.7],
        "mis_legend_loc": [0.45, 0.3],
        "multiframe_shannon_ylim": (0, 100),
    },
    "noise_decrease": {
        "seed": 1005,
        "xlabel": "Number of Samplings (m)",
        "init": lambda: (np.random.randint(0, 30, 100), np.random.randint(0, 5, 100)),
        "stream": lambda n: (np.random.randint(0, 30, n), None),
        "stream_y": lambda x: mod_10(x),
        "sb_file": "noise_decrease_SB.pdf",
        "mis_file": "noise_decrease_MIS.pdf",
        "sb_ylim": (0, 8),
        "bayes_ylim": (0, 20),
        "legend_loc": [0.01, 0.1],
        "mis_legend_loc": [0.45, 0.4],
    },
    "explorenew": {
        "seed": 1006,
        "xlabel": "Number of Explorations (m)",
        "init": lambda: (np.random.randint(0, 30, 100), None),
        "stream": lambda n: (np.random.randint(30, 60, n), None),
        "stream_y": lambda x: mod_10(x) - 10,
        "sb_file": "explorenew_SB.pdf",
        "mis_file": "explorenew_MIS.pdf",
        "sb_ylim": (0, 8),
        "bayes_ylim": None,
        "legend_loc": "best",
        "mis_legend_loc": [0.45, 0.7],
        "multiframe_shannon_ylim": (0, 100),
    },
}


def _initial_and_stream(scenario_key: str, num_steps: int):
    meta = SCENARIOS[scenario_key]
    np.random.seed(meta["seed"])
    x_init, y_init = meta["init"]()
    if y_init is None:
        y_init = mod_10(x_init)
    x_stream, y_stream = meta["stream"](num_steps)
    if y_stream is None:
        y_stream = meta["stream_y"](x_stream)
    return np.asarray(x_init), np.asarray(y_init), np.asarray(x_stream), np.asarray(y_stream)


def compute_single_step_logs(scenario_key: str, num_steps: int = 100):
    x_old, y_old, x_stream, y_stream = _initial_and_stream(scenario_key, num_steps)

    shannon_log = np.zeros(num_steps)
    bayesian_log = np.zeros(num_steps)
    mle_log = np.zeros(num_steps)

    for i, (x_new, y_new) in enumerate(zip(x_stream, y_stream)):
        gp = fit_gp(x_old, y_old, random_state=0)
        old_xy = np.column_stack((x_old, y_old))
        shannon_log[i] = shannon_surprise(old_xy, np.array([x_new, y_new]), gp=gp)
        bayesian_log[i] = bayesian_surprise(gp, np.array([x_new]), y_new)
        mle_log[i] = mle_entropy(x_old) + mle_entropy(y_old) - mle_entropy(old_xy)
        x_old = np.append(x_old, x_new)
        y_old = np.append(y_old, y_new)

    return shannon_log, bayesian_log, mle_log


def _mis_bounds(num_steps: int, n: int = 100, rho: float = 0.1):
    ind = np.arange(num_steps)
    upper = 9 * (1 / (ind + n) - 1 / n) + (
        np.sqrt(2 * np.log(2 / rho) * ind) * np.log(10) / (n + ind)
    )
    lower = 9 * (1 / (ind + n) - 1 / n) - (
        np.sqrt(2 * np.log(2 / rho) * ind) * np.log(10) / (n + ind)
    )
    return ind, upper, lower


def plot_single_step_panels(scenario_key: str, out_dir: Path, num_steps: int = 100):
    meta = SCENARIOS[scenario_key]
    shannon_log, bayesian_log, mle_log = compute_single_step_logs(scenario_key, num_steps)
    shannon_plot = np.asarray(shannon_log, dtype=float).copy()
    shannon_plot[np.isinf(shannon_plot)] = 5000
    ind, upper, lower = _mis_bounds(num_steps)

    fig1, ax1 = plt.subplots(figsize=(8, 6))
    ax2 = ax1.twinx()
    ax1.plot(ind, shannon_plot, label="Shannon Surprise", color="blue", linestyle="--")
    ax2.plot(ind, bayesian_log, label="Bayesian Surprise", color="red")
    ax1.set_xlabel(meta["xlabel"], fontsize=14)
    ax1.set_ylabel("Shannon Surprise", fontsize=14, color="blue")
    ax2.set_ylabel("Bayesian Surprise", fontsize=14, color="red")
    if meta.get("sb_ylim"):
        ax1.set_ylim(*meta["sb_ylim"])
    if meta.get("bayes_ylim"):
        ax2.set_ylim(*meta["bayes_ylim"])
    ax1.tick_params(axis="y", labelcolor="blue")
    ax2.tick_params(axis="y", labelcolor="red")
    ax1.tick_params(axis="x", labelsize=12)
    ax1.set_title("Shannon and Bayesian Surprises", fontsize=16)
    ax1.grid()
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, fontsize=14, loc=meta["legend_loc"])
    fig1.tight_layout()
    fig1.savefig(out_dir / meta["sb_file"])
    plt.close(fig1)

    fig2, ax3 = plt.subplots(figsize=(8, 6))
    ax3.plot(ind, mle_log - mle_log[0], label="Mutual Information Surprise", color="green")
    ax3.fill_between(ind, -upper, -lower, color="gray", alpha=0.5, label="MIS Bound")
    ax3.set_xlabel(meta["xlabel"], fontsize=14)
    ax3.set_ylabel("Mutual Information Surprise", fontsize=14)
    ax3.set_title("Mutual Information Surprise", fontsize=16)
    ax3.tick_params(axis="x", labelsize=12)
    ax3.grid()
    ax3.legend(fontsize=14, loc=meta["mis_legend_loc"])
    fig2.tight_layout()
    fig2.savefig(out_dir / meta["mis_file"])
    plt.close(fig2)


def kl_mvn(m0, s0, m1, s1, jitter=1e-8):
    m0 = np.atleast_1d(m0)
    m1 = np.atleast_1d(m1)
    k = m0.size
    s0 = np.atleast_2d(s0) + jitter * np.eye(k)
    s1 = np.atleast_2d(s1) + jitter * np.eye(k)
    l1 = np.linalg.cholesky(s1)
    a = np.linalg.solve(l1, s0)
    a = np.linalg.solve(l1.T, a)
    dm = (m1 - m0).reshape(-1, 1)
    v = np.linalg.solve(l1, dm)
    logdet_s0 = 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(s0))))
    logdet_s1 = 2.0 * np.sum(np.log(np.diag(l1)))
    return float(0.5 * (np.trace(a) + float(v.T @ v) - k + logdet_s1 - logdet_s0))


def multiframe_shannon(x_init, y_init, x_stream, y_stream):
    x_hist = np.asarray(x_init).copy()
    y_hist = np.asarray(y_init).copy()
    sh_step = []
    for x_new, y_new in zip(x_stream, y_stream):
        gp = fit_gp(x_hist, y_hist, n_restarts_optimizer=5, random_state=0)
        sh_step.append(shannon_surprise(np.column_stack((x_hist, y_hist)), [x_new, y_new], gp=gp))
        x_hist = np.append(x_hist, x_new)
        y_hist = np.append(y_hist, y_new)
    sh_step = np.asarray(sh_step, dtype=float)
    return sh_step, np.cumsum(sh_step)


def multiframe_shannon_moving_window(sh_step, window_size):
    out = np.full_like(sh_step, np.nan, dtype=float)
    for t in range(len(sh_step)):
        start = max(0, t - window_size + 1)
        out[t] = np.sum(sh_step[start : t + 1])
    return out


def bayesian_surprise_frame(gp_before, gp_after, x_eval):
    x_eval = np.asarray(x_eval).reshape(-1, 1)
    m0, s0 = gp_before.predict(x_eval, return_cov=True)
    m1, s1 = gp_after.predict(x_eval, return_cov=True)
    return kl_mvn(m0, s0, m1, s1)


def multiframe_bayesian_cumulative(x_init, y_init, x_stream, y_stream, x_eval):
    gp0 = fit_gp(x_init, y_init, n_restarts_optimizer=5, random_state=0)
    out = []
    x_hist = np.asarray(x_init).copy()
    y_hist = np.asarray(y_init).copy()
    for x_new, y_new in zip(x_stream, y_stream):
        x_hist = np.append(x_hist, x_new)
        y_hist = np.append(y_hist, y_new)
        out.append(
            bayesian_surprise_frame(
                gp0,
                fit_gp(x_hist, y_hist, n_restarts_optimizer=5, random_state=0),
                x_eval,
            )
        )
    return np.asarray(out)


def multiframe_bayesian_moving_window(x_init, y_init, x_stream, y_stream, x_eval, window_size):
    x_all = np.concatenate([x_init, x_stream])
    y_all = np.concatenate([y_init, y_stream])
    n0 = len(x_init)
    out = np.full(len(x_stream), np.nan, dtype=float)
    for t in range(len(x_stream)):
        end_idx = n0 + t + 1
        start_idx = n0 + max(0, t + 1 - window_size)
        gp_before = fit_gp(x_all[:start_idx], y_all[:start_idx], n_restarts_optimizer=5, random_state=0)
        gp_after = fit_gp(x_all[:end_idx], y_all[:end_idx], n_restarts_optimizer=5, random_state=0)
        out[t] = bayesian_surprise_frame(gp_before, gp_after, x_eval)
    return out


def plot_multiframe_sb(
    ind,
    shannon_vals,
    bayesian_vals,
    scenario_key,
    out_dir,
    mode,
    window_size=20,
    shannon_ylim=None,
):
    meta = SCENARIOS[scenario_key]
    shannon_vals = np.asarray(shannon_vals, dtype=float).copy()
    shannon_vals[np.isinf(shannon_vals)] = 5000
    bayesian_vals = np.asarray(bayesian_vals, dtype=float)
    mask = ind >= window_size if mode == "rolling" else np.ones_like(ind, dtype=bool)
    x = np.asarray(ind)[mask]
    sh = shannon_vals[mask]
    bay = bayesian_vals[mask]
    sh_inf = np.all(sh == 5000)

    fig, ax1 = plt.subplots(figsize=(8, 6))
    ax2 = ax1.twinx()

    if mode == "cumulative":
        filename = f"{scenario_key}_SB_cumulative.pdf"
        title = "Shannon and Bayesian Surprises (Cumulative)"
        sh_label = "Shannon Surprise"
        bay_label = "Bayesian Surprise"
    else:
        filename = f"{scenario_key}_SB_rolling.pdf"
        title = "Shannon and Bayesian Surprises (Moving Window)"
        sh_label = f"Shannon Surprise (W={window_size})"
        bay_label = f"Bayesian Surprise (W={window_size})"

    if not sh_inf:
        ax1.plot(x, sh, label=sh_label, color="blue", linestyle="--")
    ax2.plot(x, bay, label=bay_label, color="red")

    ax1.set_xlabel(meta["xlabel"], fontsize=14)
    ax1.set_ylabel("Shannon Surprise" if not sh_inf else "Shannon Surprise (Infinite)", fontsize=14, color="blue")
    ax2.set_ylabel("Bayesian Surprise", fontsize=14, color="red")
    if shannon_ylim is not None:
        ax1.set_ylim(*shannon_ylim)
    ax1.tick_params(axis="y", labelcolor="blue")
    ax2.tick_params(axis="y", labelcolor="red")
    ax1.tick_params(axis="x", labelsize=12)
    ax1.set_title(title, fontsize=16)
    ax1.grid()
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, fontsize=14, loc="upper left")
    fig.tight_layout()
    fig.savefig(out_dir / filename)
    plt.close(fig)


def plot_multiframe_panels(scenario_key: str, out_dir: Path, num_steps: int = 100, window_size: int = 20):
    x_init, y_init, x_stream, y_stream = _initial_and_stream(scenario_key, num_steps)
    x_eval_count = 200 if scenario_key != "pure_explore" else 400
    x_eval = np.linspace(
        min(float(np.min(x_init)), float(np.min(x_stream))),
        max(float(np.max(x_init)), float(np.max(x_stream))),
        x_eval_count,
    )
    sh_step, sh_cum = multiframe_shannon(x_init, y_init, x_stream, y_stream)
    sh_roll = multiframe_shannon_moving_window(sh_step, window_size)
    bay_cum = multiframe_bayesian_cumulative(x_init, y_init, x_stream, y_stream, x_eval)
    bay_roll = multiframe_bayesian_moving_window(x_init, y_init, x_stream, y_stream, x_eval, window_size)
    ind = np.arange(1, num_steps + 1)
    sh_ylim = SCENARIOS[scenario_key].get("multiframe_shannon_ylim")
    plot_multiframe_sb(ind, sh_cum, bay_cum, scenario_key, out_dir, "cumulative", window_size, sh_ylim)
    plot_multiframe_sb(ind, sh_roll, bay_roll, scenario_key, out_dir, "rolling", window_size, sh_ylim)


def generate_discrete_figure(figure_number: int, scenario_key: str):
    out_dir = ensure_dir(OUTPUT_DIR / f"figure{figure_number}")
    plot_single_step_panels(scenario_key, out_dir)
    plot_multiframe_panels(scenario_key, out_dir)
    print(f"Saved Figure {figure_number} files to {out_dir}")
