"""Returns loader + bootstrap behaviour."""
from pathlib import Path

import numpy as np
import pytest

from fire_planner.returns.bootstrap import block_bootstrap
from fire_planner.returns.shiller import ReturnsBundle, load_from_csv, synthetic_returns


def test_synthetic_returns_shape() -> None:
    b = synthetic_returns(seed=1, n_years=120)
    assert b.n_years == 120
    assert b.stock_nominal.shape == (120, )
    assert b.years[0] == 1871


def test_synthetic_deterministic_for_seed() -> None:
    a = synthetic_returns(seed=42, n_years=10)
    b = synthetic_returns(seed=42, n_years=10)
    np.testing.assert_array_equal(a.stock_nominal, b.stock_nominal)


def test_real_return_smoke() -> None:
    b = ReturnsBundle(
        years=np.array([2020], dtype=np.int32),
        stock_nominal=np.array([0.10]),
        bond_nominal=np.array([0.04]),
        cpi=np.array([0.03]),
    )
    # (1.10 / 1.03) - 1 ≈ 0.06796
    assert abs(b.stock_real()[0] - 0.06796116505) < 1e-9


def test_load_from_csv(tmp_path: Path) -> None:
    csv_path = tmp_path / "returns.csv"
    csv_path.write_text("year,stock_nominal_return,bond_nominal_return,cpi_inflation\n"
                        "1990,0.05,0.07,0.025\n"
                        "1991,-0.10,0.04,0.03\n")
    b = load_from_csv(csv_path)
    assert b.n_years == 2
    assert b.stock_nominal[1] == pytest.approx(-0.10)
    assert b.cpi[0] == pytest.approx(0.025)


def test_returns_bundle_rejects_mismatched_lengths() -> None:
    with pytest.raises(ValueError):
        ReturnsBundle(
            years=np.array([2020, 2021], dtype=np.int32),
            stock_nominal=np.array([0.1]),
            bond_nominal=np.array([0.04, 0.05]),
            cpi=np.array([0.03, 0.025]),
        )


def test_returns_bundle_rejects_empty() -> None:
    with pytest.raises(ValueError):
        ReturnsBundle(
            years=np.array([], dtype=np.int32),
            stock_nominal=np.array([]),
            bond_nominal=np.array([]),
            cpi=np.array([]),
        )


def test_bootstrap_shape() -> None:
    bundle = synthetic_returns(seed=1, n_years=150)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=100, n_years=60, block_size=5, rng=rng)
    assert paths.shape == (100, 60, 3)


def test_bootstrap_deterministic_with_seed() -> None:
    bundle = synthetic_returns(seed=1, n_years=150)
    a = block_bootstrap(bundle, n_paths=50, n_years=30, block_size=5, rng=np.random.default_rng(0))
    b = block_bootstrap(bundle, n_paths=50, n_years=30, block_size=5, rng=np.random.default_rng(0))
    np.testing.assert_array_equal(a, b)


def test_bootstrap_block_size_one_is_iid() -> None:
    """Block size 1 reduces to simple IID resampling — covariance
    structure isn't preserved, but all draws come from the source."""
    bundle = synthetic_returns(seed=2, n_years=100)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=10, n_years=20, block_size=1, rng=rng)
    src_set = set(zip(bundle.stock_nominal, bundle.bond_nominal, bundle.cpi, strict=True))
    drawn_set = set((float(s), float(b), float(c)) for path in paths for s, b, c in path)
    assert drawn_set <= src_set


def test_bootstrap_preserves_block_runs() -> None:
    """For block_size=5, every consecutive 5-year run within a path
    must equal a 5-year run from the source (mod circular)."""
    bundle = synthetic_returns(seed=3, n_years=50)
    rng = np.random.default_rng(0)
    paths = block_bootstrap(bundle, n_paths=5, n_years=15, block_size=5, rng=rng)
    src = np.stack([bundle.stock_nominal, bundle.bond_nominal, bundle.cpi], axis=-1)
    src_n = src.shape[0]
    for path in paths:
        for block_start in range(0, 15, 5):
            block = path[block_start:block_start + 5]
            # Find this block in source by matching the first row, then
            # checking consecutiveness (mod circular).
            for src_idx in range(src_n):
                circ_block = np.stack([src[(src_idx + i) % src_n] for i in range(5)])
                if np.allclose(block, circ_block):
                    break
            else:
                raise AssertionError(f"block {block_start} not a circular slice of source")


def test_bootstrap_rejects_zero_block_size() -> None:
    bundle = synthetic_returns(seed=1, n_years=30)
    with pytest.raises(ValueError):
        block_bootstrap(bundle, n_paths=10, n_years=10, block_size=0)


def test_bootstrap_n_years_not_multiple_of_block() -> None:
    """13 years from 5-year blocks: 3 blocks then truncate to 13."""
    bundle = synthetic_returns(seed=1, n_years=50)
    paths = block_bootstrap(bundle,
                            n_paths=4,
                            n_years=13,
                            block_size=5,
                            rng=np.random.default_rng(0))
    assert paths.shape == (4, 13, 3)