"""
Estimação de estados trifásica NÃO linear (Gauss-Newton, coordenadas retangulares) com pseudomedidas de
potência por nó (como em sistemas ADMS), medições SCADA na subestação e sensores sincronizados.
Estuda convergência e qualidade sob dados de carga degradados (falta, granularidade, prossumidores sem sentido
de fluxo, erros grosseiros) e o efeito de tratamento de dados e de poucos sensores.
"""
import sys, os
import numpy as np
import opendssdirect as dss
sys.path.insert(0, os.environ.get("G4W_P1", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "p1v2")))
from se3f import Feeder, _bus  # noqa: E402

F_HZ = 60.0


class FeederPV(Feeder):
    """Alimentador com geração fotovoltaica em parte das cargas (gerador PQ com fp = 1)."""

    def __init__(self, *a, pv_share=0.3, pv_ratio=(0.6, 1.2), seed=11, **k):
        super().__init__(*a, **k)
        rng = np.random.default_rng(seed)
        n = len(self.loads)
        self.pv_kw = np.zeros(n)
        pick = rng.choice(n, max(1, int(round(pv_share * n))), replace=False)
        for i in pick:
            L = self.loads[i]
            dss.Loads.Name(L["name"])
            bus1 = dss.CktElement.BusNames()[0]; nph = dss.CktElement.NumPhases(); kv = dss.Loads.kV()
            conn = "delta" if dss.Loads.IsDelta() else "wye"
            kw = L["kw"] * rng.uniform(*pv_ratio)
            self.pv_kw[i] = kw
            dss.Text.Command(f"New Generator.pv{i} bus1={bus1} phases={nph} kV={kv} kW=0 pf=1 model=1 conn={conn}")
        self.pv_idx = pick
        dss.Text.Command("set maxiterations=200")
        dss.Solution.Solve()

    def truth(self, mult, irr=0.0):
        for i in self.pv_idx:
            dss.Generators.Name(f"pv{i}"); dss.Generators.kW(self.pv_kw[i] * irr)
        for L, m in zip(self.loads, mult):
            dss.Loads.Name(L["name"]); dss.Loads.kW(L["kw"] * m); dss.Loads.kvar(L["kvar"] * m)
        dss.Solution.Solve()
        if not dss.Solution.Converged():               # nova tentativa com Newton; se falhar, o cenário é descartado
            dss.Text.Command("set algorithm=newton"); dss.Solution.Solve(); dss.Text.Command("set algorithm=normal")
            if not dss.Solution.Converged():
                return None
        return self._read_v()


class NLEstimator:
    """WLS não linear. Estado: tensões nodais (retangulares, p.u. da base de cada nó).
    Medidas: injeção nula (corrente, linear), pseudomedida de P e Q por nó com carga, tensão da fonte (fasor),
    |V| e P/Q por fase na saída do alimentador (SCADA), sensores (fasores de V e I)."""

    def __init__(self, fd, sensors=(), sig_v=0.002, sig_i=0.005, sync_us=1.0):
        self.fd = fd; n = len(fd.nodes); self.n = n
        self.vb = fd.vbase
        self.Yp = fd.Y * self.vb[None, :]           # corrente (A) por p.u. de tensão
        self.load_nodes = sorted({k for k in range(n) if np.any(fd.A[k] != 0)} - set(fd.source))
        self.zero_nodes = [k for k in range(n) if k not in fd.source and k not in self.load_nodes]
        self.sensors = list(sensors); self.sig_v, self.sig_i = sig_v, sig_i
        self.ang = 2 * np.pi * F_HZ * sync_us * 1e-6
        self.sub = fd.lines[fd.sub_meter_line]
        # matriz que leva a potência de cada carga (kVA) à injeção de potência por nó, com tensões nominais
        # (carga em delta entre a e b: S_a = V0_a conj(I_ab), S_b = -V0_b conj(I_ab), I_ab = conj(S/(V0_a-V0_b)))
        M = np.zeros((n, len(fd.loads)), complex)
        for li, L in enumerate(fd.loads):
            nb = len(L["branches"])
            for a, b in L["branches"]:
                ia = fd.idx[a]; ib = fd.idx[b] if b else None
                vab = fd.V0[ia] - (fd.V0[ib] if ib is not None else 0)
                # S do ramo = S/nb; potência "vista" no nó a: V0_a * conj(I) = V0_a/vab * S_ramo
                M[ia, li] += fd.V0[ia] / vab / nb
                if ib is not None: M[ib, li] += -fd.V0[ib] / vab / nb
        self.M = M * 1e3                            # kVA -> VA

    # ---------- modelo de medição
    def h(self, x, Zspec):
        """Zspec: dict com sigmas e valores; retorna h(x) e Jacobiano real."""
        n = self.n
        V = x[:n] + 1j * x[n:]                       # p.u.
        Vv = V * self.vb                              # V
        I = self.fd.Y @ Vv                            # A (corrente que sai da rede no nó)
        rows_h, rows_J = [], []

        def cplx_rows(val, dval_dV, dval_dVc):
            # dval = A dV + B dV*,  dV = de + j df  ->  d/de = A + B ; d/df = j(A - B)
            De = dval_dV + dval_dVc; Df = 1j * (dval_dV - dval_dVc)
            return val, De, Df

        out_h, out_J = [], []
        # 1) injeção nula (linear em V): I_k = 0  -> parte real e imaginária
        Z0 = self.Yp[self.zero_nodes]
        out_h += [I[self.zero_nodes].real, I[self.zero_nodes].imag]
        out_J += [np.hstack([Z0.real, -Z0.imag]), np.hstack([Z0.imag, Z0.real])]
        # 2) pseudomedidas de potência por nó: S_k = Vv_k conj(I_k)  (VA)
        L = self.load_nodes
        S = Vv[L] * np.conj(I[L])
        A = np.zeros((len(L), n), complex); B = np.zeros((len(L), n), complex)
        A[np.arange(len(L)), L] = np.conj(I[L]) * self.vb[L]
        B[:] = Vv[L][:, None] * np.conj(self.Yp[L])
        _, De, Df = cplx_rows(S, A, B)
        out_h += [S.real, S.imag]; out_J += [np.hstack([De.real, Df.real]), np.hstack([De.imag, Df.imag])]
        # 3) tensão da fonte (fasor, referência)
        src = self.fd.source
        out_h += [V[src].real, V[src].imag]
        E = np.eye(n)[src]; Zr = np.zeros_like(E)
        out_J += [np.hstack([E, Zr]), np.hstack([Zr, E])]
        # 4) SCADA na saída: |V| no nó de envio e P, Q por fase no trecho
        n1 = self.sub["n1"]; Hi = self.sub["Hi"] * self.vb[None, :]
        Vm = np.abs(V[n1]); En = np.eye(n)[n1]
        out_h += [Vm]; out_J += [np.hstack([En * (V[n1].real / Vm)[:, None], En * (V[n1].imag / Vm)[:, None]])]
        Il = Hi @ (V); Sl = Vv[n1] * np.conj(Il)
        A = np.zeros((len(n1), n), complex); A[np.arange(len(n1)), n1] = np.conj(Il) * self.vb[n1]
        B = Vv[n1][:, None] * np.conj(Hi)
        _, De, Df = cplx_rows(Sl, A, B)
        out_h += [Sl.real, Sl.imag]; out_J += [np.hstack([De.real, Df.real]), np.hstack([De.imag, Df.imag])]
        # 5) sensores: fasores (lineares)
        for s in self.sensors:
            Ls = self.fd.lines[s]; Es = np.eye(n)[Ls["n1"]]; Zs = np.zeros_like(Es)
            out_h += [V[Ls["n1"]].real, V[Ls["n1"]].imag]; out_J += [np.hstack([Es, Zs]), np.hstack([Zs, Es])]
            Hs = Ls["Hi"] * self.vb[None, :]; Is = Hs @ V
            out_h += [Is.real, Is.imag]; out_J += [np.hstack([Hs.real, -Hs.imag]), np.hstack([Hs.imag, Hs.real])]
        return np.concatenate(out_h), np.vstack(out_J)

    def build_z(self, Vt, pseudo_S, pseudo_sig, rng):
        """Monta z e sigmas a partir da verdade (Vt em V) e das pseudomedidas por nó (VA)."""
        fd = self.fd; n = self.n
        x_true = np.concatenate([(Vt / self.vb).real, (Vt / self.vb).imag])
        htrue, _ = self.h(x_true, None)
        z, sig = [], []
        nz = len(self.zero_nodes)
        z += [np.zeros(2 * nz)]; sig += [np.full(2 * nz, 1e-3)]           # A (injeção nula, quase exata)
        L = self.load_nodes
        z += [pseudo_S[L].real, pseudo_S[L].imag]; sig += [pseudo_sig[L], pseudo_sig[L]]
        ns = len(fd.source); p = 2 * nz + 2 * len(L)
        vs = Vt[fd.source] / self.vb[fd.source]
        z += [vs.real, vs.imag]; sig += [np.full(2 * ns, 1e-4)]; p += 2 * ns
        n1 = self.sub["n1"]; k = len(n1)
        vm = htrue[p:p + k] * (1 + 0.002 * rng.standard_normal(k)); z += [vm]; sig += [np.full(k, 0.002)]; p += k
        P = htrue[p:p + k]; Q = htrue[p + k:p + 2 * k]; Sm = np.hypot(P, Q)
        z += [P + 0.005 * Sm * rng.standard_normal(k), Q + 0.005 * Sm * rng.standard_normal(k)]
        sig += [0.005 * Sm + 1.0, 0.005 * Sm + 1.0]; p += 2 * k
        for s in self.sensors:
            Ls = fd.lines[s]; m = len(Ls["n1"]); th = self.ang * rng.standard_normal(); rot = np.exp(1j * th)
            Vs = (htrue[p:p + m] + 1j * htrue[p + m:p + 2 * m]) * (1 + self.sig_v * rng.standard_normal(m)) * rot
            z += [Vs.real, Vs.imag]; sv = np.sqrt(self.sig_v ** 2 + self.ang ** 2) * np.abs(Vs) / np.sqrt(2) + 1e-6
            sig += [sv, sv]; p += 2 * m
            Is = (htrue[p:p + m] + 1j * htrue[p + m:p + 2 * m]) * (1 + self.sig_i * rng.standard_normal(m)) * rot
            z += [Is.real, Is.imag]; si = np.sqrt(self.sig_i ** 2 + self.ang ** 2) * np.abs(Is) / np.sqrt(2) + 0.05
            sig += [si, si]; p += 2 * m
        return np.concatenate(z), np.concatenate(sig)

    def norm_residuals(self, x, z, sig):
        hx, J = self.h(x, None)
        w = 1 / sig; rw = (z - hx) * w; Jw = J * w[:, None]
        Q, _ = np.linalg.qr(Jw)
        om = np.clip(1 - np.sum(Q * Q, axis=1), 1e-12, None)
        return rw / np.sqrt(om)

    def solve_clean(self, z, sig, rows, thr=3.0, max_rounds=25):
        """Limpeza automática: maior resíduo normalizado entre as linhas 'rows' (pseudomedidas) é desacreditado
        (sigma x 100) e o estado é reestimado, até não haver resíduo acima do limiar."""
        sig = sig.copy(); removed = 0
        x, it, st, J = self.solve(z, sig)
        for _ in range(max_rounds):
            if st != "ok": break
            rn = self.norm_residuals(x, z, sig)
            cand = [r for r in rows if sig[r] < 1e9]
            k = max(cand, key=lambda r: abs(rn[r]))
            if abs(rn[k]) < thr: break
            # desacredita P e Q do mesmo nó
            pair = k + len(rows) // 2 if k < rows[0] + len(rows) // 2 else k - len(rows) // 2
            sig[k] *= 100; sig[pair] *= 100; removed += 1
            x, it, st, J = self.solve(z, sig)
        return x, it, st, J, removed, sig

    def solve(self, z, sig, maxit=30, tol=1e-5):
        n = self.n
        V0 = self.fd.V0 / self.vb
        x = np.concatenate([V0.real, V0.imag])
        w = 1 / sig
        for it in range(1, maxit + 1):
            hx, J = self.h(x, None)
            r = (z - hx) * w; Jw = J * w[:, None]
            try:
                dx, *_ = np.linalg.lstsq(Jw, r, rcond=None)
            except np.linalg.LinAlgError:
                return x, it, "singular", np.nan
            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", np.nan
            if np.max(np.abs(dx)) < tol:
                hx, _ = self.h(x, None)
                Jc = float(np.sum(((z - hx) * w) ** 2))
                return x, it, "ok", Jc
        return x, maxit, "sem_convergencia", np.nan