"""Solve each Case's FIRE number — the smallest liquid net worth at which a
Guyton-Klinger plan reaches the bar (default 99%).

The search is a binary search over ``initial_portfolio``: success is monotone in
starting capital because the GK year-0 draw is an absolute real amount, so more
seed always means a lower withdrawal rate and never a worse outcome.

Cashflows layered onto the run via the existing simulator hooks:
- pension unlock — a grown lump that becomes available at age 57
- kids ramp      — an essential per-year outflow over the child-rearing window
- home purchase  — an optional one-time outflow

See ADR-0001 for why we seed on liquid NW and model the pension as a tranche.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

from fire_planner.glide_path import get as get_glide
from fire_planner.life_events import EventInput, events_to_cashflow_array
from fire_planner.scenarios import _JURISDICTION_CONSTRUCTORS
from fire_planner.simulator import RegimeFn, constant_regime, simulate
from fire_planner.spend_model import Case
from fire_planner.strategies.guyton_klinger import GuytonKlingerStrategy

DEFAULT_BAR = 0.99
DEFAULT_PENSION_REAL_GROWTH = 0.03
DEFAULT_HI_GBP = 5_000_000.0


@dataclass(frozen=True)
class TargetInputs:
    """Everything the solver needs for one (Case × country × with-home) target."""
    case: Case
    country_slug: str
    country_display: str
    jurisdiction: str
    annual_spend_gbp: float
    horizon_years: int
    glide_name: str = "rising"
    bar: float = DEFAULT_BAR
    # Pension: locked tranche that joins at age 57.
    pension_now_gbp: float = 0.0
    pension_real_growth: float = DEFAULT_PENSION_REAL_GROWTH
    years_to_pension: int = 0
    # Kids: essential per-year outflow (Family only).
    kids_annual_gbp: float = 0.0
    kids_start_year: int = 5
    kids_end_year: int = 22
    # Optional one-time home purchase.
    with_home: bool = False
    home_amount_gbp: float = 0.0
    home_year: int = 0


@dataclass(frozen=True)
class SolveResult:
    target_nw_gbp: float
    success_at_target: float
    pension_at_unlock_gbp: float
    reached_bar: bool


def pension_at_unlock(inp: TargetInputs) -> float:
    """Current pension value compounded at the assumed real rate to age 57."""
    if inp.pension_now_gbp <= 0:
        return 0.0
    years = max(0, inp.years_to_pension)
    return inp.pension_now_gbp * (1.0 + inp.pension_real_growth) ** years


def build_cashflows(inp: TargetInputs, horizon: int) -> npt.NDArray[np.float64]:
    """Per-year real-GBP cashflow array (pension inflow, kids/home outflows)."""
    events: list[EventInput] = []
    p_at = pension_at_unlock(inp)
    if p_at > 0 and 0 <= inp.years_to_pension < horizon:
        events.append(EventInput(year_start=inp.years_to_pension, one_time_amount_gbp=p_at))
    if inp.kids_annual_gbp > 0:
        events.append(EventInput(
            year_start=inp.kids_start_year,
            year_end=inp.kids_end_year,
            delta_gbp_per_year=-inp.kids_annual_gbp,
        ))
    if inp.with_home and inp.home_amount_gbp > 0:
        events.append(EventInput(year_start=inp.home_year,
                                 one_time_amount_gbp=-inp.home_amount_gbp))
    return events_to_cashflow_array(events, horizon)


def _regime(jurisdiction: str) -> RegimeFn:
    cls = _JURISDICTION_CONSTRUCTORS.get(jurisdiction)
    if cls is None:
        raise KeyError(f"Unknown jurisdiction: {jurisdiction!r}")
    return constant_regime(cls())


def success_at_nw(
    paths: npt.NDArray[np.float64],
    initial: float,
    inp: TargetInputs,
    cashflows: npt.NDArray[np.float64],
) -> float:
    """Success rate of the GK plan seeded at ``initial`` liquid net worth."""
    result = simulate(
        paths=paths,
        initial_portfolio=initial,
        spending_target=inp.annual_spend_gbp,
        glide=get_glide(inp.glide_name),
        strategy=GuytonKlingerStrategy(),
        regime=_regime(inp.jurisdiction),
        horizon_years=inp.horizon_years,
        cashflow_adjustments=cashflows,
    )
    return result.success_rate


def solve_target_nw(
    paths: npt.NDArray[np.float64],
    inp: TargetInputs,
    *,
    lo: float = 0.0,
    hi: float = DEFAULT_HI_GBP,
    tol: float = 5_000.0,
    max_iter: int = 40,
) -> SolveResult:
    """Binary-search the smallest seed NW in ``[lo, hi]`` meeting ``inp.bar``."""
    cashflows = build_cashflows(inp, inp.horizon_years)
    p_at = pension_at_unlock(inp)

    s_hi = success_at_nw(paths, hi, inp, cashflows)
    if s_hi < inp.bar:
        # Even the ceiling can't reach the bar — report it, flagged.
        return SolveResult(hi, s_hi, p_at, reached_bar=False)
    if success_at_nw(paths, lo, inp, cashflows) >= inp.bar:
        return SolveResult(lo, 1.0, p_at, reached_bar=True)

    for _ in range(max_iter):
        if hi - lo <= tol:
            break
        mid = 0.5 * (lo + hi)
        if success_at_nw(paths, mid, inp, cashflows) >= inp.bar:
            hi = mid
        else:
            lo = mid

    return SolveResult(hi, success_at_nw(paths, hi, inp, cashflows), p_at, reached_bar=True)