"""Aggregate income streams into per-year engine arrays.

Two outputs feed the simulator:

- ``inflows``: per-year real-GBP cash that lands in the portfolio at
  start-of-year (compounds at the year's return). This includes every
  enabled stream regardless of tax treatment.
- ``taxable``: per-year real-GBP that the jurisdiction tax engine should
  charge income tax on (i.e. ``tax_treatment='income'``). Dividend / CGT
  buckets are accounted for at withdrawal time, not income time;
  ``tax_free`` is excluded entirely.

``growth_pct`` is real (after inflation), applied geometrically each
year from ``start_year``. The simulator already runs in real-GBP, so we
do not double-deflate.
"""
from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt


@dataclass(frozen=True)
class IncomeStreamInput:
    """Engine-level income stream — decoupled from ORM and Pydantic so
    callers can construct them however."""
    kind: str = "salary"
    start_year: int = 0
    end_year: int | None = None
    amount_gbp_per_year: float = 0.0
    growth_pct: float = 0.0
    tax_treatment: str = "income"
    enabled: bool = True


def streams_to_arrays(
    streams: Iterable[IncomeStreamInput],
    horizon_years: int,
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Returns ``(inflows, taxable)`` — both ``(horizon_years,)``."""
    inflows = np.zeros(horizon_years, dtype=np.float64)
    taxable = np.zeros(horizon_years, dtype=np.float64)
    for stream in streams:
        if not stream.enabled or stream.amount_gbp_per_year <= 0:
            continue
        start = max(0, int(stream.start_year))
        if start >= horizon_years:
            continue
        end_inclusive = (stream.end_year if stream.end_year is not None else horizon_years - 1)
        end_inclusive = min(int(end_inclusive), horizon_years - 1)
        if end_inclusive < start:
            continue
        for y in range(start, end_inclusive + 1):
            growth_factor = (1.0 + stream.growth_pct)**(y - start)
            value = stream.amount_gbp_per_year * growth_factor
            inflows[y] += value
            if stream.tax_treatment == "income":
                taxable[y] += value
    return inflows, taxable