from __future__ import annotations

import math
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_DUMP_PATH = ROOT / "eval" / "probe_dump.npz"


@dataclass(frozen=True)
class BocpdResult:
    step: int
    value: float
    cp_prob: float
    map_run_length: int
    change_point: bool


def logsumexp(values: np.ndarray) -> float:
    max_value = float(np.max(values))
    if not math.isfinite(max_value):
        return max_value
    return max_value + math.log(float(np.sum(np.exp(values - max_value))))


def student_t_logpdf(x: float, mu: np.ndarray, kappa: np.ndarray, alpha: np.ndarray, beta: np.ndarray) -> np.ndarray:
    nu = 2.0 * alpha
    scale = np.sqrt(beta * (kappa + 1.0) / (alpha * kappa))
    z = (x - mu) / scale
    return (
        np.vectorize(math.lgamma)((nu + 1.0) / 2.0)
        - np.vectorize(math.lgamma)(nu / 2.0)
        - 0.5 * np.log(nu * math.pi)
        - np.log(scale)
        - ((nu + 1.0) / 2.0) * np.log1p((z * z) / nu)
    )


def update_nig(
    x: float,
    mu: np.ndarray,
    kappa: np.ndarray,
    alpha: np.ndarray,
    beta: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    next_kappa = kappa + 1.0
    next_mu = (kappa * mu + x) / next_kappa
    next_alpha = alpha + 0.5
    next_beta = beta + 0.5 * kappa * (x - mu) ** 2 / next_kappa
    return next_mu, next_kappa, next_alpha, next_beta


def run_bocpd(
    values: np.ndarray,
    hazard: float = 1.0 / 50.0,
    prior_mu: float = 0.0,
    prior_kappa: float = 1.0e-3,
    prior_alpha: float = 1.0,
    prior_beta: float = 1.0,
) -> list[BocpdResult]:
    log_hazard = math.log(hazard)
    log_growth_factor = math.log1p(-hazard)

    log_run_probs = np.asarray([0.0], dtype=np.float64)
    mu = np.asarray([prior_mu], dtype=np.float64)
    kappa = np.asarray([prior_kappa], dtype=np.float64)
    alpha = np.asarray([prior_alpha], dtype=np.float64)
    beta = np.asarray([prior_beta], dtype=np.float64)

    results: list[BocpdResult] = []
    previous_map_run_length: int | None = None

    for step, x_value in enumerate(values, start=1):
        x = float(x_value)
        predictive = student_t_logpdf(x, mu, kappa, alpha, beta)

        growth_probs = log_run_probs + predictive + log_growth_factor
        cp_prob = logsumexp(log_run_probs + predictive + log_hazard)
        new_log_run_probs = np.concatenate(([cp_prob], growth_probs))
        normalizer = logsumexp(new_log_run_probs)
        new_log_run_probs -= normalizer

        grown_mu, grown_kappa, grown_alpha, grown_beta = update_nig(x, mu, kappa, alpha, beta)
        mu = np.concatenate(([prior_mu], grown_mu))
        kappa = np.concatenate(([prior_kappa], grown_kappa))
        alpha = np.concatenate(([prior_alpha], grown_alpha))
        beta = np.concatenate(([prior_beta], grown_beta))
        log_run_probs = new_log_run_probs

        map_run_length = int(np.argmax(log_run_probs))
        cp_probability = float(np.exp(log_run_probs[0]))
        change_point = bool(
            map_run_length == 0
            or (previous_map_run_length is not None and map_run_length < previous_map_run_length)
        )
        results.append(
            BocpdResult(
                step=step,
                value=x,
                cp_prob=cp_probability,
                map_run_length=map_run_length,
                change_point=change_point,
            )
        )
        previous_map_run_length = map_run_length

    return results


def main() -> None:
    dump_path = Path(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_DUMP_PATH
    if not dump_path.is_absolute():
        dump_path = ROOT / dump_path

    if not dump_path.exists():
        print(f"No probe dump found at {dump_path}")
        return

    dump = np.load(dump_path, allow_pickle=True)
    failure = str(dump.get("failure", ""))
    values = np.asarray(dump["nll_series"], dtype=np.float64)
    if values.size == 0:
        print("No NLL samples found; skipping BOCPD.")
        if failure:
            print(f"Probe failure: {failure}")
        return

    print("step | nll | cp_prob | map_run_length | change_point")
    print("-----|-----|---------|----------------|-------------")
    for result in run_bocpd(values):
        flag = "YES" if result.change_point else "no"
        print(
            f"{result.step:>4} | {result.value:.4f} | {result.cp_prob:.4f} | "
            f"{result.map_run_length:>14} | {flag}"
        )


if __name__ == "__main__":
    main()