"""Simulator behaviour: deterministic short-horizon checks, then
stochastic monotonicity + cFIREsim sanity calibration."""
from __future__ import annotations

import time

import numpy as np

from fire_planner.glide_path import static
from fire_planner.returns.bootstrap import block_bootstrap
from fire_planner.returns.shiller import ReturnsBundle, synthetic_returns
from fire_planner.simulator import default_bucket_split, simulate
from fire_planner.strategies.guyton_klinger import GuytonKlingerStrategy
from fire_planner.strategies.trinity import TrinityStrategy
from fire_planner.strategies.vpw import VpwStrategy
from fire_planner.tax.bulgaria import BulgariaTaxRegime
from fire_planner.tax.malaysia import MalaysiaTaxRegime
from fire_planner.tax.uk import UkTaxRegime


def fixed_paths(n_paths: int, n_years: int, stock_ret: float, bond_ret: float,
                cpi: float) -> np.ndarray:
    """All-paths-identical returns — deterministic regression check."""
    out = np.zeros((n_paths, n_years, 3), dtype=np.float64)
    out[..., 0] = stock_ret
    out[..., 1] = bond_ret
    out[..., 2] = cpi
    return out


def test_simulate_zero_returns_zero_inflation_drains_at_4pc() -> None:
    """0% returns + 0% inflation, 4% Trinity, 25y horizon — withdraw
    £40k/y from £1M = drain to exactly £0 in year 25. Success because
    portfolio stays positive *during* every year (clipped to 0 at end)."""
    paths = fixed_paths(n_paths=1, n_years=25, stock_ret=0.0, bond_ret=0.0, cpi=0.0)
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.6),
        strategy=TrinityStrategy(initial_rate=0.04),
        regime=MalaysiaTaxRegime(),  # 0% to keep arithmetic clean
    )
    # Year 0 withdrawal is 40k, portfolio after = 960k
    assert res.portfolio_real[0, 1] == 960_000.0
    # 25y of £40k draws against zero growth = drain to 0 by end of y25.
    assert abs(res.portfolio_real[0, 25]) < 1.0


def test_simulate_failing_path_marked_unsuccessful() -> None:
    """6% Trinity rate against 0% real return for 25y — clearly fails."""
    paths = fixed_paths(n_paths=1, n_years=25, stock_ret=0.0, bond_ret=0.0, cpi=0.0)
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=60_000.0,
        glide=static(1.0),
        strategy=TrinityStrategy(initial_rate=0.06),
        regime=MalaysiaTaxRegime(),
    )
    assert not res.success_mask[0]


def test_simulate_growing_portfolio_succeeds() -> None:
    """5% real return, 4% draw — classic surplus case."""
    paths = fixed_paths(n_paths=1, n_years=30, stock_ret=0.05, bond_ret=0.05, cpi=0.0)
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(1.0),
        strategy=TrinityStrategy(initial_rate=0.04),
        regime=MalaysiaTaxRegime(),
    )
    assert res.success_mask[0]
    # Portfolio should grow above starting value
    assert res.portfolio_real[0, 30] > 1_000_000.0


def test_savings_phase_increases_portfolio() -> None:
    """5y of savings @ £100k / 0% return → portfolio grows."""
    paths = fixed_paths(n_paths=1, n_years=5, stock_ret=0.0, bond_ret=0.0, cpi=0.0)
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=0.0,  # not drawing during accumulation
        glide=static(1.0),
        strategy=TrinityStrategy(initial_rate=0.0),
        regime=MalaysiaTaxRegime(),
        annual_savings=np.full(5, 100_000.0),
    )
    # 1M + 5×100k = 1.5M, no growth
    assert res.portfolio_real[0, 5] == 1_500_000.0


def test_uk_tax_increases_failure_rate_vs_no_tax() -> None:
    """Same scenario, UK regime should produce more or equal failures
    than Malaysia (zero tax) — paths are identical."""
    bundle = synthetic_returns(seed=1, n_years=120)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=200, n_years=30, block_size=5, rng=rng)
    common = dict(
        paths=paths,
        initial_portfolio=600_000.0,  # tighter so tax matters
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(initial_rate=0.04),
    )
    msy = simulate(**common, regime=MalaysiaTaxRegime())  # type: ignore[arg-type]
    uk = simulate(**common, regime=UkTaxRegime())  # type: ignore[arg-type]
    assert uk.success_rate <= msy.success_rate
    assert uk.median_lifetime_tax() > msy.median_lifetime_tax()


def test_vpw_never_runs_out() -> None:
    """VPW scales withdrawal with portfolio — should never fully ruin."""
    bundle = synthetic_returns(seed=2, n_years=120)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=200, n_years=60, block_size=5, rng=rng)
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=50_000.0,
        glide=static(0.7),
        strategy=VpwStrategy(),
        regime=MalaysiaTaxRegime(),
    )
    # Every path should keep some portfolio > 0 throughout (until last year).
    # Year `n-1` end may be tiny but >= 0 since VPW caps drain at 100% in y=H-1.
    assert res.portfolio_real[:, 1:-1].min() > 0


def test_simulator_deterministic_with_same_paths() -> None:
    paths = fixed_paths(n_paths=10, n_years=30, stock_ret=0.05, bond_ret=0.03, cpi=0.02)
    common = dict(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=GuytonKlingerStrategy(),
        regime=BulgariaTaxRegime(),
    )
    a = simulate(**common)  # type: ignore[arg-type]
    b = simulate(**common)  # type: ignore[arg-type]
    np.testing.assert_array_equal(a.portfolio_real, b.portfolio_real)


def test_success_rate_monotone_in_portfolio() -> None:
    """More starting wealth → higher (or equal) success rate."""
    bundle = synthetic_returns(seed=3, n_years=120)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=300, n_years=30, block_size=5, rng=rng)
    common = dict(
        paths=paths,
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(initial_rate=0.04),
        regime=MalaysiaTaxRegime(),
    )
    low = simulate(**common, initial_portfolio=600_000.0)  # type: ignore[arg-type]
    high = simulate(**common, initial_portfolio=1_500_000.0)  # type: ignore[arg-type]
    assert high.success_rate >= low.success_rate


def test_success_rate_monotone_in_spending() -> None:
    """Less spending → higher success rate."""
    bundle = synthetic_returns(seed=4, n_years=120)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=300, n_years=30, block_size=5, rng=rng)
    common = dict(
        paths=paths,
        initial_portfolio=1_000_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(initial_rate=0.04),
        regime=MalaysiaTaxRegime(),
    )
    cheap = simulate(**common, spending_target=30_000.0)  # type: ignore[arg-type]
    fat = simulate(**common, spending_target=80_000.0)  # type: ignore[arg-type]
    assert cheap.success_rate >= fat.success_rate


def test_fan_quantiles_shape() -> None:
    bundle = synthetic_returns(seed=5, n_years=120)
    paths = block_bootstrap(bundle,
                            n_paths=100,
                            n_years=20,
                            block_size=5,
                            rng=np.random.default_rng(0))
    res = simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(),
        regime=MalaysiaTaxRegime(),
    )
    p10 = res.fan_quantiles(10)
    assert p10.shape == (21, )  # n_years + 1


def test_perf_under_60s_for_10k_paths_60y() -> None:
    """Stretch goal — at 10k paths × 60y the simulator should finish
    in well under a minute on commodity hardware. Test allows 60s
    (generous; CI can vary)."""
    bundle = synthetic_returns(seed=6, n_years=150)
    paths = block_bootstrap(bundle,
                            n_paths=10_000,
                            n_years=60,
                            block_size=5,
                            rng=np.random.default_rng(0))
    start = time.perf_counter()
    simulate(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(),
        regime=MalaysiaTaxRegime(),
    )
    elapsed = time.perf_counter() - start
    assert elapsed < 60, f"too slow: {elapsed:.2f}s"


def test_convergence_5k_vs_50k_paths() -> None:
    """Success rate should be stable to within ±1.5% between 5k and
    50k paths (Monte Carlo SE ~0.5% at 10k samples)."""
    bundle = synthetic_returns(seed=7, n_years=150)
    paths_small = block_bootstrap(bundle,
                                  n_paths=5_000,
                                  n_years=30,
                                  block_size=5,
                                  rng=np.random.default_rng(0))
    paths_large = block_bootstrap(bundle,
                                  n_paths=50_000,
                                  n_years=30,
                                  block_size=5,
                                  rng=np.random.default_rng(0))
    common = dict(
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.7),
        strategy=TrinityStrategy(),
        regime=MalaysiaTaxRegime(),
    )
    small = simulate(paths=paths_small, **common)  # type: ignore[arg-type]
    large = simulate(paths=paths_large, **common)  # type: ignore[arg-type]
    assert abs(small.success_rate - large.success_rate) < 0.015


def test_default_bucket_split_smoke() -> None:
    inputs = default_bucket_split(50_000.0, year_idx=5)
    assert inputs.capital_gains == 50000


def test_returns_bundle_supplies_ie_data_columns() -> None:
    """Sanity: the bundle has stock/bond/cpi correctly aligned."""
    b = synthetic_returns(seed=8, n_years=10)
    assert isinstance(b, ReturnsBundle)
    assert len(b.stock_nominal) == 10