"""
Estimação de estados trifásica por WLS linear (fasores) com pseudomedidas de carga e sensores sincronizados.
Verdade: fluxo de potência do OpenDSS nos alimentadores de teste do IEEE (arquivos públicos).
"""
import os
import numpy as np
import opendssdirect as dss

F_HZ = 60.0


def _c(a):
    a = np.asarray(a, float)
    return a[0::2] + 1j * a[1::2]


def _bus(name):
    return name.split(".")[0].upper()


class Feeder:
    def __init__(self, path, sub_meter_line=None, source_buses=("SOURCEBUS",), sub_v_buses=()):
        self.path = path
        dss.Text.Command("clear")
        os.chdir(os.path.dirname(path))
        dss.Text.Command(f"compile [{path}]")
        dss.Solution.Solve()                       # taps assentam com as cargas nominais
        dss.Text.Command("set controlmode=off")    # taps fixos (premissa: conhecidos)
        # cargas nominais (previsão)
        self.loads = []
        i = dss.Loads.First()
        while i:
            nm = dss.Loads.Name()
            bn = dss.CktElement.BusNames()[0]
            nph = dss.CktElement.NumPhases()
            no = dss.CktElement.NodeOrder()
            b = _bus(bn)
            nd = lambda x: f"{b}.{x}" if x != 0 else None
            if dss.Loads.IsDelta():
                br = [(nd(no[0]), nd(no[1]))] if nph == 1 else \
                     [(nd(no[k]), nd(no[(k + 1) % 3])) for k in range(3)]
            else:
                br = [(f"{b}.{no[k]}", None) for k in range(nph)]
            self.loads.append(dict(name=nm, bus=b, branches=br, kw=dss.Loads.kW(), kvar=dss.Loads.kvar()))
            i = dss.Loads.Next()
        # matriz Y sem cargas (rede: linhas, trafos, reguladores, capacitores, fonte)
        dss.Text.Command("batchedit load..* enabled=no")
        dss.Solution.Solve()
        self.nodes = [n.upper() for n in dss.Circuit.YNodeOrder()]
        self.idx = {n: k for k, n in enumerate(self.nodes)}
        n = len(self.nodes)
        self.Y = _c(dss.Circuit.SystemY()).reshape(n, n)
        self.V0 = self._read_v()                   # tensões em vazio = fasores nominais
        self.vbase = np.empty(n)
        for k, nd in enumerate(self.nodes):
            dss.Circuit.SetActiveBus(_bus(nd)); self.vbase[k] = dss.Bus.kVBase() * 1e3
        dss.Text.Command("batchedit load..* enabled=yes")
        self.source = [k for k, nd in enumerate(self.nodes) if _bus(nd) in source_buses]
        # linhas (candidatas a sensor)
        self.lines = {}
        for nm in dss.Lines.AllNames():
            dss.Circuit.SetActiveElement("Line." + nm)
            bn = [_bus(x) for x in dss.CktElement.BusNames()]
            no = dss.CktElement.NodeOrder(); nph = dss.CktElement.NumPhases()
            n1 = [f"{bn[0]}.{x}" for x in no[:nph]]; n2 = [f"{bn[1]}.{x}" for x in no[nph:2 * nph]]
            yp = _c(dss.CktElement.YPrim()).reshape(2 * nph, 2 * nph)
            cols = [self.idx[x] for x in n1 + n2]
            # corrente que entra no terminal 1, por fase: I = Yprim[:nph] · [V1; V2]
            Hi = np.zeros((nph, n), complex); Hi[:, cols] = yp[:nph]
            self.lines[nm.upper()] = dict(bus1=bn[0], bus2=bn[1], n1=[self.idx[x] for x in n1], Hi=Hi)
        self.sub_meter_line = sub_meter_line.upper() if sub_meter_line else None
        self.sub_v = [k for k, nd in enumerate(self.nodes) if _bus(nd) in [b.upper() for b in sub_v_buses]]
        self._build_pseudo()

    def _read_v(self):
        names = [n.upper() for n in dss.Circuit.YNodeOrder()]
        v = _c(dss.Circuit.YNodeVArray())
        d = dict(zip(names, v))
        return np.array([d[n] for n in self.nodes])

    def _build_pseudo(self):
        """Injeção prevista por nó (média) e matriz de mapeamento ramos de carga -> nós."""
        n = len(self.nodes)
        cols, Ibr, owner = [], [], []
        A = []
        for li, L in enumerate(self.loads):
            S = (L["kw"] + 1j * L["kvar"]) * 1e3 / len(L["branches"])
            for a, b in L["branches"]:
                ia = self.idx[a]; ib = self.idx[b] if b else None
                vn = self.V0[ia] - (self.V0[ib] if ib is not None else 0)
                col = np.zeros(n, complex); col[ia] = -1
                if ib is not None: col[ib] = +1
                A.append(col); Ibr.append(np.conj(S / vn)); owner.append(li)
        self.A = np.array(A).T                      # n x nbr: injeção = A · I_ramo
        self.Ibr_nom = np.array(Ibr)
        self.br_owner = np.array(owner)
        self.inj_nom = self.A @ self.Ibr_nom

    def truth(self, mult):
        """Resolve o fluxo com kW/kvar de cada carga multiplicados por mult[i]. Retorna V (fasores, V)."""
        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()
        assert dss.Solution.Converged()
        return self._read_v()

    def reset(self):
        self.truth(np.ones(len(self.loads)))


def scenarios(fd, N, rng, sigma=0.30, common=0.15):
    k = np.clip(1 + sigma * rng.standard_normal((N, len(fd.loads))), 0.2, 2.0)
    c = np.clip(1 + common * rng.standard_normal((N, 1)), 0.6, 1.4)
    M = k * c
    V = np.array([fd.truth(m) for m in M])
    fd.reset()
    return M, V


class Estimator:
    """WLS linear: z = H x + e, com covariância R (bloco de pseudomedidas não diagonal)."""

    def __init__(self, fd, sensors, pseudo_rel=0.5, sig_v=0.002, sig_i=0.005, sig_sub=0.001, sync_us=1.0, mode="VI", sub_current=True):
        self.fd, self.sensors = fd, list(sensors)
        n = len(fd.nodes)
        self.ang = 2 * np.pi * F_HZ * sync_us * 1e-6
        self.sig_v, self.sig_i, self.sig_sub, self.pseudo_rel = sig_v, sig_i, sig_sub, pseudo_rel
        self.sig_isub = 0.005
        rows, blocks = [], []
        E = np.eye(n)
        # 1) fonte e tensão da subestação (fasores, referência de tempo da própria subestação)
        self.r_vsub = fd.source + fd.sub_v
        rows += [E[k] for k in self.r_vsub]
        # 2) corrente na saída do alimentador
        self.r_isub = []
        if fd.sub_meter_line and sub_current:
            Hi = fd.lines[fd.sub_meter_line]["Hi"]; self.r_isub = list(range(len(rows), len(rows) + len(Hi))); rows += list(Hi)
        # 3) injeções: pseudomedidas onde há carga, injeção nula (virtual) onde não há
        self.inj_nodes = [k for k in range(n) if k not in fd.source]
        self.r_inj0 = len(rows)
        rows += [fd.Y[k] for k in self.inj_nodes]
        # 4) sensores: V nos nós da barra 1 e I no terminal 1 da linha
        self.r_sens = []
        for s in self.sensors:
            L = fd.lines[s]; r0 = len(rows)
            if mode == "VI":
                rows += [E[k] for k in L["n1"]]
            rv = list(range(r0, len(rows)))
            r1 = len(rows); rows += list(L["Hi"]); ri = list(range(r1, len(rows)))
            self.r_sens.append((s, rv, ri))
        self.H = np.array(rows)

    def _R(self, Vtrue_ref):
        """Covariância das medidas (usa magnitudes nominais para dimensionar desvios)."""
        fd = self.fd; m = len(self.H)
        R = np.zeros((m, m), complex)
        for r, k in enumerate(self.r_vsub):
            R[r, r] = (self.sig_sub * abs(fd.V0[k])) ** 2
        Isub = self.H[self.r_isub] @ Vtrue_ref if self.r_isub else []
        for r, i in zip(self.r_isub, Isub):
            R[r, r] = (self.sig_isub * abs(i)) ** 2 + 1e-6
        # pseudomedidas: cov = A Σ A^H restrito aos nós de injeção, + piso para injeção nula
        sig = self.pseudo_rel * np.abs(fd.Ibr_nom)
        Ai = fd.A[self.inj_nodes]
        C = (Ai * sig ** 2) @ Ai.conj().T
        floor = getattr(self, 'floor_scale', 1.0) * 1e-6 * max(1.0, np.max(np.abs(fd.Ibr_nom))) ** 2
        R[self.r_inj0:self.r_inj0 + len(Ai), self.r_inj0:self.r_inj0 + len(Ai)] = C + floor * np.eye(len(Ai))
        for s, rv, ri in self.r_sens:
            for r in rv:
                R[r, r] = (self.sig_v ** 2 + self.ang ** 2) * abs(fd.V0[np.argmax(np.abs(self.H[r]))]) ** 2
            Is = self.H[ri] @ Vtrue_ref
            for r, i in zip(ri, Is):
                R[r, r] = (self.sig_i ** 2 + self.ang ** 2) * abs(i) ** 2 + 1e-6
        return R

    def gain(self, Vref):
        R = self._R(Vref)
        L = np.linalg.cholesky(R)
        Hw = np.linalg.solve(L, self.H)
        self._L = L
        # QR com escala das colunas pela tensão de base (evita equações normais mal condicionadas)
        D = self.fd.vbase
        Q, Rr = np.linalg.qr(Hw * D)
        G = D[:, None] * np.linalg.solve(Rr, Q.conj().T)
        self.G = G
        return G

    def measure(self, V, rng):
        """Gera medidas para um conjunto de cenários V (S x n)."""
        fd = self.fd; S = len(V); m = len(self.H)
        Z = (self.H @ V.T)                            # verdade das medidas (m x S)
        # subestação
        for r in range(len(self.r_vsub)):
            Z[r] *= 1 + self.sig_sub * rng.standard_normal(S)
        for r in self.r_isub:
            Z[r] *= 1 + self.sig_isub * rng.standard_normal(S)
        # pseudomedidas: valor previsto (fixo), independente da verdade
        Z[self.r_inj0:self.r_inj0 + len(self.inj_nodes)] = fd.inj_nom[self.inj_nodes][:, None]
        # sensores: ruído de magnitude + erro de relógio comum ao sensor
        for s, rv, ri in self.r_sens:
            th = self.ang * rng.standard_normal(S)
            rot = np.exp(1j * th)
            for r in rv: Z[r] *= (1 + self.sig_v * rng.standard_normal(S)) * rot
            for r in ri: Z[r] *= (1 + self.sig_i * rng.standard_normal(S)) * rot
        return Z

    def estimate(self, Z):
        Zw = np.linalg.solve(self._L, Z)
        return (self.G @ Zw).T                         # S x n


def metrics(fd, Vt, Ve):
    """Erros de tensão (% da base), corrente nas linhas (A) e carga por barra (kW, kvar)."""
    out = {}
    excl = set(getattr(fd, "exclude_buses", ()))
    nodes = [k for k in range(len(fd.nodes)) if k not in fd.source and fd.nodes[k].split(".")[0] not in excl]
    dv = (np.abs(Ve) - np.abs(Vt))[:, nodes] / fd.vbase[nodes] * 100
    out["v_max"] = np.max(np.abs(dv), axis=1); out["v_node"] = np.mean(np.abs(dv), axis=0)
    out["ang_max"] = np.max(np.abs(np.angle(Ve[:, nodes] / Vt[:, nodes], deg=True)), axis=1)
    # correntes
    di, it = [], []
    for nm, L in fd.lines.items():
        It = Vt @ L["Hi"].T; Ie = Ve @ L["Hi"].T
        di.append(np.max(np.abs(np.abs(Ie) - np.abs(It)), axis=1)); it.append(np.max(np.abs(It), axis=1))
    di = np.array(di).T; it = np.array(it).T
    out["i_line_abs"] = np.mean(di, axis=0); out["i_line_true"] = np.mean(it, axis=0)
    # erro relativo agregado: soma dos erros absolutos / soma das correntes (por cenário)
    out["i_max_rel"] = np.sum(di, axis=1) / np.sum(it, axis=1) * 100
    out["i_max_abs"] = np.max(di, axis=1)
    out["line_names"] = list(fd.lines)
    # carga por barra: soma de V·conj(-I_inj) nos nós da barra
    buses = sorted({L["bus"] for L in fd.loads})
    St, Se = [], []
    for b in buses:
        ks = [fd.idx[n] for n in fd.nodes if n.split(".")[0] == b]
        St.append(np.sum(Vt[:, ks] * np.conj(-(Vt @ fd.Y[ks].T)), axis=1))
        Se.append(np.sum(Ve[:, ks] * np.conj(-(Ve @ fd.Y[ks].T)), axis=1))
    St = np.array(St).T / 1e3; Se = np.array(Se).T / 1e3
    out["load_buses"] = buses; out["S_true"] = St; out["S_est"] = Se
    out["p_err_rel"] = np.sum(np.abs(Se.real - St.real), axis=1) / np.sum(np.abs(St.real), axis=1) * 100
    return out


class EstimatorSCADA:
    """Mesma formulação do Estimator (linear em V para fonte, injeções e sensores), mas a saída do alimentador
    é medida como no SCADA: |I| por fase e P, Q trifásicos totais (não lineares). Resolve por Gauss-Newton em
    coordenadas retangulares reais, partindo da solução linear sem a medição de saída."""

    def __init__(self, fd, sensors, sig_iabs=0.005, sig_pq=0.01, **kw):
        self.fd = fd
        self.lin = Estimator(fd, sensors, sub_current=False, **kw)
        L = fd.lines[fd.sub_meter_line]
        self.Hi = L["Hi"]; self.n1 = L["n1"]
        self.sig_iabs, self.sig_pq = sig_iabs, sig_pq
        self.D = fd.vbase                      # escala das colunas (p.u.)

    def gain(self, Vref):
        lin = self.lin; R = lin._R(Vref); m, n = lin.H.shape
        # forma real de ruído complexo circular: cov([Re;Im]) = 0.5 [[Re R, -Im R], [Im R, Re R]]
        Rr = 0.5 * np.block([[R.real, -R.imag], [R.imag, R.real]])
        self.Lr = np.linalg.cholesky(Rr + 1e-12 * np.eye(2 * m) * np.mean(np.diag(Rr)))
        H = lin.H
        self.Hr = np.block([[H.real, -H.imag], [H.imag, H.real]])
        self.Hw = np.linalg.solve(self.Lr, self.Hr) * np.concatenate([self.D, self.D])[None, :]
        Iref = self.Hi @ Vref; Sref = np.sum(Vref[self.n1] * np.conj(Iref))
        self.sn = np.concatenate([self.sig_iabs * np.abs(Iref) + 1e-3, [self.sig_pq * abs(Sref)] * 2])
        Ql, Rl = np.linalg.qr(self.Hw)
        self.Glin = (np.concatenate([self.D, self.D])[:, None]) * np.linalg.solve(Rl, Ql.T)   # só parte linear

    def h_nl(self, V):
        I = self.Hi @ V; S = np.sum(V[self.n1] * np.conj(I))
        return np.concatenate([np.abs(I), [S.real, S.imag]])

    def J_nl(self, V):
        n = len(V); I = self.Hi @ V; rows = []
        for p in range(len(I)):
            a = np.conj(I[p]) * self.Hi[p] / abs(I[p])
            rows.append(np.concatenate([a.real, -a.imag]))
        dS_de = np.zeros(n, complex); dS_df = np.zeros(n, complex)
        for p, k in enumerate(self.n1):
            dS_de[k] += np.conj(I[p]); dS_df[k] += 1j * np.conj(I[p])
            dS_de += V[k] * np.conj(self.Hi[p]); dS_df += -1j * V[k] * np.conj(self.Hi[p])
        rows.append(np.concatenate([dS_de.real, dS_df.real])); rows.append(np.concatenate([dS_de.imag, dS_df.imag]))
        return np.array(rows)

    def measure(self, V, rng):
        Z = self.lin.measure(V, rng)                          # linhas lineares (sem corrente de saída)
        S = len(V); Zn = []
        for s in range(S):
            h = self.h_nl(V[s]); Zn.append(h + self.sn * rng.standard_normal(len(h)))
        return Z, np.array(Zn)

    def estimate(self, Zs, maxit=15, tol=1e-6):
        Z, Zn = Zs; n = len(self.fd.nodes); Dd = np.concatenate([self.D, self.D])
        Zr = np.vstack([Z.real, Z.imag]); Zw = np.linalg.solve(self.Lr, Zr)
        X0 = (self.Glin @ Zw).T                               # solução linear sem a saída
        out = np.empty((len(Zn), n), complex)
        for s in range(len(Zn)):
            x = X0[s].copy()
            for _ in range(maxit):
                V = x[:n] + 1j * x[n:]
                rl = Zw[:, s] - np.linalg.solve(self.Lr, self.Hr @ x)
                rn = (Zn[s] - self.h_nl(V)) / self.sn
                A = np.vstack([self.Hw, (self.J_nl(V) / self.sn[:, None]) * Dd[None, :]])
                dxs, *_ = np.linalg.lstsq(A, np.concatenate([rl, rn]), rcond=None)
                dx = dxs * Dd; x = x + dx
                if np.max(np.abs(dxs)) < tol: break
            out[s] = x[:n] + 1j * x[n:]
        return out