"""Balanced-growth path of the economy with frozen AI (alpha fixed): the reference state.

With Harrod-neutral efficiency growth g on labor, land and the energy resource, all stocks, output and
consumption grow at g, prices per efficiency unit are constant, and the interest rate is
r* = (1+g)^gamma / beta - 1 (Euler with a constant price index). Each capital stock satisfies the
steady-state version of its Euler equation (investment rate x = delta + g):
    R_k = q_k (r* + delta_k) - p_k phi_k g (delta_k + g/2),  q_k = p_k (1 + phi_k g).
"""
from __future__ import annotations

import numpy as np
from scipy.optimize import fsolve

from .production import housing_rent, static_period
from .tasks import technology_from_params
from .transition import KINDS


def bgp_interest_rate(p):
    return (1.0 + p.g_Z) ** p.gamma_u / p.beta - 1.0


def target_rentals(p, r, pQ=None):
    g = p.g_Z
    price = {"N": 1.0, "Q": p.pQ0 if pQ is None else pQ, "E": 1.0, "H": 1.0}
    R = {}
    for k in KINDS:
        d, ph, pk = p.delta[k], p.phi[k], price[k]
        q = pk * (1 + ph * g)
        R[k] = q * (r + d) - pk * ph * g * (d + g / 2)
    return R


def solve_bgp(p, alpha=None, eps=None, K0=None, KQ=None, verbose=False):
    """Stationary (detrended) stocks K_N, K_E, K_H such that every capital Euler holds at r*, holding the
    compute stock at KQ (default p.KQ0) and solving the compute Euler for the price p^Q_0 in closed form:
        p^Q_0 = R_Q / [ (1+phi g)(r*+delta_Q) - phi g (delta_Q + g/2) ].
    Returns the stocks, r*, the implied pQ0 and the period dictionary (per efficiency unit, gf = 1).
    """
    tech = technology_from_params(p)
    alpha = p.alpha_path()[0] if alpha is None else alpha
    eps = p.eps0 if eps is None else eps
    KQ = p.KQ0 if KQ is None else KQ
    r = bgp_interest_rate(p)
    Rstar = target_rentals(p, r, 1.0)
    g = p.g_Z
    x = {k: p.delta[k] + g for k in KINDS}
    kk = ("N", "E", "H")

    def period(K, pQ):
        s = static_period(p, tech, K["N"], K["Q"], K["E"], K["H"], alpha, p.Z0, eps, gf=1.0)
        inv = sum((pQ if k == "Q" else 1.0) * (x[k] * K[k] + 0.5 * p.phi[k] * g ** 2 * K[k]) for k in KINDS)
        C = s["Y"] - inv
        housing_rent(p, s, C)
        return s, C

    def pQ_of(s):
        d, ph = p.delta["Q"], p.phi["Q"]
        return s["RQ"] / ((1 + ph * g) * (r + d) - ph * g * (d + g / 2))

    def resid(logK):
        K = dict(zip(kk, np.exp(logK)))
        K["Q"] = KQ
        s, C = period(K, 1.0)
        pQ = max(pQ_of(s), 1e-9)
        s, C = period(K, pQ)
        Rk = {"N": s["RN"], "E": s["RE"], "H": s["RHK"]}
        return np.array([np.log(max(Rk[k], 1e-300) / Rstar[k]) for k in kk])

    K0 = {"N": p.KN0, "E": p.KE0, "H": p.KH0} if K0 is None else K0
    lk0 = np.log([K0[k] for k in kk])
    sol, info, ier, msg = fsolve(resid, lk0, full_output=True, xtol=1e-13)
    K = dict(zip(kk, np.exp(sol)))
    K["Q"] = KQ
    s, C = period(K, 1.0)
    pQ = pQ_of(s)
    s, C = period(K, pQ)
    res = resid(sol)
    if verbose:
        print("bgp fsolve:", msg.strip(), "max|res|", np.abs(res).max(), "pQ0", pQ)
    return dict(K=K, r=r, pQ0=pQ, period=s, C=C, resid=res, ok=(ier == 1 and np.abs(res).max() < 1e-8))


def apply_bgp(p, b):
    """Return a copy of p whose initial stocks and compute price are the BGP values."""
    return p.replace(KN0=b["K"]["N"], KQ0=b["K"]["Q"], KE0=b["K"]["E"], KH0=b["K"]["H"], pQ0=b["pQ0"])