from __future__ import annotations

import math
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass

import numpy as np


REPETITIONS = 160
MASTER_SEED = 20260802


@dataclass(frozen=True)
class Cell:
    family: str
    d: int
    horizon: int
    sigma: float
    radius: float
    prior_scale: float
    anisotropy: float


def source_bound(d: int, horizon: int, sigma: float, radius: float, eigenvalues: np.ndarray) -> dict[str, float]:
    q = 24.0 * math.log(horizon) / d
    c1 = math.sqrt(1.0 + max(q, math.sqrt(q)))
    c2 = 2.0 * c1 * math.sqrt(math.log(1.0 + radius**2 * float(eigenvalues.max()) * horizon / (d * sigma**2)))
    minimax = d * sigma * math.sqrt(horizon) * c2
    trace_sqrt = float(np.sqrt(eigenvalues).sum())
    trace = float(eigenvalues.sum())
    prior = 3.0 * math.sqrt(2.0) * radius * math.sqrt(d) * trace_sqrt * c1
    tail = math.sqrt(2.0 * radius**2 * trace)
    return {
        "C1": c1,
        "C2": c2,
        "minimax_term": minimax,
        "prior_trace_term": prior,
        "tail_term": tail,
        "theorem_1_bound": minimax + prior + tail,
        "corollary_burnin_scale": d * radius * math.sqrt(trace),
        "legacy_multiplicative_burnin_scale": d * radius * math.sqrt(trace) * math.sqrt(horizon),
        "trace_sqrt": trace_sqrt,
        "trace": trace,
        "operator_norm": float(eigenvalues.max()),
    }


def covariance_eigenvalues(cell: Cell) -> np.ndarray:
    if cell.anisotropy == 1.0:
        return np.full(cell.d, cell.prior_scale**2)
    return cell.prior_scale**2 * np.geomspace(1.0, 1.0 / cell.anisotropy, cell.d)


def simulate_cell(payload: tuple[int, Cell]) -> dict[str, float | int | str | list[float]]:
    cell_index, cell = payload
    rng = np.random.default_rng(np.random.SeedSequence([MASTER_SEED, cell_index]))
    eigenvalues = covariance_eigenvalues(cell)
    covariance = np.diag(eigenvalues)
    regrets = np.empty(REPETITIONS)
    frozen_regrets = np.empty(REPETITIONS)
    for repetition in range(REPETITIONS):
        theta = rng.multivariate_normal(np.zeros(cell.d), covariance)
        mean = np.zeros(cell.d)
        posterior = covariance.copy()
        regret = 0.0
        frozen_regret = 0.0
        optimal = cell.radius * float(np.linalg.norm(theta))
        for _ in range(cell.horizon):
            sampled = rng.multivariate_normal(mean, posterior, check_valid="raise")
            sampled_norm = float(np.linalg.norm(sampled))
            action = cell.radius * sampled / max(sampled_norm, 1e-300)
            regret += optimal - float(theta @ action)

            frozen_sample = rng.multivariate_normal(np.zeros(cell.d), covariance, check_valid="raise")
            frozen_action = cell.radius * frozen_sample / max(float(np.linalg.norm(frozen_sample)), 1e-300)
            frozen_regret += optimal - float(theta @ frozen_action)

            reward = float(theta @ action + rng.normal(0.0, cell.sigma))
            cov_action = posterior @ action
            innovation = cell.sigma**2 + float(action @ cov_action)
            mean = mean + cov_action * (reward - float(action @ mean)) / innovation
            posterior = posterior - np.outer(cov_action, cov_action) / innovation
            posterior = 0.5 * (posterior + posterior.T)
        regrets[repetition] = regret
        frozen_regrets[repetition] = frozen_regret

    theorem = source_bound(cell.d, cell.horizon, cell.sigma, cell.radius, eigenvalues)
    mean_regret = float(regrets.mean())
    se = float(regrets.std(ddof=1) / math.sqrt(REPETITIONS))
    frozen_mean = float(frozen_regrets.mean())
    frozen_se = float(frozen_regrets.std(ddof=1) / math.sqrt(REPETITIONS))
    return {
        "family": cell.family,
        "d": cell.d,
        "T": cell.horizon,
        "sigma": cell.sigma,
        "radius": cell.radius,
        "prior_scale": cell.prior_scale,
        "anisotropy": cell.anisotropy,
        "prior_eigenvalues": [float(value) for value in eigenvalues],
        "repetitions": REPETITIONS,
        "mean_bayesian_regret": mean_regret,
        "regret_standard_error": se,
        "regret_ci95_low": mean_regret - 1.96 * se,
        "regret_ci95_high": mean_regret + 1.96 * se,
        "frozen_posterior_regret": frozen_mean,
        "frozen_standard_error": frozen_se,
        "regret_to_bound_ratio": mean_regret / theorem["theorem_1_bound"],
        **theorem,
    }


def experiment_cells() -> list[Cell]:
    cells = []
    for scale in (0.25, 1.0, 4.0):
        for horizon in (64, 128, 256, 512, 1024, 2048):
            cells.append(Cell("horizon", 8, horizon, 0.5, 1.0, scale, 1.0))
    for dimension in (2, 4, 8, 16):
        cells.append(Cell("dimension", dimension, 512, 0.5, 1.0, 1.0, 1.0))
    for sigma in (0.125, 0.25, 0.5, 1.0, 2.0):
        cells.append(Cell("noise", 8, 512, sigma, 1.0, 1.0, 1.0))
    for dimension in (8, 16):
        for horizon in (128, 512, 2048):
            cells.append(Cell("anisotropic", dimension, horizon, 0.5, 1.0, 4.0, 64.0))
    return cells


def log_slope(rows: list[dict], value: str) -> float:
    return float(np.polyfit(np.log([row["T"] for row in rows]), np.log([row[value] for row in rows]), 1)[0])


def run(worker_count: int) -> dict[str, object]:
    cells = experiment_cells()
    with ProcessPoolExecutor(max_workers=worker_count) as executor:
        rows = list(executor.map(simulate_cell, enumerate(cells)))

    horizon_slopes = {}
    frozen_slopes = {}
    for scale in (0.25, 1.0, 4.0):
        subset = [row for row in rows if row["family"] == "horizon" and row["prior_scale"] == scale]
        horizon_slopes[str(scale)] = log_slope(subset, "mean_bayesian_regret")
        frozen_slopes[str(scale)] = log_slope(subset, "frozen_posterior_regret")

    diffuse_ratios = []
    for horizon in (64, 128, 256, 512, 1024, 2048):
        diffuse = next(row for row in rows if row["family"] == "horizon" and row["prior_scale"] == 4.0 and row["T"] == horizon)
        tight = next(row for row in rows if row["family"] == "horizon" and row["prior_scale"] == 0.25 and row["T"] == horizon)
        excess = max(0.0, diffuse["mean_bayesian_regret"] - tight["mean_bayesian_regret"])
        diffuse_ratios.append({"T": horizon, "excess_to_additive_burnin": excess / diffuse["corollary_burnin_scale"]})

    summary = {
        "master_seed": MASTER_SEED,
        "cells": len(rows),
        "paths": len(rows) * REPETITIONS,
        "total_simulated_steps": sum(int(row["T"]) * REPETITIONS for row in rows),
        "horizon_log_slopes": horizon_slopes,
        "frozen_control_log_slopes": frozen_slopes,
        "max_regret_to_exact_theorem_1_bound": max(float(row["regret_to_bound_ratio"]) for row in rows),
        "max_relative_ci95_halfwidth": max(1.96 * float(row["regret_standard_error"]) / float(row["mean_bayesian_regret"]) for row in rows),
        "diffuse_prior_excess_ratios": diffuse_ratios,
        "max_diffuse_excess_to_additive_burnin": max(item["excess_to_additive_burnin"] for item in diffuse_ratios),
        "min_frozen_control_slope": min(frozen_slopes.values()),
        "max_adaptive_slope": max(horizon_slopes.values()),
    }
    return {"summary": summary, "rows": rows}