#!/usr/bin/env python3
"""Original synthetic mean-reversion experiment for QuantSteps.

Run: python mean-reversion.py
Tested on CPython 3.14.0; standard library only, no pip dependencies.
Quotes 0..799 train parameters. Quotes 800..1199 are held out.
Signal at close t-1 fills at close t and earns only t..t+1 P&L.
One fixed unit; $100 initial research capital; arithmetic P&L / capital.
The final unit is liquidated at the last quote with its cost included.
This is a synthetic accounting exercise, not historical market evidence.
"""

import csv
import json
import math
import random
import statistics
from pathlib import Path

SEED = 20261003
N = 1200
SPLIT = 800
CAPITAL = 100.0
ENTRY_Z = 1.5
EXIT_Z = 0.25
STOP_Z = 3.0
MAX_HOLD = 20


def make_paths():
    rng = random.Random(SEED)
    shocks = [rng.gauss(0.0, 0.6) for _ in range(N - 1)]
    stable = [100.0]
    broken = [100.0]
    for t, shock in enumerate(shocks, 1):
        stable.append(100.0 + 0.9 * (stable[-1] - 100.0) + shock)
        if t < SPLIT:
            broken.append(stable[-1])
        else:
            broken.append(broken[-1] + 0.12 + shock)
    return stable, broken


def fit_training(prices):
    train = prices[:SPLIT]
    lag, current = train[:-1], train[1:]
    xbar, ybar = statistics.mean(lag), statistics.mean(current)
    phi = sum((x - xbar) * (y - ybar) for x, y in zip(lag, current))
    phi /= sum((x - xbar) ** 2 for x in lag)
    intercept = ybar - phi * xbar
    if not 0.0 < phi < 1.0:
        raise ValueError("Training fit does not support monotone AR(1) decay")
    return {
        "phi": phi,
        "mu": intercept / (1.0 - phi),
        "scale": statistics.stdev(train),
        "half_life": math.log(0.5) / math.log(phi),
    }


def run_backtest(prices, fitted, cost_bps):
    unit, age = 0, 0
    gross, costs, turnover = 0.0, 0.0, 0.0
    entries, exits, exposed = 0, 0, 0
    equity, peak, drawdown = CAPITAL, CAPITAL, 0.0
    rows = []
    for t in range(SPLIT, len(prices) - 1):
        # Only yesterday's quote enters today's execution decision.
        z = (prices[t - 1] - fitted["mu"]) / fitted["scale"]
        old = unit
        if unit:
            converged = (unit == 1 and z >= -EXIT_Z) or (unit == -1 and z <= EXIT_Z)
            adverse = (unit == 1 and z <= -STOP_Z) or (unit == -1 and z >= STOP_Z)
            if converged or adverse or age >= MAX_HOLD:
                unit, age = 0, 0
        elif ENTRY_Z <= abs(z) < STOP_Z:
            unit = -1 if z > 0.0 else 1
            age = 0
        # No same-close reversal/re-entry after an exit.
        if old == 0 and unit != 0:
            entries += 1
        if old != 0 and unit == 0:
            exits += 1
        traded_notional = abs(unit - old) * prices[t]
        fee = traded_notional * cost_bps / 10000.0
        pnl = unit * (prices[t + 1] - prices[t])
        gross += pnl
        costs += fee
        turnover += traded_notional / CAPITAL
        # Execution cost creates an equity mark before the next holding P&L.
        equity -= fee
        peak = max(peak, equity)
        drawdown = max(drawdown, peak - equity)
        equity += pnl
        peak = max(peak, equity)
        drawdown = max(drawdown, peak - equity)
        if unit:
            age += 1
            exposed += 1
        rows.append({
            "execution_index": t, "signal_index": t - 1,
            "price": prices[t], "z_from_previous_close": z,
            "units": unit, "gross_pnl": pnl, "cost": fee,
        })
    liquidation = abs(unit) * prices[-1]
    costs += liquidation * cost_bps / 10000.0
    turnover += liquidation / CAPITAL
    equity -= liquidation * cost_bps / 10000.0
    drawdown = max(drawdown, peak - equity)
    exits += int(unit != 0)
    assert entries == exits, "Every entry must have a completed exit"
    assert math.isclose(costs / CAPITAL, turnover * cost_bps / 10000.0, abs_tol=1e-12)
    assert math.isclose(equity, CAPITAL + gross - costs, abs_tol=1e-10)
    return {
        "cost_bps": cost_bps, "gross_pct": gross / CAPITAL * 100.0,
        "cost_pct": costs / CAPITAL * 100.0,
        "net_pct": (gross - costs) / CAPITAL * 100.0,
        "turnover_x": turnover, "round_trips": entries,
        "exposed_pct": exposed / len(rows) * 100.0,
        "max_drawdown_pct_initial": drawdown / CAPITAL * 100.0,
    }, rows


def checks(stable, broken, fitted):
    assert stable[:SPLIT] == broken[:SPLIT]
    assert fit_training(broken) == fitted
    result, rows = run_backtest(stable, fitted, 5.0)
    changed = stable.copy()
    changed[1000:] = [p + 100.0 for p in changed[1000:]]
    _, changed_rows = run_backtest(changed, fitted, 5.0)
    # A future perturbation cannot alter earlier decisions or earlier P&L.
    assert rows[:199] == changed_rows[:199]
    assert [r["units"] for r in rows[:201]] == [r["units"] for r in changed_rows[:201]]
    assert all(r["signal_index"] < r["execution_index"] for r in rows)
    # Hand-check the first interval's fixed-unit accounting.
    first = rows[0]
    expected = first["units"] * (stable[SPLIT + 1] - stable[SPLIT])
    assert math.isclose(first["gross_pnl"], expected, abs_tol=1e-12)
    assert math.isclose(result["net_pct"], result["gross_pct"] - result["cost_pct"], abs_tol=1e-12)
    # A profitable next interval must not conceal the entry-fee equity dip.
    fee_dip_prices = [100.0] * 799 + [150.0, 150.0, 100.0, 100.0]
    fee_dip, _ = run_backtest(fee_dip_prices, {"mu": 100.0, "scale": 30.0}, 100.0)
    assert math.isclose(fee_dip["gross_pct"], 50.0, abs_tol=1e-12)
    assert math.isclose(fee_dip["cost_pct"], 2.5, abs_tol=1e-12)
    assert math.isclose(fee_dip["max_drawdown_pct_initial"], 1.5, abs_tol=1e-12)


def main():
    stable, broken = make_paths()
    fitted = fit_training(stable)
    checks(stable, broken, fitted)
    scenarios = []
    for name, path, cost in [
        ("Stable AR(1)", stable, 5.0),
        ("Stable AR(1), cost stress", stable, 50.0),
        ("Stable AR(1), extreme cost stress", stable, 100.0),
        ("Random walk with drift after split", broken, 5.0),
    ]:
        stats, rows = run_backtest(path, fitted, cost)
        scenarios.append({"scenario": name, **stats})
    print("Synthetic example; arithmetic P&L / $100 initial capital")
    print("Training fit:", json.dumps(fitted, sort_keys=True))
    print("scenario | cost bps | gross % | cost % | net % | turnover x | round trips | DD % initial")
    for s in scenarios:
        print(f"{s['scenario']} | {s['cost_bps']:.0f} | {s['gross_pct']:.4f} | "
              f"{s['cost_pct']:.4f} | {s['net_pct']:.4f} | {s['turnover_x']:.4f} | "
              f"{s['round_trips']} | {s['max_drawdown_pct_initial']:.4f}")
    print("Checks passed: training isolation, future perturbation, execution lag, costs, fee-dip drawdown, liquidation.")
    output = Path.cwd()
    (output / "mean-reversion-results.json").write_text(
        json.dumps({"synthetic": True, "seed": SEED, "training": fitted, "scenarios": scenarios}, indent=2),
        encoding="utf-8",
    )
    with (output / "mean-reversion-quotes.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["index", "stable_price", "regime_break_price", "sample"])
        for i, (a, b) in enumerate(zip(stable, broken)):
            writer.writerow([i, a, b, "train" if i < SPLIT else "holdout"])


if __name__ == "__main__":
    main()