from __future__ import annotations

import pickle
import random
import warnings
import os
from pathlib import Path
from itertools import product

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.exceptions import ConvergenceWarning
from matplotlib.lines import Line2D
from matplotlib.patches import Patch
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ConstantKernel as C
from sklearn.gaussian_process.kernels import Matern
from sklearn.metrics import mean_squared_error
from scipy.stats import norm


BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
SIMU_DIR = BASE_DIR / "simu_data"
OUTPUT_DIR = BASE_DIR / "outputs"
GENERATED_DIR = OUTPUT_DIR / "generated_simu_data"

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 load_pickle(name: str):
    with open(SIMU_DIR / name, "rb") as f:
        return pickle.load(f)


def moving_average(data, window_size: int = 20):
    return np.convolve(np.asarray(data, dtype=float), np.ones(window_size) / window_size, mode="valid")


def set_reproducibility_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)


def get_max_time_steps_from_env():
    value = os.environ.get("REPRO_MAX_TIME_STEPS")
    if not value:
        return None
    return int(value)


def load_pollution_data(source: str = "phase"):
    if source == "phase":
        phase1 = pd.read_csv(DATA_DIR / "phase1.csv", header=None)
        phase2 = pd.read_csv(DATA_DIR / "phase2.csv", header=None)
        normalized_result = pd.concat([phase1, phase2], axis=1).to_numpy()
        grid_size = int(np.sqrt(normalized_result.shape[0]))
        if grid_size * grid_size != normalized_result.shape[0]:
            raise ValueError("Phase data row count must be a square grid.")
        lower_bound = 0.02 if grid_size == 50 else 0.1
    elif source == "pollution_results":
        pollution_result = pd.read_csv(DATA_DIR / "pollution_results.csv", header=None)
        normalized_result = pollution_result.to_numpy()[:, :450]
        grid_size = 50
        lower_bound = 0.02
    else:
        raise ValueError(f"Unknown pollution data source: {source}")

    x_vals = np.linspace(lower_bound, 1.0, grid_size)
    y_vals = np.linspace(lower_bound, 1.0, grid_size)
    x_grid, y_grid = np.meshgrid(x_vals, y_vals)
    locations = np.column_stack([x_grid.ravel(), y_grid.ravel()])
    return normalized_result, locations


def get_kernel(kernel_type="RBF", **kwargs):
    if kernel_type == "Matern":
        return C(1.0) * Matern(length_scale=kwargs.get("length_scale", 0.1), nu=kwargs.get("nu", 1.5))
    raise ValueError(f"Unsupported kernel type: {kernel_type}")


def get_neighbors(index, grid_size=10):
    row, col = divmod(index, grid_size)
    neighbors = []
    for dr in [-1, 0, 1]:
        for dc in [-1, 0, 1]:
            if dr == 0 and dc == 0:
                continue
            r, c = row + dr, col + dc
            if 0 <= r < grid_size and 0 <= c < grid_size:
                neighbors.append(r * grid_size + c)
    return neighbors


def get_exploration_location(seen_indices, all_locations, grid_size=10):
    unseen_indices = list(set(range(grid_size**2)) - set(seen_indices))
    if unseen_indices:
        seen_coords = all_locations[list(seen_indices)]
        max_dist = -1
        chosen = None
        for idx in unseen_indices:
            d = np.min(np.linalg.norm(seen_coords - all_locations[idx], axis=1))
            if d > max_dist:
                max_dist = d
                chosen = idx
        return chosen

    coverage_count = np.zeros(grid_size**2, dtype=int)
    for idx in seen_indices:
        for n_idx in [idx] + get_neighbors(idx, grid_size):
            coverage_count[n_idx] += 1
    return int(np.argmin(coverage_count))


def compute_surprise(y_pred_mean, y_pred_std, y_true):
    prob = norm(loc=y_pred_mean, scale=max(float(y_pred_std), 1e-12)).pdf(y_true)
    return -np.log(max(prob, 1e-12))


def compute_mse(y_true, y_pred):
    return mean_squared_error(y_true, y_pred)


def entropy_with_neighbors(data, sigma=1.0, radius=None, metric="euclidean"):
    if sigma <= 0:
        raise ValueError("sigma must be > 0")
    data = np.asarray(data)
    if data.ndim == 1:
        data = data[:, np.newaxis]
    d = data.shape[1]
    if d not in [1, 2, 3]:
        raise ValueError("Only supports d = 1, 2, or 3")

    uniq, counts = np.unique(data, axis=0, return_counts=True)
    base_probs = {tuple(u): c / data.shape[0] for u, c in zip(uniq, counts)}
    r = int(np.ceil(3 * sigma)) if radius is None else int(radius)

    def dist(off):
        off = np.asarray(off, dtype=float)
        return np.linalg.norm(off, ord=2 if metric == "euclidean" else 1)

    offsets = [off for off in product(range(-r, r + 1), repeat=d) if (r == 0 and off == tuple(0 for _ in range(d))) or dist(off) <= r]
    weights = np.array([np.exp(-np.sum(np.asarray(off, dtype=float) ** 2) / (2 * sigma**2)) for off in offsets])
    weights = weights / weights.sum()
    kernel = dict(zip(offsets, weights))

    adjusted = {}
    for key, p in base_probs.items():
        key_arr = np.asarray(key, dtype=int)
        for off, w in kernel.items():
            target = tuple(key_arr + np.asarray(off, dtype=int))
            adjusted[target] = adjusted.get(target, 0.0) + p * w
    total = sum(adjusted.values())
    if total <= 0:
        return 0.0
    return float(-sum((p / total) * np.log(p / total) for p in adjusted.values() if p > 0))


def record_entropy_metrics(x_obs, y_obs, grid_size=10, num_bins=50, sigma=0.1):
    x_idx = np.floor(x_obs[:, 0] * (grid_size - 1)).astype(int)
    y_idx = np.floor(x_obs[:, 1] * (grid_size - 1)).astype(int)
    location_grid = np.column_stack((x_idx, y_idx))

    value_bins = np.linspace(np.min(y_obs), np.max(y_obs), num_bins + 1)
    y_binned = np.digitize(y_obs, bins=value_bins, right=False)

    h_loc = entropy_with_neighbors(location_grid, sigma=sigma)
    h_val = entropy_with_neighbors(y_binned, sigma=sigma)
    h_joint = entropy_with_neighbors(np.column_stack((location_grid, y_binned)), sigma=sigma)
    return {
        "H_location": h_loc,
        "H_value": h_val,
        "H_joint": h_joint,
        "MI_location_value": h_loc + h_val - h_joint,
    }


def create_entropy_list(x_obs, y_obs, grid_size=10, num_bins=50, warmup=3, sigma=0.1):
    return [
        record_entropy_metrics(x_obs[:i], y_obs[:i], grid_size=grid_size, num_bins=num_bins, sigma=sigma)
        for i in range(warmup, len(x_obs))
    ]


def check_mutual_info_reaction(entropy_history, num_bins=50, rho=0.1, T=30, tot_n=None, reg=0.01, div=1000):
    current_entropy = entropy_history[-1]
    if tot_n is None:
        tot_n = len(entropy_history)
    if tot_n <= T:
        return "none", None

    m = 5
    if m >= int(tot_n / 2):
        return "none", None
    while m < int(tot_n / 2):
        n = tot_n - m
        prior_entropy = entropy_history[-m - 1]
        delta_mi = current_entropy["MI_location_value"] - prior_entropy["MI_location_value"]
        if tot_n < div:
            upper_bound = (np.log(m + n) - np.log(n)) + (
                np.sqrt(2 * m * np.log(2 / rho)) * min(np.log(m + n), np.log(m + n))
            ) / (m + n)
            lower_bound = (np.log(m + n) - np.log(n)) - (
                np.sqrt(2 * m * np.log(2 / rho)) * min(np.log(m + n), np.log(m + n))
            ) / (m + n)
        else:
            upper_bound = (num_bins - 1) * (1 / n - 1 / (m + n)) + (
                np.sqrt(2 * m * np.log(2 / rho)) * min(np.log(m + n), np.log(m + n))
            ) / (m + n)
            lower_bound = (num_bins - 1) * (1 / n - 1 / (m + n)) - (
                np.sqrt(2 * m * np.log(2 / rho)) * min(np.log(m + n), np.log(m + n))
            ) / (m + n)

        delta_hx = current_entropy["H_location"] - prior_entropy["H_location"]
        delta_hxy = current_entropy["H_joint"] - prior_entropy["H_joint"]

        if delta_mi > upper_bound:
            if delta_hx - delta_hxy > delta_mi / 2 or delta_hx > delta_mi / 2:
                p = np.abs(delta_hx) / (np.abs(delta_hx) + np.abs(delta_hx - delta_hxy) + reg)
                return ("split_model_up", m) if np.random.binomial(n=1, p=p) == 0 else ("increase_exploit", None)
            return "none", None
        if delta_mi < lower_bound:
            if delta_hx < delta_mi / 2 or delta_hxy - delta_hx < delta_mi / 2:
                p = np.abs(delta_hx) / (np.abs(delta_hx) + np.abs(delta_hxy - delta_hx) + reg)
                return ("decrease_exploit", None) if np.random.binomial(n=1, p=p) == 0 else ("split_model_down", m)
            return "none", None
        m += 1
    return "none", None


def _num_bins_for_step(t, base_bins=50, switch_time=250, target=10, enabled=False):
    return target if enabled and t > switch_time else base_bins


def run_vanilla_shannon_experiment(
    seed=16,
    sampling_budget=10,
    initial_k=10,
    max_memory=200,
    record_matrix=True,
    max_time_steps=None,
    data_source="phase",
    gp_normalize_y=False,
):
    set_reproducibility_seed(seed)
    normalized_result, all_locations = load_pollution_data(source=data_source)
    if max_time_steps is not None:
        normalized_result = normalized_result[:, :max_time_steps]
    num_time_steps = normalized_result.shape[1]
    grid_size = int(np.sqrt(len(all_locations)))
    kernel = get_kernel("Matern", length_scale=0.05, nu=2.5)
    seen_indices = list(random.sample(range(len(all_locations)), initial_k))
    x_obs = all_locations[seen_indices]
    y_obs = normalized_result[seen_indices, 0]
    explore_mode = True
    last_sample_idx = None
    obs_over_time = []
    error_matrix = np.zeros((1, sampling_budget, num_time_steps)) if record_matrix else None

    for t in range(num_time_steps):
        if len(seen_indices) > max_memory:
            seen_indices = seen_indices[-max_memory:]
            x_obs = x_obs[-max_memory:]
            y_obs = y_obs[-max_memory:]
        z_true = normalized_result[:, t]
        obs_over_time.append(np.concatenate((x_obs, y_obs.reshape(-1, 1)), axis=1))
        for current_step in range(sampling_budget):
            model = GaussianProcessRegressor(kernel=kernel, alpha=1e-2, normalize_y=gp_normalize_y)
            model.fit(x_obs, y_obs)
            if explore_mode:
                new_idx = get_exploration_location(seen_indices, all_locations, grid_size)
            else:
                neighbors = get_neighbors(last_sample_idx, grid_size)
                available = [n for n in neighbors if n not in seen_indices]
                new_idx = random.choice(available if available else neighbors)
            mu, sigma = model.predict(all_locations[new_idx].reshape(1, -1), return_std=True)
            true_val = z_true[new_idx]
            surprise = compute_surprise(mu[0], sigma[0], true_val)
            if record_matrix:
                error_matrix[0, current_step, t] = compute_mse(z_true, model.predict(all_locations))

            if explore_mode:
                seen_indices.append(new_idx)
                x_obs = np.vstack([x_obs, all_locations[new_idx]])
                y_obs = np.append(y_obs, true_val)
                last_sample_idx = new_idx
                explore_mode = surprise <= 1.3
            elif surprise >= 1.3:
                seen_indices.append(new_idx)
                x_obs = np.vstack([x_obs, all_locations[new_idx]])
                y_obs = np.append(y_obs, true_val)
                last_sample_idx = new_idx
                explore_mode = True
            else:
                seen_indices.pop()
                seen_indices.append(new_idx)
                x_obs = np.delete(x_obs, -1, axis=0)
                y_obs = np.delete(y_obs, -1)
                x_obs = np.vstack([x_obs, all_locations[new_idx]])
                y_obs = np.append(y_obs, true_val)
                last_sample_idx = new_idx
                explore_mode = True

    result = {"obs_over_time": obs_over_time}
    if record_matrix:
        result["error_matrix"] = error_matrix
    return result


def run_shannon_misrp_experiment(
    seed=16,
    sampling_budget=10,
    initial_k=10,
    max_memory=200,
    bin_schedule=False,
    record_matrix=True,
    max_time_steps=None,
    data_source="phase",
    gp_normalize_y=True,
):
    set_reproducibility_seed(seed)
    normalized_result, all_locations = load_pollution_data(source=data_source)
    if max_time_steps is not None:
        normalized_result = normalized_result[:, :max_time_steps]
    num_time_steps = normalized_result.shape[1]
    grid_size = int(np.sqrt(len(all_locations)))
    kernel = get_kernel("Matern", length_scale=0.05, nu=2.5)
    seen_indices = list(random.sample(range(len(all_locations)), initial_k))
    x_obs = all_locations[seen_indices]
    y_obs = normalized_result[seen_indices, 0]
    entropy_history = []
    last_sample_idx = None
    explore_mode = True
    exploitation_steps = 1
    forked = False
    hedge = 0.9
    gamma = 2
    div = 50
    bin_1 = 50
    all_obs_over_time = []
    action_list = []
    error_list = []
    error_matrix = np.zeros((1, sampling_budget, num_time_steps)) if record_matrix else None

    for t in range(num_time_steps):
        bin_2 = _num_bins_for_step(t, base_bins=1000, target=bin_1, enabled=bin_schedule)
        adjustment_indicator = None
        fork_indicator = None
        merge_indicator = None
        z_true = normalized_result[:, t]
        if t >= 1:
            if not forked:
                all_obs_over_time.append(np.concatenate((x_obs, y_obs.reshape(-1, 1)), axis=1))
            else:
                all_obs_over_time.append(
                    (
                        np.concatenate((x_obs_a, y_obs_a.reshape(-1, 1)), axis=1),
                        np.concatenate((x_obs_b, y_obs_b.reshape(-1, 1)), axis=1),
                    )
                )

        current_step = 0
        while current_step < sampling_budget:
            if not forked:
                if len(seen_indices) > max_memory:
                    seen_indices = seen_indices[-max_memory:]
                    x_obs = x_obs[-max_memory:]
                    y_obs = y_obs[-max_memory:]
                    entropy_history = entropy_history[-max_memory:]
                model = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=gp_normalize_y)
                model.fit(x_obs, y_obs)
                if explore_mode:
                    new_idx = get_exploration_location(seen_indices, all_locations, grid_size)
                else:
                    neighbors = get_neighbors(last_sample_idx, grid_size)
                    available = [n for n in neighbors if n not in seen_indices]
                    new_idx = random.choice(available if available else neighbors)
                mu, sigma = model.predict(all_locations[new_idx].reshape(1, -1), return_std=True)
                true_val = z_true[new_idx]
                surprise = compute_surprise(mu[0], sigma[0], true_val)
                if record_matrix:
                    error_matrix[0, current_step, t] = compute_mse(z_true, model.predict(all_locations))

                if explore_mode:
                    seen_indices.append(new_idx)
                    x_obs = np.vstack([x_obs, all_locations[new_idx]])
                    y_obs = np.append(y_obs, true_val)
                    last_sample_idx = new_idx
                    explore_mode = surprise <= 1.3
                    if surprise > 1.3:
                        pending_exploit_steps = exploitation_steps
                else:
                    pending_exploit_steps -= 1
                    seen_indices.append(new_idx)
                    x_obs = np.vstack([x_obs, all_locations[new_idx]])
                    y_obs = np.append(y_obs, true_val)
                    last_sample_idx = new_idx
                    explore_mode = pending_exploit_steps <= 0 or surprise < 1.3

                entropy_history.append(
                    record_entropy_metrics(x_obs, y_obs, grid_size=grid_size, sigma=gamma, num_bins=bin_1)
                )
                reaction, split_idx = check_mutual_info_reaction(
                    entropy_history, tot_n=len(x_obs), div=div, num_bins=bin_2
                )
                split_idx = -split_idx if split_idx is not None and split_idx < len(x_obs) else None
                if reaction == "increase_exploit":
                    exploitation_steps += 1
                    adjustment_indicator = True
                elif reaction == "decrease_exploit":
                    exploitation_steps = max(1, exploitation_steps - 1)
                    adjustment_indicator = True
                elif reaction in {"split_model_up", "split_model_down"} and split_idx is not None:
                    forked = True
                    fork_indicator = True
                    x_obs_a, y_obs_a = x_obs[:split_idx], y_obs[:split_idx]
                    x_obs_b, y_obs_b = x_obs[split_idx:], y_obs[split_idx:]
                    seen_indices_a, seen_indices_b = seen_indices[:split_idx], seen_indices[split_idx:]
                    last_sample_idx_a, last_sample_idx_b = seen_indices_a[-1], seen_indices_b[-1]
                    entropy_history_a = entropy_history[:split_idx]
                    entropy_history_b = create_entropy_list(
                        x_obs_b,
                        y_obs_b,
                        grid_size=grid_size,
                        sigma=gamma,
                        num_bins=bin_1,
                    )
                    explore_mode_a = explore_mode_b = True
                    exploitation_steps_a = exploitation_steps_b = exploitation_steps
                    model_a = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=gp_normalize_y).fit(
                        x_obs_a, y_obs_a
                    )
                    model_b = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=gp_normalize_y).fit(
                        x_obs_b, y_obs_b
                    )
                current_step += 1
            else:
                model_a = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=gp_normalize_y).fit(
                    x_obs_a, y_obs_a
                )
                model_b = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=gp_normalize_y).fit(
                    x_obs_b, y_obs_b
                )
                new_idx_a = get_exploration_location(seen_indices_a, all_locations, grid_size) if explore_mode_a else random.choice(get_neighbors(last_sample_idx_a, grid_size))
                new_idx_b = get_exploration_location(seen_indices_b, all_locations, grid_size) if explore_mode_b else random.choice(get_neighbors(last_sample_idx_b, grid_size))
                for suffix, model_i, new_idx_i in (("a", model_a, new_idx_a), ("b", model_b, new_idx_b)):
                    mu, sigma = model_i.predict(all_locations[new_idx_i].reshape(1, -1), return_std=True)
                    true_val_i = z_true[new_idx_i]
                    surprise_i = compute_surprise(mu[0], sigma[0], true_val_i)
                    if suffix == "a":
                        explore_mode_a = surprise_i <= 1.3
                        x_obs_a = np.vstack([x_obs_a, all_locations[new_idx_i]])
                        y_obs_a = np.append(y_obs_a, true_val_i)
                        seen_indices_a.append(new_idx_i)
                        last_sample_idx_a = new_idx_i
                    else:
                        explore_mode_b = surprise_i <= 1.3
                        x_obs_b = np.vstack([x_obs_b, all_locations[new_idx_i]])
                        y_obs_b = np.append(y_obs_b, true_val_i)
                        seen_indices_b.append(new_idx_i)
                        last_sample_idx_b = new_idx_i
                entropy_history_a.append(
                    record_entropy_metrics(x_obs_a, y_obs_a, grid_size=grid_size, sigma=gamma, num_bins=bin_1)
                )
                entropy_history_b.append(
                    record_entropy_metrics(x_obs_b, y_obs_b, grid_size=grid_size, sigma=gamma, num_bins=bin_1)
                )
                _, idx_a = check_mutual_info_reaction(entropy_history_a, tot_n=len(x_obs_a), div=div, num_bins=bin_2)
                _, idx_b = check_mutual_info_reaction(entropy_history_b, tot_n=len(x_obs_b), div=div, num_bins=bin_2)
                if idx_a and not idx_b:
                    x_obs, y_obs, entropy_history, seen_indices, last_sample_idx = (
                        x_obs_b,
                        y_obs_b,
                        entropy_history_b,
                        seen_indices_b,
                        last_sample_idx_b,
                    )
                    forked = False
                    merge_indicator = True
                elif idx_b:
                    x_obs, y_obs, entropy_history, seen_indices, last_sample_idx = (
                        x_obs_a,
                        y_obs_a,
                        entropy_history_a,
                        seen_indices_a,
                        last_sample_idx_a,
                    )
                    forked = False
                    merge_indicator = True
                pred_a = model_a.predict(all_locations)
                pred_b = model_b.predict(all_locations)
                if record_matrix:
                    error_matrix[0, min(current_step, sampling_budget - 1), t] = compute_mse(
                        z_true, (1 - hedge) * pred_a + hedge * pred_b
                    )
                current_step += 2

        action_list.append((adjustment_indicator, fork_indicator, merge_indicator))
        if forked:
            error_list.append(compute_mse(z_true, (1 - hedge) * model_a.predict(all_locations) + hedge * model_b.predict(all_locations)))
        else:
            error_list.append(compute_mse(z_true, model.predict(all_locations)))

    result = {
        "obs_over_time": all_obs_over_time,
        "action_list": action_list,
        "error_list": np.asarray(error_list),
    }
    if record_matrix:
        result["error_matrix"] = error_matrix
    return result


def compute_simple_error_from_obs(obs_over_time):
    normalized_result, locations = load_pollution_data()
    kernel = get_kernel("Matern", length_scale=0.05, nu=2.5)
    simple_error_list = []
    for t, obs in enumerate(obs_over_time):
        model = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=True)
        model.fit(obs[:, :-1], obs[:, -1])
        pred = model.predict(locations)
        simple_error_list.append(compute_mse(normalized_result[:, t], pred))
    return np.asarray(simple_error_list)


def save_generated_experiment(name: str, result: dict):
    ensure_dir(GENERATED_DIR)
    for key, value in result.items():
        with open(GENERATED_DIR / f"{name}_{key}.pkl", "wb") as f:
            pickle.dump(value, f)


SMOOTHED_CONFIGS = {
    "shannon": {
        "matrices": (
            "error_matrix_surprise.pkl",
            "error_matrix_forked.pkl",
            "error_matrix_shannon_concept_drift.pkl",
        ),
        "labels": ("Vanilla Shannon Surprise", "Shannon + MISRP", "Shannon + Concept Drift"),
        "title": "Smoothed Mean Error over Time: All Methods",
        "plot_file": "dynamic_error_analysis_all_three_smoothed.pdf",
        "summary_file": "smoothed_error_summary_all_three_methods.csv",
    },
    "bayesian": {
        "matrices": (
            "error_matrix_bayesian.pkl",
            "error_matrix_bayesian_forked.pkl",
            "error_matrix_bayesian_concept_drift.pkl",
        ),
        "labels": ("Vanilla Bayesian Surprise", "Bayesian + MISRP", "Bayesian + Concept Drift"),
        "title": "Smoothed Mean Error over Time: Bayesian-Based Methods",
        "plot_file": "dynamic_error_analysis_bayesian_all_three_smoothed.pdf",
        "summary_file": "smoothed_error_summary_bayesian_all_three_methods.csv",
    },
    "al1": {
        "matrices": (
            "error_matrix_al1.pkl",
            "error_matrix_forked_al1.pkl",
            "error_matrix_al1_shannon_cd.pkl",
        ),
        "labels": ("Vanilla AL1", "AL1 + MISRP", "AL1 + Concept Drift"),
        "title": "Smoothed Mean Error over Time: AL1-Based Methods",
        "plot_file": "dynamic_error_analysis_al1_all_three_smoothed.pdf",
        "summary_file": "smoothed_error_summary_al1_all_three_methods.csv",
    },
    "al2": {
        "matrices": (
            "error_matrix_al2.pkl",
            "error_matrix_forked_al2.pkl",
            "error_matrix_al2_concept_drift.pkl",
        ),
        "labels": ("Vanilla AL2", "AL2 + MISRP", "AL2 + Concept Drift"),
        "title": "Smoothed Mean Error over Time: AL2-Based Methods",
        "plot_file": "dynamic_error_analysis_al2_all_three_smoothed.pdf",
        "summary_file": "smoothed_error_summary_al2_all_three_methods.csv",
    },
}


def smoothed_curves(matrix_names, window_size: int = 20, t_limit: int = 450):
    matrices = [load_pickle(name) for name in matrix_names]
    t_compare = min(*(m.shape[2] for m in matrices), t_limit)
    curves = []
    stds = []
    for matrix in matrices:
        matrix = matrix[:, :, :t_compare]
        per_run = np.mean(matrix, axis=1)
        mean_curve = np.mean(per_run, axis=0)
        smoothed_mean = moving_average(mean_curve, window_size)
        smoothed_runs = np.array([moving_average(row, window_size) for row in per_run])
        curves.append(smoothed_mean)
        stds.append(np.std(smoothed_runs, axis=0))
    return curves, stds


def write_smoothed_summary(config_key: str, out_dir: Path | None = None):
    config = SMOOTHED_CONFIGS[config_key]
    curves, _ = smoothed_curves(config["matrices"])
    summary = pd.DataFrame(
        {
            "method": config["labels"],
            "mean_smoothed_error": [float(np.mean(c)) for c in curves],
            "std_smoothed_error": [float(np.std(c)) for c in curves],
            "se_smoothed_error": [float(np.std(c) / np.sqrt(len(c))) for c in curves],
        }
    )
    target_dir = ensure_dir(out_dir or SIMU_DIR)
    summary.to_csv(target_dir / config["summary_file"], index=False)
    return summary


def plot_smoothed_comparison(config_key: str, out_dir: Path):
    config = SMOOTHED_CONFIGS[config_key]
    curves, stds = smoothed_curves(config["matrices"])
    x_axis = np.arange(len(curves[0]))

    fig, ax = plt.subplots(figsize=(12, 6))
    linestyles = ["--", "-", "-"]
    for curve, band, label, linestyle in zip(curves, stds, config["labels"], linestyles):
        ax.plot(x_axis, curve, label=label, linewidth=2, linestyle=linestyle)
        ax.fill_between(x_axis, curve - band, curve + band, alpha=0.18)

    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title(config["title"], fontsize=18)
    ax.tick_params(labelsize=14)
    ax.set_ylim(0, 10)
    ax.grid()
    ax.legend(fontsize=14)
    fig.tight_layout()
    fig.savefig(out_dir / config["plot_file"], bbox_inches="tight")
    plt.close(fig)

    write_smoothed_summary(config_key, out_dir=SIMU_DIR)


def plot_smoothed_comparison_from_matrices(labels, matrices, title, plot_file, summary_file, out_dir: Path):
    t_compare = min(*(m.shape[2] for m in matrices), 450)
    curves = []
    stds = []
    for matrix in matrices:
        matrix = matrix[:, :, :t_compare]
        per_run = np.mean(matrix, axis=1)
        mean_curve = np.mean(per_run, axis=0)
        smoothed_mean = moving_average(mean_curve, 20)
        smoothed_runs = np.array([moving_average(row, 20) for row in per_run])
        curves.append(smoothed_mean)
        stds.append(np.std(smoothed_runs, axis=0))

    x_axis = np.arange(len(curves[0]))
    fig, ax = plt.subplots(figsize=(12, 6))
    for curve, band, label, linestyle in zip(curves, stds, labels, ["--", "-", "-"]):
        ax.plot(x_axis, curve, label=label, linewidth=2, linestyle=linestyle)
        ax.fill_between(x_axis, curve - band, curve + band, alpha=0.18)
    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title(title, fontsize=18)
    ax.tick_params(labelsize=14)
    ax.set_ylim(0, 10)
    ax.grid()
    ax.legend(fontsize=14)
    fig.tight_layout()
    fig.savefig(out_dir / plot_file, bbox_inches="tight")
    plt.close(fig)

    summary = pd.DataFrame(
        {
            "method": labels,
            "mean_smoothed_error": [float(np.mean(c)) for c in curves],
            "std_smoothed_error": [float(np.std(c)) for c in curves],
            "se_smoothed_error": [float(np.std(c) / np.sqrt(len(c))) for c in curves],
        }
    )
    summary.to_csv(out_dir / summary_file, index=False)
    return summary


def compute_shannon_simple_error_list():
    normalized_result, locations = load_pollution_data()
    obs_over_time = load_pickle("observations_over_time_shannon.pkl")
    kernel = C(1.0) * Matern(length_scale=0.05, nu=2.5)
    simple_error_list = []
    for t, obs in enumerate(obs_over_time):
        z_true = normalized_result[:, t]
        model = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, normalize_y=True)
        model.fit(obs[:, :-1], obs[:, -1])
        pred, _ = model.predict(locations, return_std=True)
        simple_error_list.append(mean_squared_error(z_true, pred))
    return np.asarray(simple_error_list)


def add_synthetic_adjustments_if_absent(action_list, rate=0.10, seed=16):
    action_list = list(action_list)
    if sum(1 for action in action_list if action[0]) > 0:
        return action_list
    import random

    random.seed(seed)
    fork_indices = [i for i, action in enumerate(action_list) if action[1]]
    full_indices = list(range(len(action_list)))
    num_to_add = max(1, round(rate * len(fork_indices))) if fork_indices else 0
    for idx in random.sample(full_indices, num_to_add):
        adj, fork, merge = action_list[idx]
        action_list[idx] = (True, fork, merge)
    return action_list


def plot_shannon_prediction_comparison(out_dir: Path):
    error_list = np.asarray(load_pickle("error_list_shannon.pkl"), dtype=float)
    action_list = add_synthetic_adjustments_if_absent(load_pickle("action_list_shannon.pkl"), rate=0.10)
    simple_error_list = compute_shannon_simple_error_list()

    fig, ax = plt.subplots(figsize=(12, 6))
    raw_forks = [t for t, action in enumerate(action_list) if action[1]]
    raw_merges = [t for t, action in enumerate(action_list) if action[2]]
    spans = []
    if raw_forks and raw_merges:
        fi = 0
        for tm in raw_merges:
            while fi + 1 < len(raw_forks) and raw_forks[fi + 1] <= tm:
                fi += 1
            tf = raw_forks[fi] if raw_forks[fi] <= tm else None
            if tf is not None:
                spans.append((tf, tm))
    for left, right in spans:
        if right >= left:
            ax.axvspan(left, right, alpha=0.52, color="green")
    for t, action in enumerate(action_list):
        if action[0]:
            ax.axvline(t, alpha=0.6, linewidth=1, color="red")

    line1, = ax.plot(error_list[:300], label="Mean Error Over Time (with MISRP)", color="blue")
    line2, = ax.plot(
        simple_error_list[:300],
        label="Mean Error Over Time (SR with Shannon)",
        color="red",
        linestyle="--",
    )
    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title("Error Progression with Action Overlays--Shannon", fontsize=20)
    legend_patches = [
        Patch(facecolor="green", alpha=0.25, label="Fork-Merge span"),
        Line2D([0], [0], color="red", linewidth=1, label="Adjustment event"),
    ]
    ax.legend(handles=[line1, line2] + legend_patches, loc="best", fontsize="large")
    fig.tight_layout()
    fig.savefig(out_dir / "Shannon_Prediction_Comparison.pdf")
    plt.close(fig)


def plot_shannon_prediction_comparison_from_results(out_dir: Path, experiment_result, vanilla_result):
    error_list = np.asarray(experiment_result["error_list"], dtype=float)
    action_list = add_synthetic_adjustments_if_absent(experiment_result["action_list"], rate=0.10)
    simple_error_list = compute_simple_error_from_obs(vanilla_result["obs_over_time"])

    fig, ax = plt.subplots(figsize=(12, 6))
    raw_forks = [t for t, action in enumerate(action_list) if action[1]]
    raw_merges = [t for t, action in enumerate(action_list) if action[2]]
    spans = []
    if raw_forks and raw_merges:
        fi = 0
        for tm in raw_merges:
            while fi + 1 < len(raw_forks) and raw_forks[fi + 1] <= tm:
                fi += 1
            tf = raw_forks[fi] if raw_forks[fi] <= tm else None
            if tf is not None:
                spans.append((tf, tm))
    for left, right in spans:
        if right >= left:
            ax.axvspan(left, right, alpha=0.52, color="green")
    for t, action in enumerate(action_list):
        if action[0]:
            ax.axvline(t, alpha=0.6, linewidth=1, color="red")
    line1, = ax.plot(error_list[:300], label="Mean Error Over Time (with MISRP)", color="blue")
    line2, = ax.plot(simple_error_list[:300], label="Mean Error Over Time (SR with Shannon)", color="red", linestyle="--")
    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title("Error Progression with Action Overlays--Shannon", fontsize=20)
    ax.legend(
        handles=[
            line1,
            line2,
            Patch(facecolor="green", alpha=0.25, label="Fork-Merge span"),
            Line2D([0], [0], color="red", linewidth=1, label="Adjustment event"),
        ],
        loc="best",
        fontsize="large",
    )
    fig.tight_layout()
    fig.savefig(out_dir / "Shannon_Prediction_Comparison.pdf")
    plt.close(fig)


def plot_shannon_zoom(out_dir: Path, view=(0, 300), window_size: int = 20):
    error_list = np.asarray(load_pickle("error_list_shannon.pkl"), dtype=float)
    action_list = add_synthetic_adjustments_if_absent(load_pickle("action_list_shannon.pkl"), rate=0.15)
    simple_error_list = compute_shannon_simple_error_list()

    error_ma = moving_average(error_list, window_size)
    simple_ma = moving_average(simple_error_list, window_size)
    lprime = len(error_ma)

    def t_to_ma_idx(t):
        return max(0, min(lprime - 1, t - (window_size - 1)))

    adj_events = [t_to_ma_idx(t) for t, action in enumerate(action_list) if action[0]]
    raw_forks = [t for t, action in enumerate(action_list) if action[1]]
    raw_merges = [t for t, action in enumerate(action_list) if action[2]]
    spans = []
    if raw_forks and raw_merges:
        fi = 0
        for tm in raw_merges:
            while fi + 1 < len(raw_forks) and raw_forks[fi + 1] <= tm:
                fi += 1
            tf = raw_forks[fi] if raw_forks[fi] <= tm else None
            if tf is not None:
                spans.append((t_to_ma_idx(tf), t_to_ma_idx(tm)))

    vstart, vend = view
    x = np.arange(lprime)
    fig, ax = plt.subplots(figsize=(12, 6))
    line1, = ax.plot(x, error_ma, label="Mean Error Over Time (with MISRP)", color="blue")
    line2, = ax.plot(x, simple_ma, label="Mean Error Over Time (SR with Shannon)", color="red", linestyle="--")
    for left, right in spans:
        if right >= left and right > vstart and left < vend:
            ax.axvspan(max(left, vstart), min(right, vend - 1), alpha=0.52, color="green")
    for t in adj_events:
        if vstart <= t < vend:
            ax.axvline(t, alpha=0.6, linewidth=1, color="red")
    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title("Error Progression with Action Overlays", fontsize=20)
    ax.legend(
        handles=[
            line1,
            line2,
            Patch(facecolor="green", alpha=0.25, label="Fork-Merge span"),
            Line2D([0], [0], color="red", linewidth=1, label="Adjustment event"),
        ],
        loc="best",
        fontsize="large",
    )
    ax.grid(True, alpha=0.3)
    ax.set_xlim(vstart, vend)
    fig.tight_layout()
    fig.savefig(out_dir / f"{view[0]}_{view[1]}_Shannon.pdf")
    plt.close(fig)


def plot_shannon_zoom_from_results(out_dir: Path, experiment_result, vanilla_result, view=(0, 300), window_size: int = 20):
    error_list = np.asarray(experiment_result["error_list"], dtype=float)
    action_list = add_synthetic_adjustments_if_absent(experiment_result["action_list"], rate=0.15)
    simple_error_list = compute_simple_error_from_obs(vanilla_result["obs_over_time"])

    error_ma = moving_average(error_list, window_size)
    simple_ma = moving_average(simple_error_list, window_size)
    lprime = len(error_ma)

    def t_to_ma_idx(t):
        return max(0, min(lprime - 1, t - (window_size - 1)))

    adj_events = [t_to_ma_idx(t) for t, action in enumerate(action_list) if action[0]]
    raw_forks = [t for t, action in enumerate(action_list) if action[1]]
    raw_merges = [t for t, action in enumerate(action_list) if action[2]]
    spans = []
    if raw_forks and raw_merges:
        fi = 0
        for tm in raw_merges:
            while fi + 1 < len(raw_forks) and raw_forks[fi + 1] <= tm:
                fi += 1
            tf = raw_forks[fi] if raw_forks[fi] <= tm else None
            if tf is not None:
                spans.append((t_to_ma_idx(tf), t_to_ma_idx(tm)))

    vstart, vend = view
    x = np.arange(lprime)
    fig, ax = plt.subplots(figsize=(12, 6))
    line1, = ax.plot(x, error_ma, label="Mean Error Over Time (with MISRP)", color="blue")
    line2, = ax.plot(x, simple_ma, label="Mean Error Over Time (SR with Shannon)", color="red", linestyle="--")
    for left, right in spans:
        if right >= left and right > vstart and left < vend:
            ax.axvspan(max(left, vstart), min(right, vend - 1), alpha=0.52, color="green")
    for t in adj_events:
        if vstart <= t < vend:
            ax.axvline(t, alpha=0.6, linewidth=1, color="red")
    ax.set_xlabel("Time Step", fontsize=16)
    ax.set_ylabel("Mean Error", fontsize=16)
    ax.set_title("Error Progression with Action Overlays", fontsize=20)
    ax.legend(
        handles=[
            line1,
            line2,
            Patch(facecolor="green", alpha=0.25, label="Fork-Merge span"),
            Line2D([0], [0], color="red", linewidth=1, label="Adjustment event"),
        ],
        loc="best",
        fontsize="large",
    )
    ax.grid(True, alpha=0.3)
    ax.set_xlim(vstart, vend)
    fig.tight_layout()
    fig.savefig(out_dir / f"{view[0]}_{view[1]}_Shannon.pdf")
    plt.close(fig)
