"""
vita.core — the VPFM engine

The Vivanco Proximity-to-Failure Model: the math that turns a calibrated lift
into per-set estimates of fresh capacity, fatigue, proximity to failure (RIR),
and recovery. This module is the publishable scientific core of VITA. It depends
on nothing but the standard library (numpy is imported lazily, only inside
``rir_interval``).

Everything else in VITA — calibration, workouts, predictions, planning,
coaching output — is a layer on top of this. All fatigue / RIR / capacity /
recovery mathematics lives here and nowhere else.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import List, Optional, Tuple

__all__ = [
    "LoadRepCurve",
    "Params",
    "Session",
    "SetResult",
    "total_load",
    "rir_interval",
    "BODYWEIGHT_FRACTION",
]


BODYWEIGHT_FRACTION = {
    "pullup": 1.00,
    "chinup": 1.00,
    "dip": 1.00,
    "pushup": 0.64,
    "squat": 0.85,
    "bodyweight_squat": 0.85,
    "lunge": 0.80,
}


def total_load(bodyweight: float, borne_fraction: float, added: float = 0.0) -> float:
    """Total effective load  w = borne_fraction * bodyweight + added.

    Places bodyweight and weighted calisthenics on one load axis (Case B).
    `added` may be negative for band/machine assistance, but the result must
    stay positive -- assistance cannot fully cancel the bodyweight load.
    """
    if bodyweight <= 0:
        raise ValueError("bodyweight must be positive")
    if not (0.0 < borne_fraction <= 1.0):
        raise ValueError("borne_fraction must be in (0, 1]")
    w = borne_fraction * bodyweight + added
    if w <= 0:
        raise ValueError("total effective load must be positive (assistance too large)")
    return w


@dataclass
class Params:
    """Fatigue dynamics parameters. Defaults are placeholders pending data."""

    k: float = 0.4  # per-set cost coefficient, in (0, 1)
    gamma: float = 2.0  # effort-convexity exponent, >= 1
    t_rec: float = 180.0  # recovery time constant, seconds

    def __post_init__(self) -> None:
        if not (0.0 < self.k < 1.0):
            raise ValueError("k must be in (0, 1)")
        if self.gamma < 1.0:
            raise ValueError("gamma must be >= 1")
        if self.t_rec <= 0.0:
            raise ValueError("t_rec must be positive")


@dataclass
class LoadRepCurve:
    """Static load-rep law.

    Either a fitted exponential curve (tau, L0) via `from_two_sets`, or a
    fixed-load constant (bodyweight Case A) via `fixed`.
    """

    tau: Optional[float] = None
    L0: Optional[float] = None
    fixed_fresh_max: Optional[float] = None  # rho_BW for Case A

    # ---- constructors -----------------------------------------------------
    @classmethod
    def from_two_sets(
        cls,
        set_a: Tuple[float, float],
        set_b: Tuple[float, float],
    ) -> "LoadRepCurve":
        """Fit (tau, L0) from two fresh sets to failure: (load, reps).

        For weighted/assisted calisthenics, pass TOTAL effective loads
        (see `total_load`).
        """
        (wa, ra), (wb, rb) = set_a, set_b
        if wa <= 0 or wb <= 0:
            raise ValueError("loads must be positive")
        if ra <= 0 or rb <= 0:
            raise ValueError("rep counts must be positive")
        if wa == wb:
            raise ValueError(
                "calibration needs two DIFFERENT loads; a single load cannot "
                "identify the curve (this is exactly the bodyweight Case A -- "
                "use LoadRepCurve.fixed instead)"
            )
        tau = (ra - rb) / math.log(wb / wa)
        if tau <= 0:
            raise ValueError(
                "inconsistent calibration: the heavier set must allow fewer reps"
            )
        L0 = wa * math.exp(ra / tau)
        return cls(tau=tau, L0=L0)

    @classmethod
    def fixed(cls, fresh_max_reps: float) -> "LoadRepCurve":
        """Fixed-load (bodyweight) case: one measured fresh set to failure.

        The leverage constant `b` is not needed here -- the load is never
        compared across values, so it cancels.
        """
        if fresh_max_reps <= 0:
            raise ValueError("fresh max reps must be positive")
        return cls(fixed_fresh_max=float(fresh_max_reps))

    @classmethod
    def from_sets(cls, sets) -> "LoadRepCurve":
        """Fit (tau, L0) from two OR MORE fresh sets to failure: [(load, reps), ...].

        The least-squares generalization of `from_two_sets`. The load-rep law
        rho(w) = tau * ln(L0 / w) is linear in ln(load), so ordinary least
        squares on (ln load, reps) recovers the best-fit curve. For exactly two
        sets this reduces to the same line `from_two_sets` draws.

        For weighted/assisted calisthenics, pass TOTAL effective loads
        (see `total_load`).
        """
        pts = [(float(w), float(r)) for w, r in sets]
        if len(pts) < 2:
            raise ValueError("need at least two sets to fit a curve")
        for w, r in pts:
            if w <= 0:
                raise ValueError("loads must be positive")
            if r <= 0:
                raise ValueError("rep counts must be positive")

        n = len(pts)
        xs = [math.log(w) for w, _ in pts]
        ys = [r for _, r in pts]
        xbar = sum(xs) / n
        ybar = sum(ys) / n
        sxx = sum((x - xbar) ** 2 for x in xs)
        if sxx <= 0:
            raise ValueError(
                "all loads are equal; a single load cannot identify the curve "
                "(this is the bodyweight Case A -- use LoadRepCurve.fixed instead)"
            )
        sxy = sum((x - xbar) * (y - ybar) for x, y in zip(xs, ys))
        b = sxy / sxx  # slope of reps vs ln(load); should be negative
        a = ybar - b * xbar
        tau = -b
        if tau <= 0:
            raise ValueError(
                "inconsistent calibration: heavier sets must allow fewer reps"
            )
        L0 = math.exp(a / tau)
        return cls(tau=tau, L0=L0)

    # ---- queries ----------------------------------------------------------
    @property
    def is_fixed(self) -> bool:
        return self.fixed_fresh_max is not None

    def fresh_reps(self, load: Optional[float] = None) -> float:
        """rho(w): fresh max reps at a load.

        In fixed-load mode the measured fresh max is returned and `load` is
        ignored. For a fitted curve, a load in (0, L0] is required; loads at
        or above L0 return 0 (not liftable for a full rep).
        """
        if self.is_fixed:
            return self.fixed_fresh_max
        if load is None:
            raise ValueError("a load is required for a fitted curve")
        if load <= 0:
            raise ValueError("load must be positive")
        return max(0.0, self.tau * math.log(self.L0 / load))

    @property
    def one_rep_max(self) -> float:
        """Predicted 1RM = L0 * exp(-1/tau). Undefined in fixed-load mode."""
        if self.is_fixed:
            raise ValueError("1RM is undefined in fixed-load (bodyweight) mode")
        return self.L0 * math.exp(-1.0 / self.tau)

    def describe(self) -> str:
        if self.is_fixed:
            return f"fixed-load curve: fresh max = {self.fixed_fresh_max:.2f} reps"
        return (
            f"fitted curve: tau = {self.tau:.2f} reps, L0 = {self.L0:.1f}, "
            f"1RM = {self.one_rep_max:.1f}"
        )


@dataclass
class SetResult:
    """Outcome of one set, returned by Session.do_set."""

    index: int
    load: Optional[float]
    reps: float
    fresh_max: float  # rho(w)
    fatigued_max: float  # F = C * rho
    effort: float  # e
    rir: float  # Phi
    capacity_before: float
    capacity_after_set: float
    capacity_after_rest: float


class Session:
    """Runs the per-set capacity battery for one training session."""

    def __init__(
        self,
        curve: LoadRepCurve,
        params: Optional[Params] = None,
        capacity: float = 1.0,
    ) -> None:
        if not (0.0 < capacity <= 1.0):
            raise ValueError("starting capacity must be in (0, 1]")
        self.curve = curve
        self.p = params if params is not None else Params()
        self.C = float(capacity)
        self.history: List[SetResult] = []

    @property
    def capacity(self) -> float:
        return self.C

    def predicted_rir(self, reps: float, load: Optional[float] = None) -> float:
        """RIR the model expects for a contemplated set at the CURRENT capacity,
        without recording it or changing state."""
        return self.C * self.curve.fresh_reps(load) - reps

    def reps_for_rir(self, target_rir: float, load: Optional[float] = None) -> float:
        """Reps that land you at a target reps-in-reserve, at the CURRENT capacity.

        Inverts the RIR relationship (RIR = F - reps), so it accounts for the
        fatigue accumulated so far this session. Use it to prescribe each set to
        a chosen proximity to failure -- e.g. hold every working set at 2 RIR.
        Returns the exact (real-valued) rep target; round it for practice.
        """
        if target_rir < 0:
            raise ValueError("target_rir must be >= 0")
        return max(0.0, self.C * self.curve.fresh_reps(load) - target_rir)

    def do_set(
        self,
        reps: float,
        load: Optional[float] = None,
        rest: Optional[float] = None,
    ) -> SetResult:
        """Perform a set, update the capacity state, and return the result.

        `rest` is the recovery time AFTER this set in seconds (defaults to the
        recovery time constant t_rec; pass rest=0 for the final set).
        """
        if reps < 0:
            raise ValueError("reps must be non-negative")
        rho = self.curve.fresh_reps(load)
        c_before = self.C
        F = c_before * rho
        e = (reps / F) if F > 0 else math.inf
        rir = F - reps

        # A4 -- bounded convex depletion (valid for every e >= 0)
        if math.isfinite(e):
            c_plus = c_before * math.exp(-self.p.k * (e**self.p.gamma))
        else:
            c_plus = 0.0

        # A5 -- exponential recovery during rest
        dt = self.p.t_rec if rest is None else float(rest)
        if dt < 0:
            raise ValueError("rest must be non-negative")
        c_next = 1.0 - (1.0 - c_plus) * math.exp(-dt / self.p.t_rec)

        self.C = c_next
        res = SetResult(
            index=len(self.history) + 1,
            load=load,
            reps=reps,
            fresh_max=rho,
            fatigued_max=F,
            effort=e,
            rir=rir,
            capacity_before=c_before,
            capacity_after_set=c_plus,
            capacity_after_rest=c_next,
        )
        self.history.append(res)
        return res


def rir_interval(
    curve: LoadRepCurve,
    reps: float,
    load: Optional[float] = None,
    capacity: float = 1.0,
    param_cv: float = 0.05,
    noise_frac: float = 0.10,
    n: int = 10_000,
    interval: float = 0.80,
    seed: Optional[int] = None,
) -> Tuple[float, float, float]:
    """Probabilistic RIR for a single contemplated set.

    Propagates two sources of uncertainty by Monte Carlo and returns
    (median, low, high) for the central `interval` (default 80%):

      * param_cv   relative SD on the curve parameters (tau, L0, or rho_BW),
                   a placeholder for calibration uncertainty until fit from data.
      * noise_frac biological scatter SD as a fraction of F (heteroscedastic:
                   more scatter at higher reps).

    Reporting an interval -- not a single decimal -- is what makes the model
    falsifiable: check whether the stated 80% intervals contain the realized
    outcome ~80% of the time.
    """
    import numpy as np

    if not (0.0 < interval < 1.0):
        raise ValueError("interval must be in (0, 1)")
    rng = np.random.default_rng(seed)

    if curve.is_fixed:
        rho = np.maximum(
            curve.fixed_fresh_max * (1.0 + param_cv * rng.standard_normal(n)), 1e-6
        )
    else:
        if load is None:
            raise ValueError("a load is required for a fitted curve")
        tau = np.maximum(curve.tau * (1.0 + param_cv * rng.standard_normal(n)), 1e-6)
        L0 = np.maximum(
            curve.L0 * (1.0 + param_cv * rng.standard_normal(n)), load * 1.0001
        )
        rho = np.maximum(tau * np.log(L0 / load), 0.0)

    F = capacity * rho
    draws = (F - reps) + noise_frac * F * rng.standard_normal(n)

    lo_q = (1.0 - interval) / 2.0
    hi_q = 1.0 - lo_q
    return (
        float(np.median(draws)),
        float(np.quantile(draws, lo_q)),
        float(np.quantile(draws, hi_q)),
    )
