"""
vpfm.py -- The Vivanco Proximity-to-Failure Model (reference implementation)

A fatigue-adjusted estimator of reps in reserve (RIR) for resistance training.
It has two layers:

  1. A static load-rep law for the fresh lifter:
         rho(w) = tau * ln(L0 / w)
     fit from two calibration sets to failure -- OR, for fixed-load movements
     (pure bodyweight: pull-ups, push-ups, ...), a single MEASURED fresh max.

  2. A dynamic capacity state C in (0, 1], drained by hard sets and reconstituted
     during rest.

Per-set quantities:
     F(w, C) = C * rho(w)                  fatigued max reps
     e       = reps / F                    effort (proximity to failure)
     RIR     = F - reps                    reps in reserve   (Vivanco RIR, Phi)
     C_plus  = C * exp(-k * e**gamma)      depletion (A4; bounded for ALL e >= 0)
     C_next  = 1 - (1 - C_plus) * exp(-rest / t_rec)   recovery (A5)

Bodyweight movements are a special case, not a separate model:
     total effective load   w = b*m + a
       m = bodyweight, b = bodyweight-borne fraction (pull-up ~1.0, push-up ~0.64,
       bodyweight squat ~0.85), a = added load (a < 0 = assistance).
   Case A (pure bodyweight): load is fixed, so the curve cannot be fit and is not
       needed -- measure rho_BW directly with LoadRepCurve.fixed(...).
   Case B (weighted / assisted): load varies, so fit LoadRepCurve.from_two_sets(...)
       using TOTAL effective loads (use total_load(...) to compute them).

See "The Vivanco Proximity-to-Failure Model" (v0.3) for axioms and proofs.

Status: a theoretical model. Parameters (k, gamma, t_rec) and the curve must be
fit and validated against measured reps-to-failure; defaults below are placeholders.

Dependencies: standard library; numpy is used only for rir_interval (Monte Carlo).
"""

from __future__ import annotations

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

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

# Movement bodyweight-borne fractions (approximate; replace with fitted values).
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))

    # ---- 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)),
    )


# --------------------------------------------------------------------------
# Demonstration
# --------------------------------------------------------------------------
def _rule(width: int = 66) -> str:
    return "-" * width


def _print_trace(session: Session, show_load: bool = True) -> None:
    if show_load:
        print(f"  {'SET':<11}{'LOAD':>6}{'REPS':>6}{'F':>8}{'EFFORT':>8}{'RIR':>8}{'C_AFT':>8}")
    else:
        print(f"  {'SET':<11}{'REPS':>6}{'F':>8}{'EFFORT':>8}{'RIR':>8}{'C_AFT':>8}")
    print("  " + _rule(60 if show_load else 48))
    for r in session.history:
        if show_load:
            ld = "-" if r.load is None else f"{r.load:.0f}"
            print(
                f"  {r.index:<11}{ld:>6}{r.reps:>6.0f}{r.fatigued_max:>8.2f}"
                f"{r.effort:>8.3f}{r.rir:>8.2f}{r.capacity_after_rest:>8.3f}"
            )
        else:
            print(
                f"  {r.index:<11}{r.reps:>6.0f}{r.fatigued_max:>8.2f}"
                f"{r.effort:>8.3f}{r.rir:>8.2f}{r.capacity_after_rest:>8.3f}"
            )


def _demo() -> None:
    print(_rule())
    print("  THE VIVANCO PROXIMITY-TO-FAILURE MODEL  (vpfm reference demo)")
    print(_rule())

    # 1) Barbell -- fitted curve from two calibration sets ------------------
    print("\n[1] BARBELL  (fitted curve from two calibration sets to failure)\n")
    curve = LoadRepCurve.from_two_sets((225, 6), (185, 12))
    print("  calibration: 225 x 6  and  185 x 12")
    print("  ->", curve.describe())
    print()
    s = Session(curve, Params(k=0.4, gamma=2.0, t_rec=180.0))
    s.do_set(reps=5, load=135, rest=120)   # warm-up
    s.do_set(reps=6, load=205, rest=180)   # working 1
    s.do_set(reps=6, load=205, rest=0)     # working 2
    _print_trace(s, show_load=True)
    med, lo, hi = rir_interval(curve, reps=6, load=205, capacity=1.0, seed=1)
    print(f"\n  probabilistic RIR for a fresh 205 x 6:  median {med:.1f}  "
          f"(80% interval [{lo:.1f}, {hi:.1f}])")

    # 2) Bodyweight pull-ups -- Case A, fixed load --------------------------
    print("\n" + _rule())
    print("\n[2] BODYWEIGHT PULL-UPS  (Case A: fixed load -> measured fresh max)\n")
    bw = LoadRepCurve.fixed(12)            # one fresh set to failure = 12 reps
    print("  measured fresh max: 12 strict pull-ups   (no load is ever entered)")
    print("  ->", bw.describe())
    print()
    s2 = Session(bw, Params(k=0.4, gamma=2.0, t_rec=180.0))
    s2.do_set(reps=8, rest=120)
    s2.do_set(reps=8, rest=0)
    _print_trace(s2, show_load=False)
    print("\n  same 8-rep set, less reserve the second time -- the battery drained.")

    # 3) Weighted pull-ups -- Case B, variable load -------------------------
    print("\n" + _rule())
    print("\n[3] WEIGHTED PULL-UPS  (Case B: total effective load axis)\n")
    m = 175.0
    b = BODYWEIGHT_FRACTION["pullup"]
    w_body = total_load(m, b, added=0.0)     # bodyweight pull-up
    w_belt = total_load(m, b, added=45.0)    # +45 lb belt
    print(f"  bodyweight {m:.0f} lb, pull-up borne fraction b = {b:.2f}")
    print(f"  total loads: bodyweight = {w_body:.0f},  +45 belt = {w_belt:.0f}")
    curveB = LoadRepCurve.from_two_sets((w_body, 10), (w_belt, 4))
    print("  ->", curveB.describe(),
          f"  (1RM total {curveB.one_rep_max:.0f} = bodyweight + "
          f"{curveB.one_rep_max - m:.0f} lb added)")
    print()
    sB = Session(curveB, Params(k=0.4, gamma=2.0, t_rec=180.0))
    print(f"  predicted RIR, fresh bodyweight set of 6:  "
          f"{sB.predicted_rir(reps=6, load=w_body):.1f}")
    print(f"  predicted RIR, fresh +25 lb set of 4:      "
          f"{sB.predicted_rir(reps=4, load=total_load(m, b, 25)):.1f}")

    print("\n" + _rule())
    print("  Model is a formulation pending empirical validation. Parameters are placeholders.")
    print(_rule())


if __name__ == "__main__":
    _demo()