"""
vita.fitting — calibrating the fatigue model to your data

The VPFM fatigue parameters (k, gamma, t_rec) ship as honest placeholders. This
module fits them to your own logged sets and checks the fit by held-out
prediction.

The idea: a logged set's reps-in-reserve, with your calibrated lift, pins down
your capacity at that set --

    capacity_before = (rir + reps) / fresh_reps(load)

The trajectory of that capacity across a session is governed by k (fatigue rate),
gamma (how sharply effort bites), and t_rec (recovery during rest). We fit the
parameters the data can actually identify, then validate: does the fitted model
predict held-out sessions better than the generic defaults?

    fit = vita.fit_fatigue(bench, sessions)
    fit.validated          # did it beat the defaults on held-out data?
    fit.params             # the fitted Params
    w = vita.Workout(bench, params=fit.params)   # train on YOUR fatigue model

Identifiability is respected, not assumed:
  • k       fitted whenever there's data
  • t_rec   fitted only if your rest times vary (otherwise held — it's invisible)
  • gamma   fitted only when you ask (fit_gamma=True) AND effort varies, since it
            needs sets across a wide range of proximity-to-failure

All fatigue math still lives in vita.core — the fitter reuses the engine for
every forward step.
"""

from __future__ import annotations

import itertools
from dataclasses import dataclass
from typing import Callable, List, Optional, Sequence, Tuple

import numpy as np

from .core import Params, Session, rir_interval

__all__ = ["ObservedSet", "observed_session", "FatigueFit", "fit_fatigue", "coverage"]


# --------------------------------------------------------------------------
# observed data
# --------------------------------------------------------------------------
@dataclass
class ObservedSet:
    """One logged set: the load, reps done, reps-in-reserve reported, and the
    rest taken AFTER it (seconds)."""

    load: Optional[float]
    reps: float
    rir: float
    rest: float = 180.0


def observed_session(rows: Sequence[Sequence[float]]) -> List[ObservedSet]:
    """Build a session from rows of (load, reps, rir[, rest])."""
    out: List[ObservedSet] = []
    for row in rows:
        load, reps, rir = row[0], row[1], row[2]
        rest = row[3] if len(row) > 3 else 180.0
        out.append(ObservedSet(load, reps, rir, rest))
    return out


# --------------------------------------------------------------------------
# the fit report
# --------------------------------------------------------------------------
@dataclass(frozen=True)
class FatigueFit:
    """The outcome of fitting the fatigue model to logged sessions."""

    params: Params
    fitted: Tuple[str, ...]  # which parameters were actually fitted
    rmse: float  # in-sample one-step RMSE (capacity units)
    val_rmse: Optional[float]  # held-out RMSE of the fitted model
    baseline_rmse: Optional[float]  # held-out RMSE of the default params
    validated: bool  # fitted beat the defaults on held-out sessions
    n_sets: int
    messages: Tuple[str, ...]

    @property
    def improvement(self) -> Optional[float]:
        """Fractional reduction in held-out error vs the defaults (None if not validated)."""
        if self.val_rmse is None or not self.baseline_rmse:
            return None
        return 1.0 - self.val_rmse / self.baseline_rmse

    def summary(self) -> str:
        p = self.params
        held = (
            ", ".join(n for n in ("k", "gamma", "t_rec") if n not in self.fitted)
            or "none"
        )
        lines = [
            f"  Fatigue fit — {'VALIDATED' if self.validated else 'NOT validated'}",
            f"  k = {p.k:.3f} · gamma = {p.gamma:.3f} · t_rec = {p.t_rec:.0f}s",
            f"  fitted: {', '.join(self.fitted)}   held at default: {held}",
            f"  in-sample RMSE = {self.rmse:.4f} capacity · fitted on {self.n_sets} sets",
        ]
        if self.val_rmse is not None:
            imp = self.improvement
            lines.append(
                f"  held-out RMSE = {self.val_rmse:.4f} vs default {self.baseline_rmse:.4f} "
                f"({imp * 100:+.0f}% vs defaults)"
            )
        for m in self.messages:
            lines.append(f"  • {m}")
        return "\n".join(lines)

    def __str__(self) -> str:
        return self.summary()

    def __repr__(self) -> str:
        verdict = "validated" if self.validated else "unvalidated"
        return (
            f"<FatigueFit {verdict} · k={self.params.k:.2f} gamma={self.params.gamma:.2f}"
            f" t_rec={self.params.t_rec:.0f} · fitted {self.fitted}>"
        )


# --------------------------------------------------------------------------
# internals
# --------------------------------------------------------------------------
_BOUNDS = {"k": (0.02, 0.95), "gamma": (1.0, 4.0), "t_rec": (20.0, 600.0)}


def _capacity_before(lift, s: ObservedSet) -> float:
    rho = lift.fresh_reps(s.load)
    return (s.rir + s.reps) / rho if rho > 0 else float("nan")


def _make_params(vec, fit_names: Sequence[str], default: Params) -> Params:
    vals = {"k": default.k, "gamma": default.gamma, "t_rec": default.t_rec}
    for name, v in zip(fit_names, vec):
        vals[name] = float(v)
    return Params(**vals)


def _mean_sse(lift, sessions, params: Params) -> Tuple[float, int]:
    """Mean squared one-step-ahead capacity error over all session transitions."""
    sse = 0.0
    n = 0
    for sess in sessions:
        caps = [_capacity_before(lift, s) for s in sess]
        for i in range(len(sess) - 1):
            c_i, c_next = caps[i], caps[i + 1]
            if not (0.0 < c_i <= 1.5) or not np.isfinite(c_next):
                continue
            eng = Session(lift.curve, params, capacity=min(1.0, c_i))
            eng.do_set(reps=sess[i].reps, load=sess[i].load, rest=sess[i].rest)
            sse += (eng.capacity - c_next) ** 2
            n += 1
    return sse, n


def _rmse(lift, sessions, params: Params) -> float:
    sse, n = _mean_sse(lift, sessions, params)
    return float(np.sqrt(sse / n)) if n else float("inf")


def _objective(lift, sessions, fit_names: Sequence[str], default: Params) -> Callable:
    def f(vec):
        try:
            p = _make_params(vec, fit_names, default)
            sse, n = _mean_sse(lift, sessions, p)
            return (sse / n) if n else float("inf")
        except (ValueError, ZeroDivisionError):
            return float("inf")

    return f


def _grid_refine(objective, bounds, steps: int = 7, rounds: int = 4):
    """Derivative-free minimizer: evaluate a grid, zoom in around the best, repeat."""
    los = [lo for lo, _ in bounds]
    his = [hi for _, hi in bounds]
    best, best_val = None, float("inf")
    for _ in range(rounds):
        axes = [np.linspace(lo, hi, steps) for lo, hi in zip(los, his)]
        for combo in itertools.product(*axes):
            v = objective(np.array(combo))
            if v < best_val:
                best_val, best = v, np.array(combo)
        new_los, new_his = [], []
        for d, (lo, hi) in enumerate(zip(los, his)):
            span = (hi - lo) / (steps - 1)
            new_los.append(max(bounds[d][0], best[d] - span))
            new_his.append(min(bounds[d][1], best[d] + span))
        los, his = new_los, new_his
    return best, best_val


def _fit_on(lift, sessions, fit_names: Sequence[str], default: Params) -> Params:
    if not fit_names:
        return default
    bounds = [_BOUNDS[n] for n in fit_names]
    best, _ = _grid_refine(_objective(lift, sessions, fit_names, default), bounds)
    return _make_params(best, fit_names, default)


def _effort(s: ObservedSet) -> float:
    # e = reps / fatigued_max = reps / (reps + rir), directly observable
    denom = s.reps + s.rir
    return s.reps / denom if denom > 0 else 1.0


def _cv(values) -> float:
    arr = np.asarray(values, dtype=float)
    m = float(np.mean(arr))
    return float(np.std(arr) / m) if m else 0.0


# --------------------------------------------------------------------------
# optional diagnostic: interval coverage
# --------------------------------------------------------------------------
def coverage(
    lift, sessions, params: Params, interval: float = 0.80, seed: int = 0
) -> float:
    """Fraction of observed RIRs inside the model's predicted intervals.

    A diagnostic for the model's *uncertainty* settings (param_cv / noise_frac),
    not the point fit — each session is simulated forward from fresh and the
    observed RIR is checked against the predicted interval at each set.
    """
    inside, total = 0, 0
    for sess in sessions:
        eng = Session(lift.curve, params, capacity=1.0)
        for s in sess:
            _, lo, hi = rir_interval(
                lift.curve,
                reps=s.reps,
                load=s.load,
                capacity=eng.capacity,
                interval=interval,
                seed=seed,
            )
            if lo <= s.rir <= hi:
                inside += 1
            total += 1
            eng.do_set(reps=s.reps, load=s.load, rest=s.rest)
    return inside / total if total else float("nan")


# --------------------------------------------------------------------------
# the fitter
# --------------------------------------------------------------------------
def fit_fatigue(
    lift,
    sessions: List[List[ObservedSet]],
    fit_gamma: bool = False,
    fit_t_rec: Optional[bool] = None,
    default_params: Optional[Params] = None,
) -> FatigueFit:
    """Fit the fatigue parameters to logged sessions.

    `sessions` is a list of sessions, each a list of ObservedSet (or build them
    with observed_session). Needs a few within-session transitions.

    k is always fitted. t_rec is fitted only when rest times vary enough to
    identify it (fit_t_rec=None auto-detects; True/False force it). gamma is
    fitted only when fit_gamma=True and effort varies, since pinning down the
    effort-convexity needs sets across a wide range of proximity-to-failure.
    """
    default = default_params or Params()
    n_sets = sum(len(s) for s in sessions)
    transitions = sum(max(0, len(s) - 1) for s in sessions)
    if transitions < 3:
        raise ValueError(
            "need at least 3 within-session set-to-set transitions to fit "
            "(log a few multi-set sessions first)"
        )

    rests = [s.rest for sess in sessions for s in sess]
    efforts = [_effort(s) for sess in sessions for s in sess]
    rest_cv = _cv(rests)
    effort_range = (max(efforts) - min(efforts)) if efforts else 0.0

    messages: List[str] = []

    # --- decide which parameters the data can identify ---
    fit_names = ["k"]

    want_t_rec = (rest_cv > 0.15) if fit_t_rec is None else bool(fit_t_rec)
    if want_t_rec:
        fit_names.append("t_rec")
    elif fit_t_rec is None:
        messages.append(
            f"Rest barely varied ({rest_cv * 100:.0f}% CV) — t_rec isn't identifiable, "
            f"held at {default.t_rec:.0f}s. Vary your rest to fit it."
        )

    if fit_gamma:
        if effort_range > 0.15:
            fit_names.append("gamma")
        else:
            messages.append(
                f"Effort barely varied (range {effort_range:.2f}) — gamma can't be "
                f"identified, held at {default.gamma:.2f}. Mix in easy and near-failure sets."
            )

    # keep a stable order
    fit_names = [n for n in ("k", "gamma", "t_rec") if n in fit_names]

    # --- final fit on all data ---
    params = _fit_on(lift, sessions, fit_names, default)
    rmse = _rmse(lift, sessions, params)

    # --- validation by held-out prediction ---
    val_rmse = baseline_rmse = None
    validated = False
    if len(sessions) >= 4:
        hold = max(1, round(len(sessions) * 0.3))
        train, val = sessions[:-hold], sessions[-hold:]
        p_train = _fit_on(lift, train, fit_names, default)
        val_rmse = _rmse(lift, val, p_train)
        baseline_rmse = _rmse(lift, val, default)
        validated = val_rmse < baseline_rmse * 0.98  # require a small real margin
        if validated:
            messages.append(
                f"Beats the default model on held-out sessions "
                f"({(1 - val_rmse / baseline_rmse) * 100:.0f}% lower error)."
            )
        else:
            messages.append(
                "No clear win over the defaults on held-out data — your fatigue is "
                "close to the generic model, or there isn't enough data yet."
            )
    else:
        messages.append(
            "Log at least 4 sessions to validate the fit against held-out data."
        )

    if n_sets < 12:
        messages.append("Few sets overall — more sessions will firm up the fit.")

    return FatigueFit(
        params=params,
        fitted=tuple(fit_names),
        rmse=rmse,
        val_rmse=val_rmse,
        baseline_rmse=baseline_rmse,
        validated=validated,
        n_sets=n_sets,
        messages=tuple(messages),
    )
