"""Domain objects for sprint forecasting.

The code uses arrays for performance, but keeps a clear domain-level boundary.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Sequence

import numpy as np


DeveloperId = str


@dataclass(frozen=True)
class WorkItemEstimate:
    """Three-point estimate for a remaining work item.

    Attributes
    ----------
    item_id:
        Unique identifier for the item.
    owner_id:
        Developer responsible for the item.
    optimistic:
        Optimistic estimate in hours.
    most_likely:
        Modal / most-likely estimate in hours.
    pessimistic:
        Pessimistic estimate in hours.
    ai_intensity:
        Planned AI usage intensity for the item.
        Examples: 0/1 (label) or a graded score.
    """

    item_id: str
    owner_id: DeveloperId
    optimistic: float
    most_likely: float
    pessimistic: float
    ai_intensity: float = 0.0

    def __post_init__(self) -> None:
        if self.optimistic <= 0 or self.most_likely <= 0 or self.pessimistic <= 0:
            raise ValueError("Triad values must be positive.")
        if not (self.optimistic < self.most_likely < self.pessimistic):
            raise ValueError(
                "Triad must satisfy optimistic < most_likely < pessimistic."
            )
        if self.ai_intensity < 0:
            raise ValueError("ai_intensity must be non-negative.")


@dataclass(frozen=True)
class BlockedStateSpec:
    """Blocked/disruption-state specification.

    We represent shared disruptions by a categorical state B_t in {1..K}.

    Attributes
    ----------
    multipliers:
        Array of length K with productivity multipliers g_k.
    probabilities:
        Array of shape (K, T) with per-day probabilities pi_{k,t}.
        Each column must sum to 1.
    """

    multipliers: np.ndarray
    probabilities: np.ndarray

    @property
    def n_states(self) -> int:
        return int(self.multipliers.shape[0])

    @property
    def n_days(self) -> int:
        return int(self.probabilities.shape[1])

    def __post_init__(self) -> None:
        g = np.asarray(self.multipliers, dtype=float).reshape(-1)
        p = np.asarray(self.probabilities, dtype=float)

        if g.ndim != 1:
            raise ValueError("multipliers must be a 1D array")
        if p.ndim != 2:
            raise ValueError("probabilities must be a 2D array of shape (K, T)")
        if p.shape[0] != g.shape[0]:
            raise ValueError(
                "probabilities first dimension must match multipliers length"
            )
        if np.any(g <= 0):
            raise ValueError("All multipliers g_k must be positive.")
        if np.any(p < 0):
            raise ValueError("Blocked-state probabilities must be non-negative.")
        col_sums = p.sum(axis=0)
        if not np.allclose(col_sums, 1.0, atol=1e-9):
            raise ValueError("Each day t must satisfy sum_k pi_{k,t} = 1.")

        object.__setattr__(self, "multipliers", g)
        object.__setattr__(self, "probabilities", p)


@dataclass(frozen=True)
class SprintPlan:
    """All inputs needed to simulate remaining capacity.

    Arrays are stored in developer-major form with shape (n_devs, n_days).

    Attributes
    ----------
    developer_ids:
        Sequence of developer identifiers.
    planned_hours:
        h_{i,t} of shape (n, T).
    availability_prob:
        q_{i,t} of shape (n, T).
    dev_day_ai_intensity:
        u_{i,t} of shape (n, T).
    focus_log_mu:
        Base mu^{(F)}_{i,t} (on log scale) of shape (n, T).
    focus_log_sigma:
        sigma^{(F)}_{i,t} (on log scale) of shape (n, T).
    blocked:
        Blocked state model parameters.

    Notes
    -----
    Day indices in the model are 1..T (paper notation), but arrays are 0-indexed.
    """

    developer_ids: Sequence[DeveloperId]
    planned_hours: np.ndarray
    availability_prob: np.ndarray
    dev_day_ai_intensity: np.ndarray
    focus_log_mu: np.ndarray
    focus_log_sigma: np.ndarray
    blocked: BlockedStateSpec

    def __post_init__(self) -> None:
        dev_ids = list(self.developer_ids)
        if len(dev_ids) == 0:
            raise ValueError("developer_ids must be non-empty")

        h = np.asarray(self.planned_hours, dtype=float)
        q = np.asarray(self.availability_prob, dtype=float)
        u = np.asarray(self.dev_day_ai_intensity, dtype=float)
        mu = np.asarray(self.focus_log_mu, dtype=float)
        sig = np.asarray(self.focus_log_sigma, dtype=float)

        if h.ndim != 2:
            raise ValueError("planned_hours must have shape (n, T)")
        n, T = h.shape

        arrays_to_check = [
            (q, "availability_prob"),
            (u, "dev_day_ai_intensity"),
            (mu, "focus_log_mu"),
            (sig, "focus_log_sigma"),
        ]

        for arr, name in arrays_to_check:
            if arr.shape != (n, T):
                raise ValueError(f"{name} must have shape {(n, T)}")

        if len(dev_ids) != n:
            raise ValueError("developer_ids length must match planned_hours rows")

        if np.any(h < 0):
            raise ValueError("planned_hours must be non-negative")
        if np.any((q < 0) | (q > 1)):
            raise ValueError("availability_prob must be in [0,1]")
        if np.any(u < 0):
            raise ValueError("dev_day_ai_intensity must be non-negative")
        if np.any(sig < 0):
            raise ValueError("focus_log_sigma must be non-negative")

        if self.blocked.n_days != T:
            raise ValueError("blocked.probabilities must have T columns")

        object.__setattr__(self, "developer_ids", tuple(dev_ids))
        object.__setattr__(self, "planned_hours", h)
        object.__setattr__(self, "availability_prob", q)
        object.__setattr__(self, "dev_day_ai_intensity", u)
        object.__setattr__(self, "focus_log_mu", mu)
        object.__setattr__(self, "focus_log_sigma", sig)

    @property
    def n_devs(self) -> int:
        return int(self.planned_hours.shape[0])

    @property
    def n_days(self) -> int:
        return int(self.planned_hours.shape[1])

    def developer_index(self, developer_id: DeveloperId) -> int:
        try:
            return self.developer_ids.index(developer_id)
        except ValueError as exc:
            raise KeyError(f"Unknown developer_id: {developer_id}") from exc

    def iter_days(self) -> Iterable[int]:
        return range(self.n_days)
