"""Full static block for one period: tasks + structures + land + energy + compute utilization + housing.

Given predetermined stocks (K_N, K_Q, K_E, K_H), the period's exogenous state (alpha_t, Z_t, eps_t) and
parameters, `static_period` returns all quantities and prices of the period except those that depend on
household consumption (housing rent), which `housing_rent` adds once total consumption is known.
`income_identity` checks the national-accounting identity GDP = factor payments.
"""
from __future__ import annotations

import numpy as np
from scipy.optimize import brentq

from .tasks import Technology, technology_from_params


def energy_supply(p, KE, gf=1.0):
    """Energy output; the resource endowment is measured in efficiency units R^E * gf."""
    return p.ZE * KE ** (1.0 - p.eta) * (p.RE * gf) ** p.eta


def static_period(p, tech: Technology, KN, KQ, KE, KH, alpha_t, Z_t, eps_t, x0=None, gf=1.0):
    """gf = (1+g_Z)^t: Harrod-neutral efficiency factor applied to labor, land and the energy resource."""
    E = energy_supply(p, KE, gf) if p.nu_E > 0 else 0.0
    H_eff = np.asarray(p.H, dtype=float) * gf
    T_eff = p.T_land * gf
    TH_eff = p.TH_land * gf

    def run(u):
        Q = u * KQ
        EY = E - eps_t * Q if p.nu_E > 0 else 0.0
        sol = tech.solve(Q, H_eff, KN, T_eff, EY, Z_t, alpha_t, x0=x0)
        return sol, EY

    u = 1.0
    if p.nu_E > 0 and KQ > 0:
        # feasibility: energy left for final production must be positive
        u_max = min(1.0, (1 - 1e-10) * E / (eps_t * KQ)) if eps_t * KQ > 0 else 1.0
        sol, EY = run(u_max)
        net = sol["A"] - eps_t * sol["PE"]
        if net < 0.0:
            # capacity is idle until the marginal unit just covers its energy cost: bisection in log u
            lo = np.log(min(1e-12, 0.5 * u_max))
            hi = np.log(u_max)
            f_lo = (lambda s: s[0]["A"] - eps_t * s[0]["PE"])(run(np.exp(lo)))
            if f_lo <= 0:
                u = np.exp(lo)
            else:
                for _ in range(60):
                    mid = 0.5 * (lo + hi)
                    fm = (lambda s: s[0]["A"] - eps_t * s[0]["PE"])(run(np.exp(mid)))
                    if fm > 0:
                        lo = mid
                    else:
                        hi = mid
                    if hi - lo < 1e-11:
                        break
                u = np.exp(lo)
            sol, EY = run(u)
        else:
            u = u_max
    else:
        sol, EY = run(1.0)

    Q = u * KQ
    PE = sol["PE"] if p.nu_E > 0 else 0.0
    RQ = u * (sol["A"] - eps_t * PE)            # net rental per unit of installed capacity
    RE = (1.0 - p.eta) * PE * E / KE if (p.nu_E > 0 and KE > 0) else 0.0
    rentR = p.eta * PE * E if p.nu_E > 0 else 0.0
    HS = KH ** (1.0 - p.nu_H) * TH_eff ** p.nu_H
    out = dict(Y=sol["Y"], w=sol["w"], A=sol["A"], RN=sol["RN"], RT=sol["RT"], PE=PE, RQ=RQ, RE=RE, rentR=rentR,
               util=u, Qtot=Q, E=E, EY=EY, EQ=eps_t * Q, HS=HS, L=sol["L"], Qi=sol["Q"], X=sol["X"], F=sol["F"],
               XB=sol["XB"], x=sol["x"], success=sol["success"], KN=KN, KQ=KQ, KE=KE, KH=KH, gf=gf,
               T_eff=T_eff, TH_eff=TH_eff, H_eff=H_eff,
               wage_bill=float(np.sum(sol["w"] * H_eff)))
    return out


def housing_rent(p, per, C_total):
    """Rent per unit of housing services when total goods consumption is C_total (Cobb-Douglas share)."""
    RH = p.theta_H / (1.0 - p.theta_H) * C_total / per["HS"]
    per["RH"] = RH
    per["RHK"] = (1.0 - p.nu_H) * RH * per["HS"] / per["KH"]      # rental of housing structures
    per["RTH"] = p.nu_H * RH * per["HS"] / per["TH_eff"]           # residential land rent per efficiency unit
    per["GDP"] = per["Y"] + RH * per["HS"]
    return per


def income_identity(p, per, markup_profit=0.0):
    """Residual of GDP - (wages + rentals + rents + profits). Zero in the model; must fail if accounting is broken."""
    inc = (per["wage_bill"] + per["RN"] * per["KN"] + per["RT"] * per["T_eff"] + per["RQ"] * per["KQ"]
           + per["RE"] * per["KE"] + per["rentR"] + per["RHK"] * per["KH"] + per["RTH"] * per["TH_eff"])
    # RQ already includes the markup profit when markup_Q > 1 (owners receive A - eps PE per unit)
    return per["GDP"] - inc


def value_added_by_sector(p, per):
    """Sectoral value added (final goods, compute, energy, housing); sums to GDP."""
    final = per["Y"] - per["A"] * per["Qtot"] - per["PE"] * per["EY"]
    compute = per["A"] * per["Qtot"] - per["PE"] * per["EQ"]
    energy = per["PE"] * per["E"]
    housing = per["RH"] * per["HS"]
    return dict(final=final, compute=compute, energy=energy, housing=housing, total=final + compute + energy + housing)