"""
Pension valuation (annuity evaluator) service.

Computes present value for:
- Single-life level annuity with optional COLA and discounting
- Joint-survivor annuity with survivor continuation percentage

Survival probabilities are sourced from `number_tables` if available
for the requested month range, using the ratio NA_t / NA_0 for the
specified sex and race. If monthly entries are missing and life table
values are available, a simple exponential survival curve is derived
from life expectancy (LE) to approximate monthly survival.

Rates are provided as percentages (e.g., 3.0 = 3%).
"""

from __future__ import annotations

from dataclasses import dataclass
import math
from typing import Dict, List, Optional, Tuple

from sqlalchemy.orm import Session

from app.models.pensions import LifeTable, NumberTable


class InvalidCodeError(ValueError):
    pass


_RACE_MAP: Dict[str, str] = {
    "W": "w",  # White
    "B": "b",  # Black
    "H": "h",  # Hispanic
    "A": "a",  # All races
}

_SEX_MAP: Dict[str, str] = {
    "M": "m",
    "F": "f",
    "A": "a",  # All sexes
}


def _normalize_codes(sex: str, race: str) -> Tuple[str, str, str]:
    sex_u = (sex or "A").strip().upper()
    race_u = (race or "A").strip().upper()
    if sex_u not in _SEX_MAP:
        raise InvalidCodeError("Invalid sex code; expected one of M, F, A")
    if race_u not in _RACE_MAP:
        raise InvalidCodeError("Invalid race code; expected one of W, B, H, A")
    return _RACE_MAP[race_u] + _SEX_MAP[sex_u], sex_u, race_u


def _to_monthly_rate(annual_percent: float) -> float:
    """Convert an annual percentage (e.g. 6.0) to monthly effective rate."""
    annual_rate = float(annual_percent or 0.0) / 100.0
    if annual_rate <= -1.0:
        # Avoid invalid negative base
        raise ValueError("Annual rate too negative")
    return (1.0 + annual_rate) ** (1.0 / 12.0) - 1.0


def _load_monthly_na_series(
    db: Session,
    *,
    sex: str,
    race: str,
    start_month: int,
    months: int,
    interpolate_missing: bool = False,
    interpolation_method: str = "linear",  # "linear" or "step"
) -> Optional[List[float]]:
    """Return NA series for months [start_month, start_month + months - 1].

    Values are floats for the column `na_{suffix}`. If any month in the
    requested range is missing, returns None to indicate fallback.
    """
    if months <= 0:
        return []

    suffix, _, _ = _normalize_codes(sex, race)
    na_col = f"na_{suffix}"

    month_values: Dict[int, float] = {}
    rows: List[NumberTable] = (
        db.query(NumberTable)
        .filter(NumberTable.month >= start_month, NumberTable.month < start_month + months)
        .all()
    )
    for row in rows:
        value = getattr(row, na_col, None)
        if value is not None:
            month_values[int(row.month)] = float(value)

    # Build initial series with possible gaps
    series_vals: List[Optional[float]] = []
    for m in range(start_month, start_month + months):
        series_vals.append(month_values.get(m))

    if any(v is None for v in series_vals) and interpolate_missing:
        # Linear interpolation for internal gaps
        if (interpolation_method or "linear").lower() == "step":
            # Step-wise: carry forward previous known; if leading gaps, use next known
            # Fill leading gaps
            first_known = None
            for idx, val in enumerate(series_vals):
                if val is not None:
                    first_known = float(val)
                    break
            if first_known is None:
                return None
            for i in range(len(series_vals)):
                if series_vals[i] is None:
                    # find prev known
                    prev_val = None
                    for k in range(i - 1, -1, -1):
                        if series_vals[k] is not None:
                            prev_val = float(series_vals[k])
                            break
                    if prev_val is not None:
                        series_vals[i] = prev_val
                    else:
                        # Use first known for leading gap
                        series_vals[i] = first_known
        else:
            for i in range(len(series_vals)):
                if series_vals[i] is None:
                    # find prev
                    prev_idx = None
                    for k in range(i - 1, -1, -1):
                        if series_vals[k] is not None:
                            prev_idx = k
                            break
                    # find next
                    next_idx = None
                    for k in range(i + 1, len(series_vals)):
                        if series_vals[k] is not None:
                            next_idx = k
                            break
                    if prev_idx is None or next_idx is None:
                        return None
                    v0 = float(series_vals[prev_idx])
                    v1 = float(series_vals[next_idx])
                    frac = (i - prev_idx) / (next_idx - prev_idx)
                    series_vals[i] = v0 + (v1 - v0) * frac

    if any(v is None for v in series_vals):
        return None

    return [float(v) for v in series_vals]  # type: ignore


def _approximate_survival_from_le(le_years: float, months: int) -> List[float]:
    """Approximate monthly survival probabilities using an exponential model.

    Given life expectancy in years (LE), approximate a constant hazard rate
    such that expected remaining life equals LE. For a memoryless exponential
    distribution, E[T] = 1/lambda. We discretize monthly: p_survive(t) = exp(-lambda * t_years).
    """
    if le_years is None or le_years <= 0:
        # No survival; return zero beyond t=0
        return [1.0] + [0.0] * (max(0, months - 1))

    lam = 1.0 / float(le_years)
    series: List[float] = []
    for idx in range(months):
        t_years = idx / 12.0
        series.append(float(pow(2.718281828459045, -lam * t_years)))
    return series


def _load_life_expectancy(db: Session, *, age: int, sex: str, race: str) -> Optional[float]:
    suffix, _, _ = _normalize_codes(sex, race)
    le_col = f"le_{suffix}"
    row: Optional[LifeTable] = db.query(LifeTable).filter(LifeTable.age == age).first()
    if not row:
        return None
    val = getattr(row, le_col, None)
    return float(val) if val is not None else None


def _to_survival_probabilities(
    db: Session,
    *,
    start_age: Optional[int],
    sex: str,
    race: str,
    term_months: int,
    interpolation_method: str = "linear",
) -> List[float]:
    """Build per-month survival probabilities p(t) for t in [0, term_months-1].

    Prefer monthly NumberTable NA series if contiguous; otherwise approximate
    from LifeTable life expectancy at `start_age`.
    """
    if term_months <= 0:
        return []

    # Try exact monthly NA series first
    na_series = _load_monthly_na_series(
        db,
        sex=sex,
        race=race,
        start_month=0,
        months=term_months,
        interpolate_missing=True,
        interpolation_method=interpolation_method,
    )
    if na_series is not None and len(na_series) > 0:
        base = na_series[0]
        if base is None or base <= 0:
            # Degenerate base; fall back
            na_series = None
        else:
            probs = [float(v) / float(base) for v in na_series]
            # Clamp to [0,1]
            return [0.0 if p < 0.0 else (1.0 if p > 1.0 else p) for p in probs]

    # Fallback to LE approximation
    le_years = _load_life_expectancy(db, age=int(start_age or 0), sex=sex, race=race)
    return _approximate_survival_from_le(le_years if le_years is not None else 0.0, term_months)


def _present_value_from_stream(
    payments: List[float],
    *,
    discount_monthly: float,
    cola_monthly: float,
) -> float:
    """PV of a cash-flow stream with monthly discount and monthly COLA growth applied."""
    pv = 0.0
    growth_factor = 1.0
    discount_factor = 1.0
    for idx, base_payment in enumerate(payments):
        if idx == 0:
            growth_factor = 1.0
            discount_factor = 1.0
        else:
            growth_factor *= (1.0 + cola_monthly)
            discount_factor *= (1.0 + discount_monthly)
        pv += (base_payment * growth_factor) / discount_factor
    return float(pv)


def _compute_growth_factor_at_month(
    month_index: int,
    *,
    cola_annual_percent: float,
    cola_mode: str,
    cola_cap_percent: Optional[float] = None,
) -> float:
    """Compute nominal COLA growth factor at month t relative to t=0.

    cola_mode:
      - "monthly": compound monthly using effective monthly rate derived from annual percent
      - "annual_prorated": step annually, prorate linearly within the year
    """
    annual_pct = float(cola_annual_percent or 0.0)
    if cola_cap_percent is not None:
        try:
            annual_pct = min(annual_pct, float(cola_cap_percent))
        except Exception:
            pass

    if month_index <= 0 or annual_pct == 0.0:
        return 1.0

    if (cola_mode or "monthly").lower() == "annual_prorated":
        years_completed = month_index // 12
        remainder_months = month_index % 12
        a = annual_pct / 100.0
        step = (1.0 + a) ** years_completed
        prorata = 1.0 + a * (remainder_months / 12.0)
        return float(step * prorata)
    else:
        # monthly compounding from annual percent
        m = _to_monthly_rate(annual_pct)
        return float((1.0 + m) ** month_index)


@dataclass
class SingleLifeInputs:
    monthly_benefit: float
    term_months: int
    start_age: Optional[int]
    sex: str
    race: str
    discount_rate: float = 0.0  # annual percent
    cola_rate: float = 0.0  # annual percent
    defer_months: float = 0.0  # months to delay first payment (supports fractional)
    payment_period_months: int = 1  # months per payment (1=monthly, 3=quarterly, etc.)
    certain_months: int = 0  # months guaranteed from commencement regardless of mortality
    cola_mode: str = "monthly"  # "monthly" or "annual_prorated"
    cola_cap_percent: Optional[float] = None
    interpolation_method: str = "linear"
    max_age: Optional[int] = None


def present_value_single_life(db: Session, inputs: SingleLifeInputs) -> float:
    """Compute PV of a single-life level annuity under mortality and economic assumptions."""
    if inputs.monthly_benefit < 0:
        raise ValueError("monthly_benefit must be non-negative")
    if inputs.term_months < 0:
        raise ValueError("term_months must be non-negative")

    if inputs.payment_period_months <= 0:
        raise ValueError("payment_period_months must be >= 1")
    if inputs.defer_months < 0:
        raise ValueError("defer_months must be >= 0")
    if inputs.certain_months < 0:
        raise ValueError("certain_months must be >= 0")

    # Survival probabilities for participant
    # Adjust term if max_age is provided and start_age known
    term_months = inputs.term_months
    if inputs.max_age is not None and inputs.start_age is not None:
        max_months = max(0, (int(inputs.max_age) - int(inputs.start_age)) * 12)
        term_months = min(term_months, max_months)

    p_survive = _to_survival_probabilities(
        db,
        start_age=inputs.start_age,
        sex=inputs.sex,
        race=inputs.race,
        term_months=term_months,
        interpolation_method=inputs.interpolation_method,
    )

    i_m = _to_monthly_rate(inputs.discount_rate)
    period = int(inputs.payment_period_months)
    t0 = int(math.ceil(inputs.defer_months))
    t = t0
    guarantee_end = float(inputs.defer_months) + float(inputs.certain_months)

    pv = 0.0
    first = True
    while t < term_months:
        p_t = p_survive[t] if t < len(p_survive) else 0.0
        base_amt = inputs.monthly_benefit * float(period)
        # Pro-rata first payment if deferral is fractional
        if first:
            frac_defer = float(inputs.defer_months) - math.floor(float(inputs.defer_months))
            pro_rata = 1.0 - (frac_defer / float(period)) if frac_defer > 0 else 1.0
        else:
            pro_rata = 1.0
        eff_base = base_amt * pro_rata
        amount = eff_base if t < guarantee_end else eff_base * p_t
        growth = _compute_growth_factor_at_month(
            t,
            cola_annual_percent=inputs.cola_rate,
            cola_mode=inputs.cola_mode,
            cola_cap_percent=inputs.cola_cap_percent,
        )
        discount = (1.0 + i_m) ** t
        pv += (amount * growth) / discount
        t += period
        first = False
    return float(pv)


@dataclass
class JointSurvivorInputs:
    monthly_benefit: float
    term_months: int
    participant_age: Optional[int]
    participant_sex: str
    participant_race: str
    spouse_age: Optional[int]
    spouse_sex: str
    spouse_race: str
    survivor_percent: float  # as percent (0-100)
    discount_rate: float = 0.0  # annual percent
    cola_rate: float = 0.0  # annual percent
    defer_months: float = 0.0
    payment_period_months: int = 1
    certain_months: int = 0
    cola_mode: str = "monthly"
    cola_cap_percent: Optional[float] = None
    survivor_basis: str = "contingent"  # "contingent" or "last_survivor"
    survivor_commence_participant_only: bool = False
    interpolation_method: str = "linear"
    max_age: Optional[int] = None


def present_value_joint_survivor(db: Session, inputs: JointSurvivorInputs) -> Dict[str, float]:
    """Compute PV for a joint-survivor annuity.

    Expected monthly payment at time t:
      E[Payment_t] = B * P(both alive at t) + B * s * P(spouse alive only at t)
                   = B * [ (1 - s) * P(both alive) + s * P(spouse alive) ]
    where s = survivor_percent (0..1)
    """
    if inputs.monthly_benefit < 0:
        raise ValueError("monthly_benefit must be non-negative")
    if inputs.term_months < 0:
        raise ValueError("term_months must be non-negative")
    if inputs.survivor_percent < 0 or inputs.survivor_percent > 100:
        raise ValueError("survivor_percent must be between 0 and 100")

    if inputs.payment_period_months <= 0:
        raise ValueError("payment_period_months must be >= 1")
    if inputs.defer_months < 0:
        raise ValueError("defer_months must be >= 0")
    if inputs.certain_months < 0:
        raise ValueError("certain_months must be >= 0")

    # Adjust term if max_age is provided and participant_age known
    term_months = inputs.term_months
    if inputs.max_age is not None and inputs.participant_age is not None:
        max_months = max(0, (int(inputs.max_age) - int(inputs.participant_age)) * 12)
        term_months = min(term_months, max_months)

    p_part = _to_survival_probabilities(
        db,
        start_age=inputs.participant_age,
        sex=inputs.participant_sex,
        race=inputs.participant_race,
        term_months=term_months,
        interpolation_method=inputs.interpolation_method,
    )
    p_sp = _to_survival_probabilities(
        db,
        start_age=inputs.spouse_age,
        sex=inputs.spouse_sex,
        race=inputs.spouse_race,
        term_months=term_months,
        interpolation_method=inputs.interpolation_method,
    )

    s_frac = float(inputs.survivor_percent) / 100.0

    i_m = _to_monthly_rate(inputs.discount_rate)
    period = int(inputs.payment_period_months)
    t0 = int(math.ceil(inputs.defer_months))
    t = t0
    guarantee_end = float(inputs.defer_months) + float(inputs.certain_months)

    pv_total = 0.0
    pv_both = 0.0
    pv_surv = 0.0
    first = True
    while t < term_months:
        p_part_t = p_part[t] if t < len(p_part) else 0.0
        p_sp_t = p_sp[t] if t < len(p_sp) else 0.0
        p_both = p_part_t * p_sp_t
        p_sp_only = p_sp_t - p_both
        base_amt = inputs.monthly_benefit * float(period)
        # Pro-rata first payment if deferral is fractional
        if first:
            frac_defer = float(inputs.defer_months) - math.floor(float(inputs.defer_months))
            pro_rata = 1.0 - (frac_defer / float(period)) if frac_defer > 0 else 1.0
        else:
            pro_rata = 1.0
        both_amt = base_amt * pro_rata * p_both
        if inputs.survivor_commence_participant_only:
            surv_basis_prob = p_part_t
        else:
            surv_basis_prob = p_sp_only
        surv_amt = base_amt * pro_rata * s_frac * surv_basis_prob
        if (inputs.survivor_basis or "contingent").lower() == "last_survivor":
            # Last-survivor: pay full while either is alive, then 0
            # E[Payment_t] = base_amt * P(participant alive OR spouse alive)
            p_either = p_part_t + p_sp_t - p_both
            total_amt = base_amt * pro_rata * p_either
            # Components are less meaningful; keep mortality-only decomposition
        else:
            # Contingent: full while both alive, survivor_percent to spouse when only spouse alive
            total_amt = base_amt * pro_rata if t < guarantee_end else (both_amt + surv_amt)

        growth = _compute_growth_factor_at_month(
            t,
            cola_annual_percent=inputs.cola_rate,
            cola_mode=inputs.cola_mode,
            cola_cap_percent=inputs.cola_cap_percent,
        )
        discount = (1.0 + i_m) ** t

        pv_total += (total_amt * growth) / discount
        # Components exclude guarantee to reflect mortality-only decomposition
        pv_both += (both_amt * growth) / discount
        pv_surv += (surv_amt * growth) / discount
        t += period
        first = False

    return {
        "pv_total": float(pv_total),
        "pv_participant_component": float(pv_both),
        "pv_survivor_component": float(pv_surv),
    }


__all__ = [
    "SingleLifeInputs",
    "JointSurvivorInputs",
    "present_value_single_life",
    "present_value_joint_survivor",
    "InvalidCodeError",
]