"""
Estimação depois do conserto: MT pelo WLS não linear do Paper 2 (cadastro da MT + SCADA + K sensores) e BT pela
propagação da tensão estimada da MT no modelo de BT do cadastro, com as cargas previstas.
"""
import os, json, tempfile
import numpy as np
import opendssdirect as dss
from senl import NLEstimator
from drpdrc import Feeder34
from rede import (Circuit, consumers, consumer_node, STEPS, MT_CUSTOMERS, mt_reading, mt_wrapper, HERE)

ORDER = json.load(open(os.path.join(HERE, "results_ieee34.json")))["modes"]["VI"]["greedy_order"]


class WarmNL(NLEstimator):
    def solve(self, z, sig, x0=None, maxit=20, tol=1e-5):
        n = self.n
        if x0 is None:
            V0 = self.fd.V0 / self.vb; x = np.concatenate([V0.real, V0.imag])
        else:
            x = x0.copy()
        w = 1 / sig
        for it in range(1, maxit + 1):
            hx, J = self.h(x, None)
            r = (z - hx) * w; Jw = J * w[:, None]
            dx, *_ = np.linalg.lstsq(Jw, r, rcond=None)
            x = x + dx
            vm = np.hypot(x[:n], x[n:])
            if not np.all(np.isfinite(x)) or vm.min() < 0.5 or vm.max() > 1.5:
                return x, it, "divergiu"
            if np.max(np.abs(dx)) < tol:
                return x, it, "ok"
        return x, maxit, "sem_convergencia"


class BTModel:
    """Blocos da matriz Y do cadastro completo (MT + BT) para propagar a tensão da MT até os consumidores."""

    def __init__(self, spec_cad, mt_nodes):
        ck = Circuit(spec_cad)
        d = ck.d
        d.Text.Command("batchedit load..* enabled=no")
        for g in d.Generators.AllNames():
            if g != "none": d.Text.Command(f"edit generator.{g} enabled=no")
        d.Solution.Solve()
        nodes = [n.upper() for n in d.Circuit.YNodeOrder()]
        Yv = np.asarray(d.Circuit.SystemY(), float); n = len(nodes)
        Y = (Yv[0::2] + 1j * Yv[1::2]).reshape(n, n)
        idx = {nd: k for k, nd in enumerate(nodes)}
        mtset = set(mt_nodes)
        self.P = [idx[nd] for nd in mt_nodes]
        self.B = [k for k, nd in enumerate(nodes) if nd not in mtset]
        self.bidx = {nodes[k]: j for j, k in enumerate(self.B)}
        self.YBB = Y[np.ix_(self.B, self.B)]; self.YBP = Y[np.ix_(self.B, self.P)]
        self.lu = np.linalg.inv(self.YBB)
        cons = consumers(spec_cad)
        self.cons_pos = np.array([self.bidx[consumer_node(c, t)] for _, t, c in cons])

    def propagate(self, Vp, S_cons, iters=4):
        """Vp: (nP, T) tensões de MT; S_cons: (nC, T) potência consumida (VA, complexa). Retorna |V| (nC, T)."""
        b = -self.YBP @ Vp
        VB = self.lu @ b
        for _ in range(iters):
            I = np.zeros_like(VB)
            Vc = VB[self.cons_pos]
            np.add.at(I, self.cons_pos, -np.conj(S_cons / Vc))
            VB = self.lu @ (b + I)
        return np.abs(VB[self.cons_pos])


def pseudo_mt(fd, spec_cad, kwc_typ, kvc_typ, kwm_typ, kvm_typ):
    """Pseudomedidas por carga original do IEEE 34 (kW, kvar): soma prevista dos consumidores de cada grupo de BT,
    repartida entre as cargas originais do grupo pela proporção do cadastro; cargas pontuais = consumidor de MT."""
    names = [L["name"] for L in fd.loads]; pos = {nm: i for i, nm in enumerate(names)}
    P = np.zeros((len(names),) + kwc_typ.shape[1:]); Q = np.zeros_like(P)
    j = 0
    for g in spec_cad["groups"]:
        n = sum(len(t["cons"]) for t in g["trafos"])
        gp = kwc_typ[j:j + n].sum(axis=0); gq = kvc_typ[j:j + n].sum(axis=0); j += n
        tot = sum(fd.loads[pos[m]]["kw"] for m in g["members"])
        for m in g["members"]:
            sh = fd.loads[pos[m]]["kw"] / tot
            P[pos[m]] += gp * sh; Q[pos[m]] += gq * sh
    for i, m in enumerate(spec_cad["mt"]):
        P[pos[m["name"]]] += kwm_typ[i]; Q[pos[m["name"]]] += kvm_typ[i]
    return P, Q


def estimate_week(spec_cad, Vtruth, truth_nodes, taps_seq, pseudo, pv_mt_typ, Ks=(0, 2, 4), rel=0.30, seed=0, workdir=None):
    """Roda o estimador da MT passo a passo e propaga para a BT.
    pseudo: dict com kwc, kvc (consumidores, STEPS x nC), kwm, kvm (MT, STEPS x nM) previstos.
    Retorna dict K -> dict(bt=|V| (STEPS x nC), mt={bus: (STEPS x 3) p.u.}, fail=n)."""
    wd = workdir or tempfile.mkdtemp()
    fd = Feeder34(mt_wrapper(spec_cad, os.path.join(wd, "mt_cad.dss")))
    vb = fd.vbase; n = len(fd.nodes)
    tmap = {nd: k for k, nd in enumerate(truth_nodes)}
    sel = np.array([tmap[nd] for nd in fd.nodes])
    ests = {K: WarmNL(fd, ORDER[:K]) for K in Ks}
    P, Q = pseudo_mt(fd, spec_cad, pseudo["kwc"].T, pseudo["kvc"].T, pseudo["kwm"].T, pseudo["kvm"].T)   # (nL, T)
    sgL = rel * np.hypot(P, Q)
    # geração de MT conhecida no cadastro (capacidade x irradiância típica), injeção trifásica equilibrada
    pv_nodes = []
    for pv in spec_cad.get("pv_mt", []):
        pv_nodes.append([fd.idx[f"{pv['bus']}.{p}".upper()] for p in (1, 2, 3)])
    out = {K: dict(Vmt=np.full((STEPS, n), np.nan + 0j), fail=0) for K in Ks}
    x_prev = {K: None for K in Ks}
    rng = np.random.default_rng(seed)
    for k in range(STEPS):
        Y = fd.Y_for(taps_seq[k]); fd.Y = Y
        Vt = Vtruth[k][sel]
        for K, est in ests.items():
            est.Yp = Y * vb[None, :]
            S = -(est.M @ (P[:, k] + 1j * Q[:, k]))
            sig = np.sqrt((np.abs(est.M) ** 2) @ sgL[:, k] ** 2) + 50.0
            for pvn, ptyp in zip(pv_nodes, pv_mt_typ):
                S[pvn] += ptyp[k] * 1e3 / 3; sig[pvn] = np.hypot(sig[pvn], 0.3 * ptyp[k] * 1e3 / 3 + 1e3)
            z, sg = est.build_z(Vt, S, sig, rng)
            x, it, st = est.solve(z, sg, x0=x_prev[K])
            if st != "ok":
                x, it, st = est.solve(z, sg, x0=None, maxit=30)
            if st == "ok":
                out[K]["Vmt"][k] = (x[:n] + 1j * x[n:]) * vb; x_prev[K] = x
            else:
                out[K]["fail"] += 1; x_prev[K] = None
    # propagação para a BT (cadastro). As cargas previstas dos consumidores são reescaladas, por barra de MT, pela
    # razão entre a carga ativa estimada na barra e a prevista (o estimador corrige a carga onde há informação).
    bt = BTModel(spec_cad, fd.nodes)
    Scons0 = (pseudo["kwc"] + 1j * pseudo["kvc"]).T * 1e3           # (nC, T) VA
    cons = consumers(spec_cad)
    buses = sorted({g["bus"] for g in spec_cad["groups"]})
    bnodes = {b: [i for i, nd in enumerate(fd.nodes) if nd.split(".")[0] == b.upper()] for b in buses}
    cbus = np.array([buses.index(g["bus"]) for g, _, _ in cons])
    Ppseudo = np.array([(-(ests[Ks[0]].M @ (P + 1j * Q)))[bnodes[b]].real.sum(axis=0) for b in buses])   # W (negativo = consumo)
    res = {}
    for K in Ks:
        Vp = out[K]["Vmt"].T                                          # (n, T)
        ok = ~np.isnan(Vp).any(axis=0)
        Vv = np.nan_to_num(Vp)
        Inj = fd.Y @ Vv                                               # corrente que sai da rede no nó
        Sinj = Vv * np.conj(Inj)                                      # potência consumida no nó (VA)
        Pest = np.array([Sinj[bnodes[b]].real.sum(axis=0) for b in buses])
        ratio = np.clip(np.where(np.abs(Ppseudo) > 1.0, Pest / np.where(np.abs(Ppseudo) > 1.0, Ppseudo, 1.0), 1.0), 0.3, 3.0)
        Scons = Scons0 * ratio[cbus]
        vbt = np.full((Scons.shape[0], STEPS), np.nan)
        vbt[:, ok] = bt.propagate(Vp[:, ok], Scons[:, ok])
        mt = {b: mt_reading(out[K]["Vmt"], fd.idx, b) for b in MT_CUSTOMERS}
        res[K] = dict(bt=vbt.T, mt=mt, fail=out[K]["fail"])
    return res


# ====================================================================== v2: modelo realista (taps desconhecidos)
REGS = ["reg1a", "reg1b", "reg1c", "reg2a", "reg2b", "reg2c"]
REG_NODES = {"reg1a": ("814.1", "814R.1"), "reg1b": ("814.2", "814R.2"), "reg1c": ("814.3", "814R.3"),
             "reg2a": ("852.1", "852R.1"), "reg2b": ("852.2", "852R.2"), "reg2c": ("852.3", "852R.3")}
TAP_PRIOR_SD = 0.02    # incerteza da posição de tap prevista pelo cadastro (~3 degraus de 0,625 %)
MODEL_ERR_INFL = 5.0   # incerteza dos sensores x5 nos pesos: absorve erro de modelo (impedâncias do cadastro)
ITS = []
RETRY = os.environ.get('G4W_RETRY', '1') == '1'   # 0: sem segunda tentativa (mais rápido)


class TapNL(WarmNL):
    """WLS não linear com a relação de transformação dos 6 reguladores como estado extra.
    Y(t) = Y_ref + ΔY(t): para o regulador com admitância série y e tap t no enrolamento 2,
    Y_ij = -y/t e Y_jj = y/t² (+ termos que não dependem de t). Prior: tap previsto pelo cadastro."""

    def __init__(self, fd, sensors, Yref, tref):
        super().__init__(fd, sensors)
        self.Yref = Yref; self.tref = np.array(tref, float)
        self.rp = []
        zpos = {k: r for r, k in enumerate(self.zero_nodes)}
        for r_, nm in enumerate(REGS):
            i, j = (fd.idx[a] for a in REG_NODES[nm])
            y = -Yref[i, j] * self.tref[r_]
            self.rp.append((i, j, y, zpos.get(i), zpos.get(j)))
        self.nz = len(self.zero_nodes)

    def inflate(self, sg, f=None):
        f = MODEL_ERR_INFL if f is None else f
        sg = sg.copy(); p = len(sg)
        for s_ in reversed(self.sensors):
            m = len(self.fd.lines[s_]["n1"])
            sg[p - 4 * m:p] *= f; p -= 4 * m
        return sg

    def Y_of(self, t):
        Y = self.Yref.copy()
        for (i, j, y, _, _), tt, t0 in zip(self.rp, t, self.tref):
            d_ij = -y / tt + y / t0; d_jj = y / tt ** 2 - y / t0 ** 2
            Y[i, j] += d_ij; Y[j, i] += d_ij; Y[j, j] += d_jj
        return Y

    def solve_tap(self, z, sig, t_prior, x0=None, maxit=25, tol=1e-4):
        n = self.n; nt = len(REGS)
        if x0 is None:
            V0 = self.fd.V0 / self.vb; x = np.concatenate([V0.real, V0.imag, t_prior])
        else:
            x = x0.copy()
        zz = np.concatenate([z, t_prior]); w = 1 / np.concatenate([sig, np.full(nt, TAP_PRIOR_SD)])
        for it in range(1, maxit + 1):
            t = np.clip(x[2 * n:], 0.88, 1.12)
            Y = self.Y_of(t); self.fd.Y = Y; self.Yp = Y * self.vb[None, :]
            hx, J = self.h(x[:2 * n], None)
            V = (x[:n] + 1j * x[n:2 * n]) * self.vb
            Jt = np.zeros((len(hx), nt))
            for p, (i, j, y, ri, rj) in enumerate(self.rp):
                dIi = (y / t[p] ** 2) * V[j]
                dIj = (y / t[p] ** 2) * V[i] - (2 * y / t[p] ** 3) * V[j]
                if ri is not None: Jt[ri, p] = dIi.real; Jt[ri + self.nz, p] = dIi.imag
                if rj is not None: Jt[rj, p] = dIj.real; Jt[rj + self.nz, p] = dIj.imag
            H = np.vstack([np.hstack([J, Jt]), np.hstack([np.zeros((nt, 2 * n)), np.eye(nt)])])
            hh = np.concatenate([hx, t])
            r = (zz - hh) * w
            dx, *_ = np.linalg.lstsq(H * w[:, None], r, rcond=None)
            x = x + dx
            vm = np.hypot(x[:n], x[n:2 * n])
            if not np.all(np.isfinite(x)) or vm.min() < 0.5 or vm.max() > 1.5:
                return x, it, "divergiu"
            if np.max(np.abs(dx)) < tol:
                return x, it, "ok"
        return x, maxit, "sem_convergencia"


def expected_taps(spec_cad, P, Q, steps):
    """Posição de tap que o cadastro prevê: fluxo de potência do modelo de MT do cadastro, com controle automático
    dos reguladores e as cargas previstas. Retorna lista de tuplas (6 reguladores)."""
    from rede import dss_commands, FEEDER
    d = dss.NewContext(); os.chdir(os.path.dirname(FEEDER))
    for c in dss_commands(spec_cad, bt=False): d.Text.Command(c)
    d.Text.Command("set controlmode=static"); d.Text.Command("set maxcontroliter=40")
    names = []
    i = d.Loads.First()
    while i: names.append(d.Loads.Name()); i = d.Loads.Next()
    out = []
    for k in steps:
        for nm, p, q in zip(names, P[:, k], Q[:, k]):
            d.Loads.Name(nm); d.Loads.kW(max(p, 0.0)); d.Loads.kvar(q)
        d.Solution.Solve()
        tp = []
        for r in REGS:
            d.Transformers.Name(r); d.Transformers.Wdg(2); tp.append(d.Transformers.Tap())
        out.append(tuple(tp))
    return out, names


def estimate_week_v2(spec_cad, Vtruth, truth_nodes, pseudo, pv_mt_typ, Ks=(0, 2, 4), rel=0.30, seed=0,
                     stride=1, workdir=None, taps_known=None):
    """Como estimate_week, mas os taps dos reguladores são desconhecidos: o estimador parte do tap que o cadastro prevê
    (fluxo do cadastro com controle automático) e o estima junto com o estado. taps_known: se dado, usa os taps reais
    (versão otimista, para comparação)."""
    wd = workdir or tempfile.mkdtemp()
    fd = Feeder34(mt_wrapper(spec_cad, os.path.join(wd, "mt_cad.dss")))
    vb = fd.vbase; n = len(fd.nodes)
    tmap = {nd: k for k, nd in enumerate(truth_nodes)}
    sel = np.array([tmap[nd] for nd in fd.nodes])
    xfm1 = spec_cad["xfm1_tap"]
    tref = (1.0,) * 6
    Yref = fd.Y_for(tref + (xfm1,))
    ests = {K: TapNL(fd, ORDER[:K], Yref, tref) for K in Ks}
    P, Q = pseudo_mt(fd, spec_cad, pseudo["kwc"].T, pseudo["kvc"].T, pseudo["kwm"].T, pseudo["kvm"].T)
    sgL = rel * np.hypot(P, Q)
    steps = list(range(0, STEPS, stride))
    texp, _ = expected_taps(spec_cad, P, Q, steps)
    pv_nodes = [[fd.idx[f"{pv['bus']}.{p}".upper()] for p in (1, 2, 3)] for pv in spec_cad.get("pv_mt", [])]
    out = {K: dict(Vmt=np.full((STEPS, n), np.nan + 0j), tap=np.full((STEPS, 6), np.nan), fail=0) for K in Ks}
    x_prev = {K: None for K in Ks}
    rng = np.random.default_rng(seed)
    for si, k in enumerate(steps):
        Vt = Vtruth[k][sel]
        tprior = np.array(texp[si]) if taps_known is None else np.array(taps_known[k][:6])
        for K, est in ests.items():
            S = -(est.M @ (P[:, k] + 1j * Q[:, k]))
            sig = np.sqrt((np.abs(est.M) ** 2) @ sgL[:, k] ** 2) + 50.0
            for pvn, ptyp in zip(pv_nodes, pv_mt_typ):
                S[pvn] += ptyp[k] * 1e3 / 3; sig[pvn] = np.hypot(sig[pvn], 0.3 * ptyp[k] * 1e3 / 3 + 1e3)
            # medições: geradas a partir da verdade (o Y verdadeiro não é necessário: h só usa Y nas injeções)
            z, sg = est.build_z(Vt, S, sig, rng)
            sg = est.inflate(sg)
            x0 = x_prev[K]
            if x0 is not None: x0 = x0.copy(); x0[2 * n:] = 0.5 * x0[2 * n:] + 0.5 * tprior
            x, it, st = est.solve_tap(z, sg, tprior, x0=x0)
            ITS.append((K, it, st))
            if st != "ok" and RETRY:
                x, it, st = est.solve_tap(z, sg, tprior, x0=None, maxit=40)
                ITS.append((K, it, "retry_" + st))
            if st == "ok":
                out[K]["Vmt"][k] = (x[:n] + 1j * x[n:2 * n]) * vb; out[K]["tap"][k] = x[2 * n:]; x_prev[K] = x
            else:
                out[K]["fail"] += 1; x_prev[K] = None
    fd.Y = Yref
    bt = BTModel(spec_cad, fd.nodes)
    Scons0 = (pseudo["kwc"] + 1j * pseudo["kvc"]).T * 1e3
    cons = consumers(spec_cad)
    buses = sorted({g["bus"] for g in spec_cad["groups"]})
    bnodes = {b: [i for i, nd in enumerate(fd.nodes) if nd.split(".")[0] == b.upper()] for b in buses}
    cbus = np.array([buses.index(g["bus"]) for g, _, _ in cons])
    Ppseudo = np.array([(-(ests[Ks[0]].M @ (P + 1j * Q)))[bnodes[b]].real.sum(axis=0) for b in buses])
    res = {}
    for K in Ks:
        Vp = out[K]["Vmt"].T
        ok = ~np.isnan(Vp).any(axis=0)
        Vv = np.nan_to_num(Vp)
        Sinj = Vv * np.conj(Yref @ Vv)            # barras de carga não dependem dos taps
        Pest = np.array([Sinj[bnodes[b]].real.sum(axis=0) for b in buses])
        ratio = np.clip(np.where(np.abs(Ppseudo) > 1.0, Pest / np.where(np.abs(Ppseudo) > 1.0, Ppseudo, 1.0), 1.0), 0.3, 3.0)
        Scons = Scons0 * ratio[cbus]
        vbt = np.full((Scons.shape[0], STEPS), np.nan)
        vbt[:, ok] = bt.propagate(Vp[:, ok], Scons[:, ok])
        mt = {b: mt_reading(out[K]["Vmt"], fd.idx, b) for b in MT_CUSTOMERS}
        res[K] = dict(bt=vbt.T, mt=mt, fail=out[K]["fail"], tap=out[K]["tap"])
    return res