"""CES aggregators with analytic first derivatives.

ces(weights, x, rho) = (sum_i w_i x_i^rho)^(1/rho) for rho != 0, and the Cobb-Douglas limit
prod_i x_i^{w_i} (weights normalized to one) for rho == 0. Zero-weight inputs are ignored, which lets the
same code run the pure Restrepo benchmark (no structures, land or energy) and the full model.
"""
from __future__ import annotations

import numpy as np

_FLOOR = 1e-300


def ces(weights, x, rho: float):
    w = np.asarray(weights, dtype=float)
    x = np.asarray(x, dtype=float)
    m = w > 0
    if not m.any():
        return 0.0
    w, x = w[m], np.maximum(x[m], _FLOOR)
    if abs(rho) < 1e-12:
        ws = w / w.sum()
        return float(np.exp(np.sum(ws * np.log(x))))
    return float(np.sum(w * x ** rho) ** (1.0 / rho))


def ces_grad(weights, x, rho: float):
    """Vector of partial derivatives d ces / d x_i (zero for zero-weight inputs)."""
    w = np.asarray(weights, dtype=float)
    x = np.asarray(x, dtype=float)
    g = np.zeros_like(x)
    m = w > 0
    if not m.any():
        return g
    val = ces(w, x, rho)
    xm = np.maximum(x[m], _FLOOR)
    if abs(rho) < 1e-12:
        ws = w[m] / w[m].sum()
        g[m] = ws * val / xm
    else:
        g[m] = w[m] * xm ** (rho - 1.0) * val ** (1.0 - rho)
    return g


def rho_of_sigma(sigma: float) -> float:
    """rho = (sigma-1)/sigma; sigma = 1 -> 0 (Cobb-Douglas); sigma = inf -> 1 (linear)."""
    if np.isinf(sigma):
        return 1.0
    return (sigma - 1.0) / sigma