"""Task-based production with compute (Model A, §1.3 of docs/model_specification.md).

Y = e^z K^alpha T^(1-alpha),   ln T = ∫ ln y(x) dx  (Cobb–Douglas over tasks x in [0,1]).
Tasks x <= iota are produced by compute with productivity psi(x) = psi0 * exp(-psi1 x); the rest by labour.
With unit elasticity across tasks the bundle is

    T = Psi(iota) * (Q_P/iota)^iota * (L/(1-iota))^(1-iota),   Psi(iota) = psi0^iota * exp(-psi1 iota^2 / 2).

The automated set is [0, iota] with iota = min(iota_bar, iota*) and  psi(iota*) = p_Q / w  (threshold condition),
which combined with the factor-price conditions gives  psi(iota) = iota/(1-iota) * L/Q_P   (A.3).

All functions are pure numpy; no state.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.optimize import brentq


@dataclass(frozen=True)
class TaskTech:
    alpha: float = 0.33
    iota_bar: float = 0.5
    psi0: float = 1.0
    psi1: float = 3.0

    def psi(self, x):
        return self.psi0 * np.exp(-self.psi1 * np.asarray(x, dtype=float))

    def Psi(self, iota):
        """exp(∫_0^iota ln psi(x) dx) = psi0^iota * exp(-psi1 iota^2/2)."""
        iota = np.asarray(iota, dtype=float)
        return self.psi0 ** iota * np.exp(-self.psi1 * iota ** 2 / 2.0)

    # ------------------------------------------------------------------ automation threshold
    def iota_of(self, QP: float, L: float) -> float:
        """Automated task share from (A.3): psi(iota) = iota/(1-iota) * L/QP, capped at iota_bar.

        LHS is decreasing in iota, RHS increasing from 0 to +inf, hence a unique interior root for QP > 0.
        """
        if QP <= 1e-9 * L:
            return 0.0
        f = lambda i: np.log(self.psi(i)) - np.log(i / (1.0 - i)) - np.log(L / QP)
        # f(0+) = +inf, f(1-) = -inf  -> unique root
        lo, hi = 1e-9, 1.0 - 1e-9
        if f(lo) <= 0:
            return 0.0
        root = brentq(f, lo, hi, xtol=1e-14, rtol=1e-12, maxiter=200)
        return float(min(root, self.iota_bar))

    # ------------------------------------------------------------------ bundle and output
    def bundle(self, QP: float, L: float, iota: float | None = None) -> float:
        if iota is None:
            iota = self.iota_of(QP, L)
        if iota <= 0.0:
            return float(L)
        return float(self.Psi(iota) * (QP / iota) ** iota * (L / (1.0 - iota)) ** (1.0 - iota))

    def output(self, z: float, K: float, QP: float, L: float, iota: float | None = None) -> float:
        T = self.bundle(QP, L, iota)
        return float(np.exp(z) * K ** self.alpha * T ** (1.0 - self.alpha))

    # ------------------------------------------------------------------ prices (competitive, CRS)
    def prices(self, z: float, K: float, QP: float, L: float) -> dict:
        """Rental rate of capital, wage, compute rental, output, automated share and factor shares.

        When iota is at the frontier iota_bar (corner), the compute price is still the marginal product of compute
        in the automated tasks, (1-alpha) iota Y / QP, and the threshold condition holds with inequality.
        """
        iota = self.iota_of(QP, L)
        Y = self.output(z, K, QP, L, iota)
        RK = self.alpha * Y / K
        w = (1.0 - self.alpha) * (1.0 - iota) * Y / L
        pQ = (1.0 - self.alpha) * iota * Y / QP if QP > 0 else 0.0
        return dict(Y=Y, iota=iota, RK=RK, w=w, pQ=pQ,
                    share_K=self.alpha, share_Q=(1.0 - self.alpha) * iota, share_L=(1.0 - self.alpha) * (1.0 - iota))

    def production_compute_demand(self, pQ: float, w: float, L: float, z: float, K: float) -> float:
        """Compute demanded by the firm at rental pQ given the wage (used for compute-market clearing).

        From the FOCs: pQ = (1-alpha) iota Y/QP and w = (1-alpha)(1-iota) Y/L  ->  QP = iota/(1-iota) * L * w/pQ,
        with iota from the threshold psi(iota) = pQ/w (capped at iota_bar). If pQ/w > psi(0), no task is automated.
        """
        ratio = pQ / w
        if ratio >= self.psi0:
            return 0.0
        iota = min(self.iota_bar, float(-np.log(ratio / self.psi0) / self.psi1))
        return float(iota / (1.0 - iota) * L / ratio)