from __future__ import annotations

import math
import random
from dataclasses import dataclass


def mean_std(values: list[float]) -> tuple[float, float]:
    n = len(values)
    if n == 0:
        return 0.0, 1.0
    mean = sum(values) / n
    if n == 1:
        return mean, 1.0
    var = sum((x - mean) ** 2 for x in values) / (n - 1)
    return mean, math.sqrt(var) if var > 0 else 1.0


def zscores(values: list[float | None]) -> list[float | None]:
    present = [v for v in values if v is not None]
    mean, std = mean_std(present)
    out: list[float | None] = []
    for v in values:
        if v is None:
            out.append(None)
        else:
            out.append((v - mean) / std)
    return out


def ranks_desc(values: list[float | None]) -> list[int | None]:
    indexed = [(i, v) for i, v in enumerate(values) if v is not None]
    indexed.sort(key=lambda iv: -iv[1])
    out: list[int | None] = [None] * len(values)
    for rank, (i, _) in enumerate(indexed, 1):
        out[i] = rank
    return out


@dataclass
class FilterResult:
    x: float | None
    p: float | None
    xs: list[float | None]
    ps: list[float | None]


def kalman_filter(zs: list[float | None], q: float, r: float) -> FilterResult:
    x = None
    p = None
    xs: list[float | None] = []
    ps: list[float | None] = []
    for z in zs:
        if x is None:
            if z is None:
                xs.append(None)
                ps.append(None)
                continue
            x, p = z, r
            xs.append(x)
            ps.append(p)
            continue
        x_pred = x
        p_pred = p + q
        if z is None:
            x, p = x_pred, p_pred
            xs.append(x)
            ps.append(p)
            continue
        k = p_pred / (p_pred + r)
        x = x_pred + k * (z - x_pred)
        p = (1 - k) * p_pred
        xs.append(x)
        ps.append(p)
    return FilterResult(x=x, p=p, xs=xs, ps=ps)


def one_step_rmse(series_list: list[list[float | None]], q: float, r: float) -> float | None:
    se = 0.0
    n = 0
    for zs in series_list:
        x = None
        p = None
        for z in zs:
            if x is None:
                if z is None:
                    continue
                x, p = z, r
                continue
            x_pred = x
            p_pred = p + q
            if z is None:
                x, p = x_pred, p_pred
                continue
            se += (x_pred - z) ** 2
            n += 1
            k = p_pred / (p_pred + r)
            x = x_pred + k * (z - x_pred)
            p = (1 - k) * p_pred
    if n == 0:
        return None
    return math.sqrt(se / n)


def fit_qr(series_list: list[list[float | None]]) -> tuple[float, float, float]:
    usable = [s for s in series_list if sum(1 for z in s if z is not None) >= 2]
    best = (1e9, 0.08, 0.25)
    for q in (0.03, 0.05, 0.08, 0.10, 0.15, 0.22, 0.30):
        for r in (0.12, 0.18, 0.25, 0.32, 0.40, 0.55):
            rmse = one_step_rmse(usable, q, r)
            if rmse is not None and rmse < best[0]:
                best = (rmse, q, r)
    return best[1], best[2], best[0]


def holdout_hits(
    series_list: list[list[float | None]],
    last_ranks: list[int | None],
    true_ranks: list[int | None],
    q: float,
    r: float,
    target_k: int,
) -> dict:
    preds: list[tuple[int, float]] = []
    last_pred: list[tuple[int, int]] = []
    for i, zs in enumerate(series_list):
        if true_ranks[i] is None:
            continue
        train = zs[:-1]
        result = kalman_filter(train, q, r)
        if result.x is None:
            continue
        preds.append((i, result.x))
        if last_ranks[i] is not None:
            last_pred.append((i, last_ranks[i]))
    preds.sort(key=lambda t: -t[1])
    kf_set = {i for i, _ in preds[:target_k]}
    last_pred.sort(key=lambda t: t[1])
    last_set = {i for i, _ in last_pred[:target_k]}
    actual = {i for i, rk in enumerate(true_ranks) if rk is not None and rk <= target_k}
    kf_hit = len(kf_set & actual)
    last_hit = len(last_set & actual)
    return {
        "n_actual": len(actual),
        "kf_hit": kf_hit,
        "last_hit": last_hit,
        "kf_precision": kf_hit / target_k if target_k else 0,
        "last_precision": last_hit / target_k if target_k else 0,
    }


def monte_carlo(
    means: list[float],
    variances: list[float],
    class_ids: list[str],
    target_k: int,
    n_sim: int = 600,
    seed: int = 42,
) -> tuple[list[float], dict[str, dict]]:
    rng = random.Random(seed)
    n = len(means)
    hits = [0] * n
    sigmas = [math.sqrt(max(v, 1e-8)) for v in variances]
    unique = list(dict.fromkeys(class_ids))
    buckets = {c: [] for c in unique}
    for _ in range(n_sim):
        zs = [means[i] + sigmas[i] * rng.gauss(0, 1) for i in range(n)]
        order = sorted(range(n), key=lambda i: zs[i], reverse=True)
        counts = {c: 0 for c in unique}
        for i in order[:target_k]:
            hits[i] += 1
            counts[class_ids[i]] += 1
        for c, v in counts.items():
            buckets[c].append(v)
    out = {}
    for c, arr in buckets.items():
        arr.sort()
        m = n_sim
        out[c] = {
            "p20": arr[max(0, int(0.20 * m) - 1)],
            "p50": arr[max(0, int(0.50 * m) - 1)],
            "p80": arr[max(0, int(0.80 * m) - 1)],
            "mean": sum(arr) / m,
        }
    return [h / n_sim for h in hits], out


def hamilton(values: list[float], total: int) -> list[int]:
    if total <= 0 or not values:
        return [0] * len(values)
    s = sum(values)
    if s <= 0:
        out = [0] * len(values)
        return out
    scaled = [v * total / s for v in values]
    floors = [int(math.floor(v + 1e-9)) for v in scaled]
    remain = total - sum(floors)
    order = sorted(
        range(len(scaled)),
        key=lambda i: (scaled[i] - floors[i], scaled[i]),
        reverse=True,
    )
    for i in order[: max(0, remain)]:
        floors[i] += 1
    return floors