"""Small-data, session-based top-set forecasts; no future observations in fitting.

The empirical transition model is a recency-weighted nearest-neighbour model.
It learns joint load/rep changes, including repeats and regressions, rather than
converting a fitted E1RM line into an alleged workout. See docs/forecasting.md.
"""
from datetime import date, datetime, timedelta
from functools import lru_cache
from collections import Counter
import math

import numpy as np

BASE = 101.3
REP_FACTOR = 2.67123


def e1rm(weight, reps):
    return np.asarray(weight) * 100 / (BASE - REP_FACTOR * np.asarray(reps))


def sessions(dates, weights, reps):
    """One best set per date, with deterministic ties and fractional loads intact."""
    by_date = {}
    for day, weight, repetitions in zip(dates, weights, reps):
        if isinstance(day, datetime):
            day = day.date()
        try:
            weight, repetitions = float(weight), float(repetitions)
        except (TypeError, ValueError):
            continue
        if (not isinstance(day, date) or not math.isfinite(weight)
                or not math.isfinite(repetitions) or weight <= 0
                or not repetitions.is_integer() or not 1 <= repetitions <= 20):
            continue
        row = (day, weight, int(repetitions))
        if day not in by_date or (float(e1rm(weight, repetitions)), weight, repetitions) > (
                float(e1rm(by_date[day][1], by_date[day][2])), *by_date[day][1:]):
            by_date[day] = row
    return tuple(sorted(by_date.values()))


def profile(rows):
    """Summarise the recent training regime without counting recovery as a PR."""
    recent = rows[-24:]
    weights = np.array([r[1] for r in recent])
    reps = np.array([r[2] for r in recent])
    gaps = np.array([(b[0] - a[0]).days for a, b in zip(recent, recent[1:])])
    dw = np.diff(weights)
    increases = dw[(dw > 0.01) & (dw <= weights[:-1] * 0.15)]
    # Use an observed increment, not an average that invents plate sizes.
    increment = Counter(np.round(increases, 3)).most_common(1)[0][0] if len(increases) else 0.0
    ceiling_samples = reps[:-1][(dw > 0.01) & (dw <= weights[:-1] * 0.15)]
    ceiling = int(round(float(np.median(ceiling_samples)))) if len(ceiling_samples) >= 2 else int(max(reps))
    block = [rows[-1]]
    for row in reversed(rows[:-1]):
        if abs(row[1] - weights[-1]) > 0.01:
            break
        block.insert(0, row)
    best, stalls, run = block[0][2], 0, 0
    for row in block[1:]:
        if row[2] > best:
            best, run = row[2], 0
        else:
            stalls, run = stalls + 1, run + 1
    return {
        'increment': increment, 'ceiling': ceiling,
        'cadence_days': max(1, int(round(float(np.median(gaps[-12:]))))),
        'block_sessions': len(block), 'block_stalls': stalls, 'stall_run': run,
        'block_start': block[0][0], 'block_start_reps': block[0][2],
        'block_best_reps': best, 'recent_sessions': len(recent),
        'recent_load_repeats': int(np.sum(abs(dw) < 0.01)),
        'transitions': len(dw), 'ceiling_evidence': len(ceiling_samples),
    }


def training_blocks(rows):
    """Detect layoffs online using only cadence known before each gap.

    A gap exceeds three usual intervals (at least 28, at most 84 days).
    Sparse starts use 42 days. The cadence restarts with each new block.
    """
    blocks = []
    gaps = []
    for row in rows:
        if not blocks:
            blocks.append([row])
            continue
        gap = (row[0] - blocks[-1][-1][0]).days
        threshold = max(28, min(84, 3 * float(np.median(gaps[-8:])))) if len(gaps) >= 3 else 42
        if gap > threshold:
            blocks.append([row])
            gaps = []
        else:
            blocks[-1].append(row)
            gaps.append(gap)
    return tuple(tuple(block) for block in blocks)


class TransitionModel:
    """Kernel-smoothed joint transitions with an explicit repeat-set prior."""
    def __init__(self, rows):
        self.rows = rows
        self.profile = profile(rows)
        recent = rows[-32:]
        w = np.array([r[1] for r in recent])
        r = np.array([r[2] for r in recent])
        gap = np.array([(b[0] - a[0]).days for a, b in zip(recent, recent[1:])])
        # A long layoff starts a different training block, not an ordinary step.
        keep = gap <= max(28, self.profile['cadence_days'] * 3)
        self.source_w, self.source_r = w[:-1][keep], r[:-1][keep]
        self.dw = np.diff(w)[keep]
        self.dr = np.diff(r)[keep]
        self.reset_r = r[1:][keep]
        self.age = np.arange(len(w) - 2, -1, -1)[keep]
        stall_runs, run, best, load = [], 0, int(r[0]), w[0]
        for index, (weight, reps) in enumerate(zip(w[:-1], r[:-1])):
            if abs(weight - load) > 0.01:
                load, best, run = weight, int(reps), 0
            elif reps > best:
                best, run = int(reps), 0
            elif index:
                run += 1
            stall_runs.append(run)
        self.stalls = np.array(stall_runs)[keep]

    def outcomes(self, weights, reps, stalls, step=1):
        weights, reps, stalls = map(np.atleast_1d, (weights, reps, stalls))
        # Learn from nearby rep counts, loads and stall lengths. Recent sessions
        # have greater influence; the current long plateau outweighs old PRs.
        distance = (abs(reps[:, None] - self.source_r) / 2
                    + abs(np.log(weights[:, None] / self.source_w)) / 0.15
                    + abs(np.minimum(stalls[:, None], 5) - np.minimum(self.stalls, 5)) / 4)
        probabilities = np.exp(-distance) * 0.5 ** (self.age / 8)
        increase = self.dw > 0.01
        decrease = self.dw < -0.01
        probabilities[:, increase] *= (reps >= self.profile['ceiling'])[:, None]
        # Calendar horizons beyond ~12 sessions are scenarios. Damp the rate of
        # further load advances rather than extrapolating beginner gains forever.
        probabilities[:, increase] *= 0.98 ** max(0, step - 12)
        delta_w = np.where(increase, self.profile['increment'], self.dw)
        delta_w = np.maximum(delta_w, -weights[:, None] * 0.15)
        out_w = np.maximum(weights[:, None] + delta_w, 0.5)
        out_r = np.clip(reps[:, None] + self.dr, 1, 20)
        out_r[:, increase | decrease] = np.clip(self.reset_r[increase | decrease], 1, 20)
        # Rep gains below a learned ceiling, not unlimited rep inflation.
        out_r = np.where((self.dw == 0) & (self.dr > 0),
                         np.minimum(out_r, np.maximum(reps[:, None], self.profile['ceiling'])), out_r)
        probabilities = np.column_stack([np.ones(len(weights)), probabilities])
        out_w = np.column_stack([weights, out_w])
        out_r = np.column_stack([reps, out_r])
        probabilities /= probabilities.sum(axis=1, keepdims=True)
        return np.round(out_w, 3), out_r.astype(int), probabilities

    def next_set(self):
        w, r, p = self.outcomes(self.rows[-1][1], self.rows[-1][2], self.profile['stall_run'])
        # Combine identical outcomes before choosing the most likely joint set.
        outcomes = {}
        for weight, reps, probability in zip(w[0], r[0], p[0]):
            key = (float(weight), int(reps))
            outcomes[key] = outcomes.get(key, 0) + float(probability)
        ranked = sorted(outcomes.items(), key=lambda item: -item[1])
        return ranked

    def simulate(self, steps, paths=512):
        rng = np.random.default_rng(20260906)
        weights = np.full(paths, self.rows[-1][1], dtype=float)
        reps = np.full(paths, self.rows[-1][2], dtype=int)
        stalls = np.full(paths, self.profile['stall_run'], dtype=int)
        best_reps = np.full(paths, self.profile['block_best_reps'], dtype=int)
        result = [(weights.copy(), reps.copy())]
        for step in range(1, steps + 1):
            ow, oreps, p = self.outcomes(weights, reps, stalls, step)
            choices = (rng.random(paths)[:, None] > np.cumsum(p, axis=1)).sum(axis=1)
            choices = np.minimum(choices, p.shape[1] - 1)
            nw, nr = ow[np.arange(paths), choices], oreps[np.arange(paths), choices]
            changed = abs(nw - weights) > 0.01
            stalls = np.where(changed | (nr > best_reps), 0, stalls + 1)
            best_reps = np.where(changed, nr, np.maximum(best_reps, nr))
            weights, reps = nw, nr
            result.append((weights.copy(), reps.copy()))
        return result


def benchmark(rows, method, steps=1):
    """Forecast-only baselines: repeat last set or a damped recent E1RM trend."""
    weight, reps = rows[-1][1:]
    if method == 'repeat':
        return weight, reps
    sample = rows if method == 'linear' else rows[-12:]
    x = np.array([(r[0] - sample[0][0]).days for r in sample])
    y = e1rm([r[1] for r in sample], [r[2] for r in sample])
    slope = float(np.polyfit(x, y, 1)[0])
    days = profile(rows)['cadence_days']
    advance = steps if method == 'linear' else sum(0.9 ** i for i in range(1, steps + 1))
    value = max(0, float(y[-1]) + slope * days * advance)
    load = value * (BASE - REP_FACTOR * reps) / 100
    increment = profile(rows)['increment'] or 0.5
    return max(0.5, round(load / increment) * increment), reps


def backtest(rows, limit=20):
    """Expanding-window one-session holdouts. No random train/test split."""
    errors = {name: [] for name in ('transitions', 'repeat', 'damped', 'linear')}
    for end in range(max(8, len(rows) - limit), len(rows)):
        train, actual = rows[:end], rows[end]
        if (actual[0] - train[-1][0]).days > max(28, profile(train)['cadence_days'] * 3):
            continue
        predicted = {'transitions': TransitionModel(train).next_set()[0][0]}
        predicted.update({name: benchmark(train, name) for name in errors if name != 'transitions'})
        for name, (weight, reps) in predicted.items():
            errors[name].append((abs(weight - actual[1]), abs(reps - actual[2]),
                                 abs(float(e1rm(weight, reps) - e1rm(actual[1], actual[2])))))
    return {name: {'count': len(values),
                   'weight_mae': round(float(np.mean(values, axis=0)[0]), 2),
                   'reps_mae': round(float(np.mean(values, axis=0)[1]), 2),
                   'e1rm_mae': round(float(np.mean(values, axis=0)[2]), 2),
                   'error_p80': round(float(np.quantile(np.array(values)[:, 2], .8)), 2)}
            for name, values in errors.items() if values}


def _point(day, weight, reps, lower, upper, session_index):
    return {'date': day, 'value': round(float(e1rm(weight, reps)), 1),
            'lower': round(float(lower), 1), 'upper': round(float(upper), 1),
            'projected_weight': float(weight), 'projected_repetitions': int(reps),
            'set_display': f'{weight:g} kg × {int(reps)}', 'session': session_index}


def representative_set(weights, reps):
    """Median load, then conditional median reps at that load.

    Picking by E1RM alone can jump between distant load/rep combinations with
    almost identical strength. Use the actual joint load/rep distribution.
    """
    loads = np.unique(weights)
    weight = float(loads[np.argmin(abs(loads - np.median(weights)))])
    conditional = reps[weights == weight]
    observed_reps = np.unique(conditional)
    repetitions = int(observed_reps[np.argmin(abs(observed_reps - np.median(conditional)))])
    return weight, repetitions


@lru_cache(maxsize=128)
def _forecast(rows, horizon_weeks):
    p = profile(rows)
    model = TransitionModel(rows)
    cadence = p['cadence_days']
    steps = math.ceil(horizon_weeks * 7 / cadence)
    simulations = model.simulate(steps)
    validation = backtest(rows)
    # Keep the displayed scenario band at least as broad as recent held-out
    # errors. This is an empirical safeguard, not a calibrated 80% guarantee.
    error_floor = validation.get('transitions', {}).get('error_p80', 0)
    points = []
    for index, (weights, reps) in enumerate(simulations):
        if index * cadence > horizon_weeks * 7:
            break
        values = e1rm(weights, reps)
        lower, median, upper = np.quantile(values, [0.1, 0.5, 0.9])
        if index:
            lower, upper = max(0, min(lower, median - error_floor)), max(upper, median + error_floor)
        weight, repetitions = representative_set(weights, reps)
        points.append(_point(rows[-1][0] + timedelta(days=index * cadence), weight, repetitions,
                             min(lower, float(e1rm(weight, repetitions))),
                             max(upper, float(e1rm(weight, repetitions))), index))
    # Extend the chart to the selected calendar horizon without inventing a
    # workout between the learned session dates.
    end_date = rows[-1][0] + timedelta(weeks=horizon_weeks)
    chart_points = list(points)
    if points[-1]['date'] != end_date:
        chart_points.append(dict(points[-1], date=end_date))
    ranking = model.next_set()
    (weight, repetitions), probability = ranking[0]
    next_values = e1rm(simulations[1][0], simulations[1][1])
    next_point = _point(rows[-1][0] + timedelta(days=cadence), weight, repetitions,
                        min(np.quantile(next_values, .1), e1rm(weight, repetitions)),
                        max(np.quantile(next_values, .9), e1rm(weight, repetitions)), 1)
    next_point['lower'] = round(max(0, min(next_point['lower'], next_point['value'] - error_floor)), 1)
    next_point['upper'] = round(max(next_point['upper'], next_point['value'] + error_floor), 1)
    if len(points) > 1:
        points[1].update(next_point)
    return {
        'available': True, 'metric': 'E1RM', 'model': 'transitions',
        'horizon_weeks': horizon_weeks, 'observed_points': len(rows),
        'last_date': rows[-1][0], 'last_value': points[0]['value'],
        'projected_value': points[-1]['value'],
        'projected_change': round(points[-1]['value'] - points[0]['value'], 1),
        'weekly_change': round((points[-1]['value'] - points[0]['value']) / horizon_weeks, 1),
        'end_date': end_date, 'points': chart_points, 'session_points': points[1:],
        'next_session': next_point, 'profile': p, 'validation': validation,
        'alternatives': [{'weight': w, 'reps': r, 'probability': round(prob * 100)}
                         for (w, r), prob in ranking[:3]],
        'long_range': horizon_weeks > 12,
    }


def build_forecast(values, dates, horizon_weeks=8, metric='E1RM', weights=None, repetitions=None):
    # Copy the cached result because the rendering layer adds equivalent sets.
    from copy import deepcopy
    try:
        horizon_weeks = max(4, min(52, int(horizon_weeks)))
    except (TypeError, ValueError):
        horizon_weeks = 8
    rows = sessions(dates, weights if weights is not None else [],
                    repetitions if repetitions is not None else [])
    if len(rows) < 4:
        return {'available': False, 'reason': 'Add at least 4 distinct workout dates with valid sets (1–20 reps).'}
    dated = [d.date() if isinstance(d, datetime) else d for d in dates if isinstance(d, date)]
    if dated and max(dated) > rows[-1][0]:
        return {'available': False, 'reason': 'The latest reading needs a positive load and 1–20 reps for this forecast.'}
    blocks = training_blocks(rows)
    rows = blocks[-1]
    if len(rows) < 4:
        return {'available': False, 'reason': 'Add 4 sessions after the recent training break to learn the new pattern.'}
    result = deepcopy(_forecast(rows, horizon_weeks))
    result['training_blocks'] = len(blocks)
    result['training_start'] = rows[0][0]
    result['stale'] = (date.today() - rows[-1][0]).days > max(28, result['profile']['cadence_days'] * 3)
    return result