"""Tests for fire_planner.goals_eval — parametrised over goal kinds."""
from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal

import numpy as np
import pytest

from fire_planner.goals_eval import evaluate_goals
from fire_planner.simulator import SimulationResult


@dataclass
class _Goal:
    kind: str
    name: str
    target_amount_gbp: Decimal | None = None
    target_year: int | None = None
    comparator: str = ">="
    success_threshold: Decimal = Decimal("0.95")
    enabled: bool = True


def _make_result(
    portfolio_paths: list[list[float]],
    withdrawal_paths: list[list[float]] | None = None,
) -> SimulationResult:
    """Build a SimulationResult from explicit per-path arrays."""
    portfolio = np.asarray(portfolio_paths, dtype=np.float64)
    n_paths, ncols = portfolio.shape
    n_years = ncols - 1
    if withdrawal_paths is None:
        wd = np.zeros((n_paths, n_years), dtype=np.float64)
    else:
        wd = np.asarray(withdrawal_paths, dtype=np.float64)
    tax = np.zeros((n_paths, n_years), dtype=np.float64)
    success_mask = portfolio[:, 1:-1].min(axis=1) > 0.0 if ncols >= 3 else np.ones(
        n_paths, dtype=bool)
    return SimulationResult(
        portfolio_real=portfolio,
        withdrawal_real=wd,
        tax_real=tax,
        success_mask=success_mask,
    )


def test_target_nw_by_year_exact_count() -> None:
    # 4 paths, 3 years. At year 2: [200, 1500, 2500, 3000]. Target ≥ £2M
    # → 2/4 hit → probability 0.5.
    portfolio = [
        [1000, 500, 200, 100],
        [1000, 1200, 1500, 1700],
        [1000, 2000, 2500, 2800],
        [1000, 2400, 3000, 3500],
    ]
    result = _make_result(portfolio)
    goal = _Goal(kind="target_nw_by_year",
                 name="≥ £2M at y2",
                 target_amount_gbp=Decimal("2000"),
                 target_year=2,
                 comparator=">=",
                 success_threshold=Decimal("0.4"))
    [eval_] = evaluate_goals(result, [goal])
    assert eval_.probability == pytest.approx(0.5)
    assert eval_.passed is True


def test_never_run_out_full_horizon() -> None:
    # 4 paths over 4 years. Path 0 hits 0 at year 2. Path 1 hits 0 at
    # year 3. Path 2 + 3 stay positive throughout.
    portfolio = [
        [1000, 500, 0, 0, 0],
        [1000, 800, 600, 0, 0],
        [1000, 900, 800, 700, 600],
        [1000, 1100, 1200, 1300, 1400],
    ]
    result = _make_result(portfolio)
    goal = _Goal(kind="never_run_out",
                 name="don't ruin",
                 target_year=None,
                 success_threshold=Decimal("0.5"))
    [eval_] = evaluate_goals(result, [goal])
    assert eval_.probability == pytest.approx(0.5)
    assert eval_.passed is True


def test_target_real_income_uses_path_median() -> None:
    portfolio = [
        [1000, 1000, 1000],
        [1000, 1000, 1000],
        [1000, 1000, 1000],
    ]
    withdrawals = [
        [40_000, 40_000],
        [60_000, 60_000],
        [80_000, 80_000],
    ]
    result = _make_result(portfolio, withdrawals)
    goal = _Goal(kind="target_real_income",
                 name="≥ £50k income",
                 target_amount_gbp=Decimal("50000"),
                 target_year=0,
                 comparator=">=",
                 success_threshold=Decimal("0.5"))
    [eval_] = evaluate_goals(result, [goal])
    assert eval_.probability == pytest.approx(2 / 3)
    assert eval_.passed is True


def test_disabled_goals_skipped() -> None:
    portfolio = [[1000, 500, 0]]
    result = _make_result(portfolio)
    enabled = _Goal(kind="never_run_out", name="active")
    disabled = _Goal(kind="never_run_out", name="muted", enabled=False)
    evals = evaluate_goals(result, [enabled, disabled])
    assert [e.name for e in evals] == ["active"]


def test_unknown_kind_returns_zero() -> None:
    portfolio = [[1000, 1500, 2000]]
    result = _make_result(portfolio)
    goal = _Goal(kind="not_implemented", name="???")
    [eval_] = evaluate_goals(result, [goal])
    assert eval_.probability == 0.0
    assert eval_.passed is False