"""Workload (demand) modelling and simulation."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Sequence

import numpy as np

from .distributions import WeightedLogNormalQuantileFitter
from .domain import WorkItemEstimate
from .models.bias import BiasModelParameters


@dataclass(frozen=True)
class WorkloadSampler:
    """Precomputed sampler for total workload W = sum_j X_j.

    The sampler uses the identity:
        X_j = exp(mu_j + sigma_j * Z_j)
    where Z_j ~ Normal(0,1).

    The same standard normal draws Z_j can be reused across counterfactual
    scenarios to reduce Monte Carlo variance (common random numbers).
    """

    mu: np.ndarray
    sigma: np.ndarray

    def __post_init__(self) -> None:
        mu = np.asarray(self.mu, dtype=float).reshape(-1)
        sigma = np.asarray(self.sigma, dtype=float).reshape(-1)
        if mu.shape != sigma.shape:
            raise ValueError("mu and sigma must have the same shape")
        if np.any(sigma <= 0):
            raise ValueError("All sigma values must be positive")
        object.__setattr__(self, "mu", mu)
        object.__setattr__(self, "sigma", sigma)

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

    def sample_total(self, rng: np.random.Generator, size: int) -> np.ndarray:
        """Sample total workload for each Monte Carlo draw."""
        if size <= 0:
            raise ValueError("size must be positive")
        if self.n_items == 0:
            return np.zeros(size, dtype=float)

        z = rng.normal(loc=0.0, scale=1.0, size=(size, self.n_items))
        return np.sum(np.exp(self.mu + self.sigma * z), axis=1)

    def sample_total_with_normals(self, z: np.ndarray) -> np.ndarray:
        """Sample total workload using pre-generated standard normals.

        Parameters
        ----------
        z:
            Array of shape (size, n_items) of standard normal variates.
        """
        z = np.asarray(z, dtype=float)
        if z.ndim != 2 or z.shape[1] != self.n_items:
            raise ValueError(
                f"z must have shape (size, {self.n_items}), got {z.shape}"
            )
        return np.sum(np.exp(self.mu + self.sigma * z), axis=1)


class WorkloadModelBuilder:
    """Build workload samplers from (bias-corrected) three-point estimates."""

    def __init__(
        self,
        fitter: WeightedLogNormalQuantileFitter | None = None,
    ):
        self.fitter = fitter or WeightedLogNormalQuantileFitter()

    def build_sampler(
        self,
        items: Sequence[WorkItemEstimate],
        bias: BiasModelParameters,
        *,
        override_item_ai_intensity: float | None = None,
    ) -> WorkloadSampler:
        """Create a sampler for W from remaining items.

        Parameters
        ----------
        override_item_ai_intensity:
            If provided, this value replaces each item's ai_intensity.
            Use 0.0 to construct a "noAI" counterfactual on the demand side.
        """
        if len(items) == 0:
            return WorkloadSampler(
                mu=np.array([], dtype=float),
                sigma=np.array([], dtype=float),
            )

        mu_list: list[float] = []
        sigma_list: list[float] = []

        for item in items:
            ai_u = (
                item.ai_intensity
                if override_item_ai_intensity is None
                else float(override_item_ai_intensity)
            )
            corrected = WorkItemEstimate(
                item_id=item.item_id,
                owner_id=item.owner_id,
                optimistic=item.optimistic,
                most_likely=item.most_likely,
                pessimistic=item.pessimistic,
                ai_intensity=ai_u,
            )
            a_p, m_p, b_p = bias.correct_triad(corrected)
            params = self.fitter.fit(a=a_p, m=m_p, b=b_p)
            mu_list.append(params.mu)
            sigma_list.append(params.sigma)

        return WorkloadSampler(
            mu=np.array(mu_list, dtype=float),
            sigma=np.array(sigma_list, dtype=float),
        )
