"""
Paper 3 (G4W) — Observabilidade e compensações por DRP/DRC.
Caso exemplo: semana regulatória (1008 leituras de 10 min) antes e depois de um conserto na MT, no IEEE 34.
Verdade: OpenDSS (reguladores em controle automático passo a passo). Estimação: WLS não linear do Paper 2
(pseudomedidas de P/Q por nó a partir de curvas típicas e do cadastro, SCADA na saída) com 0, 2 ou 4 sensores.
"""
import os, sys
import numpy as np
import opendssdirect as dss
from senl import NLEstimator  # noqa: E402  (usa se3f.Feeder)
from se3f import Feeder, _c, _bus  # noqa: E402

STEPS = 1008            # 7 dias x 144 leituras de 10 min
DT_H = 1 / 6

# Faixas para 2,3 kV <= Vn < 69 kV (PRODIST Módulo 8, Anexo 8.A — a conferir no texto oficial)
MT_ADQ = (0.93, 1.05)
MT_PREC_LOW = 0.90      # precária: 0,90 <= V < 0,93 ; crítica: V < 0,90 ou V > 1,05


def classify_mt(v):
    """v: tensões em p.u. da referência. Retorna (precária, crítica) booleanos."""
    crit = (v < MT_PREC_LOW) | (v > MT_ADQ[1])
    prec = (v >= MT_PREC_LOW) & (v < MT_ADQ[0])
    return prec, crit


def drp_drc(vpu):
    """vpu: (STEPS, nfases) tensões de leitura em p.u. DRP/DRC (%) pela maior contagem entre as fases."""
    prec, crit = classify_mt(vpu)
    nlp = prec.sum(axis=0).max(); nlc = crit.sum(axis=0).max()
    return 100 * nlp / len(vpu), 100 * nlc / len(vpu)


# ---------------------------------------------------------------- curvas de carga (sintéticas, documentadas)
def typical_shapes():
    """Curvas típicas normalizadas (pico = 1) por classe, resolução de 10 min, dia útil e fim de semana."""
    h = np.arange(144) / 6
    res = 0.35 + 0.15 * np.exp(-((h - 7.5) / 1.5) ** 2) + 0.65 * np.exp(-((h - 19.5) / 2.0) ** 2) + 0.1 * np.exp(-((h - 13) / 2) ** 2)
    com = 0.30 + 0.70 * (1 / (1 + np.exp(-(h - 8.0) * 2))) * (1 / (1 + np.exp((h - 18.5) * 2)))
    ind = 0.55 + 0.45 * (1 / (1 + np.exp(-(h - 6.5) * 3))) * (1 / (1 + np.exp((h - 22.0) * 3)))
    S = {"res": res / res.max(), "com": com / com.max(), "ind": ind / ind.max()}
    weekend = {"res": 1.05, "com": 0.55, "ind": 0.75}
    return S, weekend


def week_curves(classes, rng, start_dow=0, noise=0.12, rho=0.9):
    """Curvas verdadeiras por carga (STEPS x nL): típica x fator diário x ruído AR(1)."""
    S, wk = typical_shapes()
    nL = len(classes); out = np.empty((STEPS, nL)); typ = np.empty((STEPS, nL))
    e = np.zeros(nL)
    for d in range(7):
        dow = (start_dow + d) % 7; we = dow >= 5
        day = np.clip(1 + 0.06 * rng.standard_normal(), 0.85, 1.15)       # clima do dia (comum)
        for t in range(144):
            k = d * 144 + t
            base = np.array([S[c][t] * (wk[c] if we else 1.0) for c in classes])
            e = rho * e + np.sqrt(1 - rho ** 2) * noise * rng.standard_normal(nL)
            out[k] = base * day * (1 + e)
            typ[k] = base
    return np.clip(out, 0.05, None), typ


# ---------------------------------------------------------------- rede com estado de reguladores
class Feeder34(Feeder):
    """IEEE 34 com taps de regulador variáveis: Y é reconstruída (e guardada) para cada combinação de taps."""

    REGS = ["reg1a", "reg1b", "reg1c", "reg2a", "reg2b", "reg2c", "xfm1"]

    def __init__(self, path):
        super().__init__(path, sub_meter_line="L1", sub_v_buses=("800",), source_buses=("SOURCEBUS",))
        self._ycache = {}
        self.Y_nom = self.Y

    def taps(self):
        out = []
        for r in self.REGS:
            dss.Transformers.Name(r); dss.Transformers.Wdg(2); out.append(round(dss.Transformers.Tap(), 5))
        return tuple(out)

    def set_taps(self, taps):
        for r, t in zip(self.REGS, taps):
            dss.Transformers.Name(r); dss.Transformers.Wdg(2); dss.Transformers.Tap(t)

    def Y_for(self, taps):
        if taps not in self._ycache:
            cur = self.taps()
            dss.Text.Command("set controlmode=off")
            self.set_taps(taps)
            dss.Text.Command("batchedit load..* enabled=no"); dss.Solution.Solve()
            names = [n.upper() for n in dss.Circuit.YNodeOrder()]
            assert names == self.nodes
            n = len(names)
            self._ycache[taps] = _c(dss.Circuit.SystemY()).reshape(n, n)
            dss.Text.Command("batchedit load..* enabled=yes")
            self.set_taps(cur)
        return self._ycache[taps]

    def solve_step(self, kw, kvar, reg2_auto=True, reg2_fixed=1.0, reg2_vreg=124.0, xfm1_tap=1.0):
        """Fluxo com cargas dadas. Reguladores 1 sempre automáticos; regulador 2 automático ou travado."""
        for L, p, q in zip(self.loads, kw, kvar):
            dss.Loads.Name(L["name"]); dss.Loads.kW(p); dss.Loads.kvar(q)
        dss.Text.Command("set controlmode=static")
        for r in ("creg2a", "creg2b", "creg2c"):
            dss.RegControls.Name(r); dss.CktElement.Enabled(reg2_auto); dss.RegControls.ForwardVreg(reg2_vreg)
        if not reg2_auto:
            for r in ("reg2a", "reg2b", "reg2c"):
                dss.Transformers.Name(r); dss.Transformers.Wdg(2); dss.Transformers.Tap(reg2_fixed)
        dss.Transformers.Name("xfm1"); dss.Transformers.Wdg(2); dss.Transformers.Tap(xfm1_tap)
        dss.Solution.Solve()
        assert dss.Solution.Converged()
        return self._read_v(), self.taps()


def reading_voltages(fd, V, points):
    """Tensões de leitura (p.u.) em cada ponto de atendimento: entre fases se a carga é trifásica/delta em 24,9 kV
    ou 4,16 kV; fase-neutro se é monofásica estrela. points: lista de (bus, tipo, nós)."""
    out = []
    for bus, kind, nodes in points:
        ks = [fd.idx[n] for n in nodes]
        if kind == "ll":
            vv = [abs(V[..., ks[a]] - V[..., ks[(a + 1) % len(ks)]]) for a in range(len(ks))] if len(ks) == 3 else \
                 [abs(V[..., ks[0]] - V[..., ks[1]])]
            base = fd.vbase[ks[0]] * np.sqrt(3)
        else:
            vv = [abs(V[..., k]) for k in ks]; base = fd.vbase[ks[0]]
        out.append(np.stack(vv, axis=-1) / base)
    return out


def supply_points(fd):
    """Um ponto de atendimento por barra com carga."""
    pts = {}
    for L in fd.loads:
        b = L["bus"]
        for a, c in L["branches"]:
            pts.setdefault(b, set()).add(a)
            if c: pts[b].add(c)
    out = []
    for b, ns in sorted(pts.items()):
        ns = sorted(ns)
        dss.Circuit.SetActiveBus(b); kv = dss.Bus.kVBase()
        delta_or_3ph = any(c for L in fd.loads if L["bus"] == b for _, c in L["branches"]) or len(ns) == 3
        out.append((b, "ll" if delta_or_3ph and len(ns) >= 2 else "ln", ns))
    return out