import math
from typing import Any, Dict, List, Optional, Sequence

from core import clamp


def brier_score(preds: Sequence[float], labels: Sequence[int]) -> float:
    if len(preds) != len(labels):
        raise ValueError("brier_score: preds and labels length mismatch")
    if not preds:
        return 0.0
    s = 0.0
    for p, y in zip(preds, labels):
        pp = clamp(float(p), 0.0, 1.0)
        yy = 1.0 if int(y) else 0.0
        d = pp - yy
        s += d * d
    return s / float(len(preds))


def log_loss(preds: Sequence[float], labels: Sequence[int], eps: float = 1e-12) -> float:
    if len(preds) != len(labels):
        raise ValueError("log_loss: preds and labels length mismatch")
    if not preds:
        return 0.0
    s = 0.0
    for p, y in zip(preds, labels):
        pp = clamp(float(p), 0.0, 1.0)
        pp = clamp(pp, eps, 1.0 - eps)
        yy = 1 if int(y) else 0
        s += -(yy * math.log(pp) + (1 - yy) * math.log(1.0 - pp))
    return s / float(len(preds))


def summarize_u(u: Sequence[float]) -> Dict[str, float]:
    if not u:
        return {"mean": 0.0, "sum": 0.0, "sum_sq": 0.0, "du_abs_sum": 0.0}

    s = 0.0
    s2 = 0.0
    du = 0.0
    prev = float(u[0])
    for i, x in enumerate(u):
        xx = clamp(float(x), 0.0, 1.0)
        s += xx
        s2 += xx * xx
        if i > 0:
            du += abs(xx - prev)
            prev = xx
    return {"mean": s / float(len(u)), "sum": s, "sum_sq": s2, "du_abs_sum": du}


def recovery_time(in_e: Sequence[bool], t0: int) -> Optional[int]:
    n = len(in_e)
    if t0 < 0 or t0 >= n:
        return None

    t_leave: Optional[int] = None
    for t in range(t0, n):
        if not bool(in_e[t]):
            t_leave = t
            break

    if t_leave is None:
        return 0

    for t in range(t_leave + 1, n):
        if bool(in_e[t]):
            return t - t0

    return None


def recovery_times(in_e: Sequence[bool], shock_times: Sequence[int]) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    for t0 in shock_times:
        rt = recovery_time(in_e, int(t0))
        out.append({"t0": int(t0), "recovery": rt})
    return out