"""Static task block: output maximization over labor and compute allocations (Restrepo's work equation).

The competitive equilibrium of the static block is the solution of the concave program

    max_{L_i, Q_i >= 0}  Y = Z * ( G(V(X_B({X_i}), K_N), T, E_Y) + sum_a gamma_a X_a )
    s.t.  X_i = L_i + Q_i / alpha_i,  sum_{i in Omega(s)} L_i = H_s,  sum_i Q_i = Q,

and factor prices are the multipliers (envelope theorem): w_s = max_{i in Omega(s)} F_i,
A = max_i F_i / alpha_i, where F_i = dY/dX_i. Tasks with alpha_i = inf cannot use compute.

Nests (zero weights switch a nest off, so the pure Restrepo benchmark is the case nu_N = nu_T = nu_E = 0):
    X_B = (sum_i nu_i X_i^rho_B)^(1/rho_B)
    V   = (nu_N K_N^rho_N + (1-nu_N) X_B^rho_N)^(1/rho_N)
    G   = (nu_T T^rho_T + nu_E E_Y^rho_T + nu_V V^rho_T)^(1/rho_T),   nu_V = 1 - nu_T - nu_E
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np
from scipy.optimize import minimize, fsolve

from .ces import ces, ces_grad

_XFLOOR = 1e-10


@dataclass
class Technology:
    nu: np.ndarray            # task weights (bottleneck), zeros for accessory
    gamma: np.ndarray         # linear accessory coefficients, zeros for bottleneck
    skill: np.ndarray         # skill index of each task
    rho_B: float
    nu_N: float
    rho_N: float
    nu_T: float
    nu_E: float
    rho_T: float

    @property
    def n(self):
        return len(self.nu)

    @property
    def nskill(self):
        return int(self.skill.max()) + 1

    # ----- nested aggregator and its marginal products -----
    def aggregate(self, X, KN, T, EY, Z):
        """Return dict with Y, XB, V, G and marginal values F_i = dY/dX_i, R_N, R_T, P_E."""
        X = np.maximum(np.asarray(X, dtype=float), _XFLOOR)
        XB = ces(self.nu, X, self.rho_B)
        if self.nu_N > 0:
            V = ces([self.nu_N, 1 - self.nu_N], [KN, XB], self.rho_N)
            gV = ces_grad([self.nu_N, 1 - self.nu_N], [KN, XB], self.rho_N)
            dV_dKN, dV_dXB = gV[0], gV[1]
        else:
            V, dV_dKN, dV_dXB = XB, 0.0, 1.0
        nu_V = 1.0 - self.nu_T - self.nu_E
        if self.nu_T > 0 or self.nu_E > 0:
            wts = [self.nu_T, self.nu_E, nu_V]
            G = ces(wts, [T, EY, V], self.rho_T)
            gG = ces_grad(wts, [T, EY, V], self.rho_T)
            dG_dT, dG_dE, dG_dV = gG[0], gG[1], gG[2]
        else:
            G, dG_dT, dG_dE, dG_dV = V, 0.0, 0.0, 1.0
        acc = float(np.sum(self.gamma * X))
        Y = Z * (G + acc)
        dXB = ces_grad(self.nu, X, self.rho_B)
        F = Z * (dG_dV * dV_dXB * dXB + self.gamma)
        return dict(Y=Y, XB=XB, V=V, G=G, F=F,
                    RN=Z * dG_dV * dV_dKN, RT=Z * dG_dT, PE=Z * dG_dE)

    # ----- the static program -----
    def solve(self, Q, H, KN, T, EY, Z, alpha, x0: Optional[np.ndarray] = None, tol=1e-14):
        """Solve the static allocation. Variables are scaled: l_i = L_i/H_{s(i)}, q_i = Q_i/Q in [0,1]."""
        n = self.n
        alpha = np.asarray(alpha, dtype=float)
        fin = np.isfinite(alpha)
        inv_alpha = np.where(fin, 1.0 / np.where(fin, alpha, 1.0), 0.0)
        H = np.asarray(H, dtype=float)
        Hs = H[self.skill]                      # supply of the skill owning each task
        nS = self.nskill
        use_Q = bool(fin.any()) and Q > 0
        Qs = float(Q) if use_Q else 0.0

        def X_of(x):
            L = x[:n] * Hs
            Qa = x[n:] * Qs if use_Q else np.zeros(n)
            return L, Qa, L + Qa * inv_alpha

        # objective scale: output at the interior starting point
        def _x0():
            l0 = np.array([1.0 / np.sum(self.skill == s) for s in self.skill])
            if use_Q:
                q0 = np.where(fin, 1.0 / fin.sum(), 0.0)
                return np.concatenate([l0, q0])
            return l0

        if x0 is None or len(x0) != (2 * n if use_Q else n):
            x0 = _x0()
        scale = max(self.aggregate(X_of(_x0())[2], KN, T, EY, Z)["Y"], 1e-12)

        def obj(x):
            return -self.aggregate(X_of(x)[2], KN, T, EY, Z)["Y"] / scale

        def jac(x):
            F = self.aggregate(X_of(x)[2], KN, T, EY, Z)["F"]
            g = np.concatenate([F * Hs, F * inv_alpha * Qs]) if use_Q else F * Hs
            return -g / scale

        cons = []
        for s in range(nS):
            m = (self.skill == s).astype(float)
            mm = np.concatenate([m, np.zeros(n)]) if use_Q else m
            cons.append({"type": "eq", "fun": (lambda x, mm=mm: 1.0 - mm @ x), "jac": (lambda x, mm=mm: -mm)})
        if use_Q:
            mq = np.concatenate([np.zeros(n), fin.astype(float)])
            cons.append({"type": "eq", "fun": (lambda x, mq=mq: 1.0 - mq @ x), "jac": (lambda x, mq=mq: -mq)})

        bounds = [(0.0, 1.0)] * n
        if use_Q:
            bounds += [((0.0, 1.0) if fin[i] else (0.0, 0.0)) for i in range(n)]
        x0 = np.clip(x0, 0.0, 1.0)
        res = minimize(obj, x0, jac=jac, bounds=bounds, constraints=cons, method="SLSQP",
                       options={"ftol": tol, "maxiter": 1000})
        x, nit = res.x, res.nit
        res2 = minimize(obj, x, jac=jac, bounds=bounds, constraints=cons, method="SLSQP",
                        options={"ftol": tol, "maxiter": 1000})
        if res2.fun <= res.fun:
            x, res = res2.x, res2
        x = self._polish(x, n, use_Q, fin, inv_alpha, Hs, Qs, X_of, KN, T, EY, Z, nS)
        L, Qa, X = X_of(x)
        agg = self.aggregate(X, KN, T, EY, Z)
        F = agg["F"]
        w = np.array([F[self.skill == s].max() for s in range(nS)])
        A = float(np.max((F * inv_alpha)[fin])) if use_Q else 0.0
        out = dict(L=L, Q=Qa, X=X, w=w, A=A, Y=agg["Y"], XB=agg["XB"], V=agg["V"], G=agg["G"], F=F,
                   RN=agg["RN"], RT=agg["RT"], PE=agg["PE"], x=x, success=bool(res.success), nit=nit)
        return out

    def _polish(self, x, n, use_Q, fin, inv_alpha, Hs, Qs, X_of, KN, T, EY, Z, nS, tol_active=1e-7):
        """Newton polish of the KKT equalities on the active set identified by SLSQP.

        Free variables: l_i with l_i > tol, q_i with q_i > tol; unknown prices w_s, A. Equations: F_i = w_s for
        free labor, F_i/alpha_i = A for free compute, adding-up constraints. The result does not depend on the
        SLSQP warm start (to machine precision), which keeps finite-difference Jacobians of the dynamic model clean.
        """
        freeL = np.where(x[:n] > tol_active)[0]
        freeQ = np.where(x[n:] > tol_active)[0] if use_Q else np.array([], dtype=int)
        nL, nQ = len(freeL), len(freeQ)
        if nL == 0:
            return x

        def eqs(z):
            xx = x.copy()
            xx[freeL] = z[:nL]
            if use_Q:
                xx[n + freeQ] = z[nL:nL + nQ]
            w = z[nL + nQ:nL + nQ + nS]
            A = z[-1] if use_Q else 0.0
            F = self.aggregate(X_of(xx)[2], KN, T, EY, Z)["F"]
            r = [F[i] - w[self.skill[i]] for i in freeL]
            r += [F[i] * inv_alpha[i] - A for i in freeQ]
            r += [np.sum(xx[:n][self.skill == s]) - 1.0 for s in range(nS)]
            if use_Q:
                r.append(np.sum(xx[n:][fin]) - 1.0)
            return np.array(r)

        F0 = self.aggregate(X_of(x)[2], KN, T, EY, Z)["F"]
        w0 = np.array([F0[self.skill == s].max() for s in range(nS)])
        z0 = np.concatenate([x[freeL], x[n + freeQ] if use_Q else [], w0, [np.max((F0 * inv_alpha)[fin])] if use_Q else []])
        # scale: prices can be tiny; solve in logs of prices
        try:
            z, info, ier, msg = fsolve(eqs, z0, full_output=True, xtol=1e-14)
        except Exception:
            return x
        # accept on the residual, not on fsolve's status flag (it complains about xtol at machine precision)
        if np.max(np.abs(eqs(z))) > 1e-11 * max(1.0, np.max(np.abs(z))) or np.any(z[:nL + nQ] < -1e-12) or np.any(z[:nL + nQ] > 1 + 1e-12):
            return x
        xx = x.copy()
        xx[freeL] = np.clip(z[:nL], 0.0, 1.0)
        if use_Q:
            xx[n + freeQ] = np.clip(z[nL:nL + nQ], 0.0, 1.0)
        # accept only if the polished point is at least as good and KKT-consistent (no inactive variable wants in)
        F = self.aggregate(X_of(xx)[2], KN, T, EY, Z)["F"]
        w = z[nL + nQ:nL + nQ + nS]
        A = z[-1] if use_Q else 0.0
        if np.any(F - w[self.skill] > 1e-9 * np.abs(w[self.skill]) + 1e-14):
            return x
        if use_Q and np.any((F * inv_alpha)[fin] - A > 1e-9 * abs(A) + 1e-14):
            return x
        return xx

    def check_kkt(self, sol, H, Q, alpha, tol=1e-6):
        """Complementary slackness of the task allocation; returns max violation."""
        F, L, Qa, w, A = sol["F"], sol["L"], sol["Q"], sol["w"], sol["A"]
        alpha = np.asarray(alpha, dtype=float)
        viol = 0.0
        for i in range(self.n):
            s = self.skill[i]
            # labor used only where F_i = w_s
            if L[i] > tol * max(H[s], 1.0):
                viol = max(viol, abs(F[i] - w[s]) / max(w[s], 1e-12))
            else:
                viol = max(viol, max(F[i] - w[s], 0.0) / max(w[s], 1e-12))
            if np.isfinite(alpha[i]) and Q > 0:
                if Qa[i] > tol * max(Q, 1.0):
                    viol = max(viol, abs(F[i] / alpha[i] - A) / max(A, 1e-12))
                else:
                    viol = max(viol, max(F[i] / alpha[i] - A, 0.0) / max(A, 1e-12))
        return viol


def technology_from_params(p) -> Technology:
    ts = p.tasks
    return Technology(nu=np.array(ts.nu, dtype=float), gamma=np.array(ts.gamma, dtype=float),
                      skill=np.array(ts.skill, dtype=int), rho_B=p.rho_B, nu_N=p.nu_N, rho_N=p.rho_N,
                      nu_T=p.nu_T, nu_E=p.nu_E, rho_T=p.rho_T)


def restrepo_benchmark(nu, gamma, skill, rho_B, alpha, H, Q_grid, Z=1.0):
    """Pure endowment economy of Restrepo (2025): Y = Z (X_B + sum gamma_a X_a), compute an endowment Q.

    Returns a list of static solutions along Q_grid (used for the replication of Propositions 1-5)."""
    tech = Technology(nu=np.asarray(nu, float), gamma=np.asarray(gamma, float), skill=np.asarray(skill, int),
                      rho_B=rho_B, nu_N=0.0, rho_N=0.0, nu_T=0.0, nu_E=0.0, rho_T=0.0)
    out, x0 = [], None
    for Q in Q_grid:
        sol = tech.solve(Q, H, KN=0.0, T=0.0, EY=0.0, Z=Z, alpha=alpha, x0=x0)
        x0 = sol["x"]
        H_arr = np.asarray(H, float)
        wages_bill = float(np.sum(sol["w"] * H_arr))
        sol["labor_share"] = wages_bill / sol["Y"]
        sol["compute_share"] = sol["A"] * Q / sol["Y"]
        sol["Q_total"] = Q
        out.append(sol)
    return out