"""
Rede de teste com média e baixa tensão para o Paper 3.

- MT: IEEE 34 do OpenDSS (fonte 1,03 p.u.), cargas pontuais = consumidores de MT.
- BT: cada carga distribuída do IEEE 34 (agrupada por barra e conexão) vira um ou mais transformadores monofásicos
  com derivação central (127/254 V), rede secundária radial e consumidores de 127 V, cada um com a sua curva.
- Dois mundos: VERDADE (rede como ela está, com o problema) e CADASTRO (o que a distribuidora acredita).
  O estimador da MT usa o modelo da MT do cadastro; a tensão na BT é estimada propagando a tensão estimada da MT
  pelo modelo de BT do cadastro com as cargas previstas.
"""
import os, copy, math
import numpy as np
import opendssdirect as dss

HERE = os.path.dirname(os.path.abspath(__file__))
FEEDER = os.path.join(HERE, "feeders_p3", "ieee34_pu1.03.dss")

KVA_STD = [5, 10, 15, 25, 37.5, 50, 75, 100]
COND = {"1/0": (0.54, 0.09), "2": (0.85, 0.10), "4": (1.35, 0.11)}   # ohm/km (fase), multiplex de alumínio (aprox.)
NEUTRAL_FACTOR = 1.5     # corrente de retorno pelo neutro (aprox.): r_eff = r x 1,5
STEPS = 1008

# Faixas de tensão (PRODIST Módulo 8, Anexo 8.A — a conferir no texto oficial)
BT_V = dict(adq=(117.0, 133.0), prec_lo=110.0, prec_hi=135.0)          # 127 V
MT_PU = dict(adq=(0.93, 1.05), prec_lo=0.90)                             # 2,3 kV <= Vn < 69 kV


def classify(v, kind):
    if kind == "bt":
        crit = (v < BT_V["prec_lo"]) | (v > BT_V["prec_hi"])
        prec = ~crit & ((v < BT_V["adq"][0]) | (v > BT_V["adq"][1]))
    else:
        crit = (v < MT_PU["prec_lo"]) | (v > MT_PU["adq"][1])
        prec = ~crit & (v < MT_PU["adq"][0])
    return prec, crit


def drp_drc(v, kind):
    """v: (STEPS, nfases). Retorna (DRP, DRC) em %: maior contagem entre as fases (PRODIST)."""
    v = np.asarray(v, float)
    if v.ndim == 1: v = v[:, None]
    ok = ~np.isnan(v).any(axis=1); v = v[ok]
    prec, crit = classify(v, kind)
    return 100 * prec.sum(axis=0).max() / len(v), 100 * crit.sum(axis=0).max() / len(v)


# ---------------------------------------------------------------- curvas
def shapes():
    h = np.arange(144) / 6
    res = 0.30 + 0.15 * np.exp(-((h - 7.0) / 1.2) ** 2) + 0.70 * np.exp(-((h - 19.5) / 1.8) ** 2) + 0.10 * np.exp(-((h - 12.5) / 1.5) ** 2)
    com = 0.25 + 0.75 / (1 + np.exp(-(h - 8.0) * 2)) / (1 + np.exp((h - 18.5) * 2))
    ind = 0.50 + 0.50 / (1 + np.exp(-(h - 6.5) * 3)) / (1 + np.exp((h - 22.0) * 3))
    S = {"res": res / res.max(), "com": com / com.max(), "ind": ind / ind.max()}
    sun = np.clip(np.sin(np.pi * (h - 6) / 12), 0, None) ** 1.3
    return S, {"res": 1.05, "com": 0.55, "ind": 0.75}, sun


def week_profiles(classes, rng, start_dow=0, noise=0.15, rho=0.92):
    """Verdade (STEPS x n) e curva típica que a distribuidora conhece (STEPS x n), mais irradiância verdadeira e típica."""
    S, wk, sun = shapes()
    n = len(classes); out = np.empty((STEPS, n)); typ = np.empty((STEPS, n)); irr = np.empty(STEPS); irr_typ = np.empty(STEPS)
    e = np.zeros(n); cl = 0.0
    for d in range(7):
        we = (start_dow + d) % 7 >= 5
        day = np.clip(1 + 0.07 * rng.standard_normal(), 0.8, 1.2)
        cloud_day = rng.uniform(0.45, 1.0)
        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 + math.sqrt(1 - rho ** 2) * noise * rng.standard_normal(n)
            out[k] = base * day * (1 + e); typ[k] = base
            cl = 0.8 * cl + 0.2 * rng.standard_normal()
            irr[k] = sun[t] * np.clip(cloud_day + 0.15 * cl, 0.05, 1.0); irr_typ[k] = sun[t] * 0.75
    return np.clip(out, 0.03, None), typ, irr, irr_typ


# ---------------------------------------------------------------- especificação da rede (cadastro base)
def build_spec(seed=1):
    """Lê o IEEE 34 e monta os grupos de BT. A especificação é o 'cadastro' de base."""
    rng = np.random.default_rng(seed)
    dss.Text.Command("clear"); os.chdir(os.path.dirname(FEEDER)); dss.Text.Command(f"compile [{FEEDER}]")
    groups = {}; mt = []
    i = dss.Loads.First()
    while i:
        nm = dss.Loads.Name(); bus = dss.CktElement.BusNames()[0]
        kw, kvar = dss.Loads.kW(), dss.Loads.kvar()
        if nm.startswith("d"):
            no = dss.CktElement.NodeOrder(); b = bus.split(".")[0]
            ll = dss.Loads.IsDelta() and len(no) > 1 and no[1] != 0
            nodes = (no[0], no[1]) if ll else (no[0],)
            key = (b, nodes)
            g = groups.setdefault(key, dict(bus=b, nodes=nodes, ll=ll, kv=24.9 if ll else 14.376, kw=0.0, kvar=0.0, members=[]))
            g["kw"] += kw; g["kvar"] += kvar; g["members"].append(nm)
        else:
            mt.append(dict(name=nm, bus=bus, kw=kw, kvar=kvar, cls=str(rng.choice(["ind", "com"]))))
        i = dss.Loads.Next()
    G = []
    for gi, (key, g) in enumerate(sorted(groups.items())):
        nt = max(1, math.ceil(g["kw"] / 40.0))
        trafos = []
        for ti in range(nt):
            kw_t = g["kw"] / nt; kvar_t = g["kvar"] / nt
            nc = int(np.clip(round(kw_t / 3.0), 1, 12))
            nsp = int(np.clip(math.ceil(nc / 3), 1, 4))
            kva = next(k for k in KVA_STD if k >= 1.15 * math.hypot(kw_t, kvar_t)) if 1.15 * math.hypot(kw_t, kvar_t) <= 100 else 100
            cons = []
            share = rng.dirichlet(np.ones(nc) * 3.0)
            for ci in range(nc):
                cons.append(dict(name=f"c{gi}_{ti}_{ci}", span=int(min(nsp, 1 + ci // 3)), leg=1 + ci % 2,
                                 kw=float(kw_t * share[ci]), kvar=float(kvar_t * share[ci]),
                                 cls=str(rng.choice(["res", "com"], p=[0.8, 0.2]))))
            trafos.append(dict(name=f"t{gi}_{ti}", kva=kva, tap=1.0, cond="1/0",
                               spans=[float(rng.uniform(0.03, 0.045)) for _ in range(nsp)], cons=cons))
        G.append(dict(g, id=gi, trafos=trafos))
    return dict(groups=G, mt=mt, regs=dict(vreg1=122.0, vreg2=122.0), xfm1_tap=1.025, s890_kw=250.0,
                caps={"c844": True, "c848": True}, line_scale={}, pv_mt=[], extra_caps=[])


def consumers(spec):
    out = []
    for g in spec["groups"]:
        for t in g["trafos"]:
            for c in t["cons"]:
                out.append((g, t, c))
    return out


# ---------------------------------------------------------------- comandos OpenDSS
def dss_commands(spec, bt=True):
    """Comandos para montar o circuito (verdade ou cadastro) a partir da especificação."""
    cmd = [f"compile [{FEEDER}]",
           f"edit regcontrol.creg1a vreg={spec['regs']['vreg1']}", f"edit regcontrol.creg1b vreg={spec['regs']['vreg1']}",
           f"edit regcontrol.creg1c vreg={spec['regs']['vreg1']}"]
    for ph in "abc":
        v = spec["regs"].get(f"vreg2{ph}", spec["regs"]["vreg2"])
        cmd.append(f"edit regcontrol.creg2{ph} vreg={v}")
        if spec["regs"].get(f"stuck2{ph}") is not None:
            cmd += [f"edit regcontrol.creg2{ph} enabled=no", f"edit transformer.reg2{ph} wdg=2 tap={spec['regs'][f'stuck2{ph}']}"]
    cmd.append(f"edit transformer.xfm1 wdg=2 tap={spec['xfm1_tap']}")
    cmd.append(f"edit load.s890 kW={spec['s890_kw']} kvar={spec['s890_kw'] * 0.5}")
    for c, on in spec["caps"].items():
        if not on: cmd.append(f"edit capacitor.{c} enabled=no")
    lerr = spec.get("line_err", {}); lbase = spec.get("line_base", {})
    for ln in set(spec["line_scale"]) | set(lerr):
        L = spec["line_scale"][ln]["len"] if ln in spec["line_scale"] else lbase[ln]
        cmd.append(f"edit line.{ln} length={L * lerr.get(ln, 1.0)}")   # recondutoramento e/ou erro de cadastro da impedância
    for k, cap in enumerate(spec["extra_caps"]):
        cmd.append(f"new capacitor.cx{k} bus1={cap['bus']} phases=3 kvar={cap['kvar']} kv=24.9")
    for k, pv in enumerate(spec["pv_mt"] if bt else []):
        cmd.append(f"new generator.pvmt{k} bus1={pv['bus']} phases=3 kV=24.9 kW=0 kvar=0 model=1")
    if not bt:
        return cmd
    for g in spec["groups"]:
        for m in g["members"]:
            cmd.append(f"edit load.{m} enabled=no")
        prim = f"{g['bus']}." + ".".join(str(x) for x in g["nodes"]) + ("" if g["ll"] else ".0")
        kvp = 24.9 if g["ll"] else 14.376
        for t in g["trafos"]:
            tp = t.get("prim", prim)
            sec = f"{t['name']}_0"
            cmd.append(f"new transformer.{t['name']} phases=1 windings=3 buses=[{tp} {sec}.1.0 {sec}.0.2] conns=[wye wye wye] "
                       f"kvs=[{kvp} 0.127 0.127] kvas=[{t['kva']} {t['kva']} {t['kva']}] %rs=[0.6 1.2 1.2] xhl=2.04 xht=2.04 xlt=1.36 "
                       f"taps=[1 {t['tap']} {t['tap']}] %noloadloss=0.2")
            r, x = COND[t["cond"]]; r *= NEUTRAL_FACTOR
            prev = sec
            for k, L in enumerate(t["spans"], start=1):
                nb = f"{t['name']}_{k}"
                cmd.append(f"new line.{t['name']}_l{k} phases=2 bus1={prev}.1.2 bus2={nb}.1.2 length={L * t.get('len_mult', 1.0)} units=km "
                           f"r1={r} x1={x} r0={r} x0={x} c1=0 c0=0")
                prev = nb
            for c in t["cons"]:
                node = f"{t['name']}_{c['span']}"
                if c.get("bad_r"):
                    cb = f"{c['name']}_x"
                    cmd.append(f"new line.{c['name']}_ramal phases=1 bus1={node}.{c['leg']} bus2={cb}.1 length=1 units=km r1={c['bad_r']} x1=0.001 r0={c['bad_r']} x0=0.001 c1=0 c0=0")
                    bus1 = f"{cb}.1"
                else:
                    bus1 = f"{node}.{c['leg']}"
                cmd.append(f"new load.{c['name']} phases=1 bus1={bus1} kV=0.127 kW={c['kw']} kvar={c['kvar']} model=1 vminpu=0.85 vmaxpu=1.2")
                if c.get("pv_kw"):
                    cmd.append(f"new generator.pv_{c['name']} phases=1 bus1={bus1} kV=0.127 kW=0 kvar=0 model=1")
    cmd.append("set controlmode=static"); cmd.append("set maxcontroliter=40"); cmd.append("set maxiterations=100")
    return cmd


def consumer_node(c, t):
    return (f"{c['name']}_x.1" if c.get("bad_r") else f"{t['name']}_{c['span']}.{c['leg']}").upper()


class Circuit:
    """Um contexto OpenDSS independente (verdade ou cadastro)."""

    def __init__(self, spec, bt=True):
        self.d = dss.NewContext()
        self.spec = spec
        os.chdir(os.path.dirname(FEEDER))
        for c in dss_commands(spec, bt):
            self.d.Text.Command(c)
        self.d.Solution.Solve()
        self.cons = consumers(spec) if bt else []
        self.load_names = [c["name"] for _, _, c in self.cons]
        self.mt = spec["mt"]
        self.nodes = [n.upper() for n in self.d.Circuit.YNodeOrder()]
        self.idx = {n: k for k, n in enumerate(self.nodes)}

    def set_loads(self, kw_c, kvar_c, kw_mt, kvar_mt, pv_c=None, pv_mt=None):
        L = self.d.Loads
        for nm, p, q in zip(self.load_names, kw_c, kvar_c):
            L.Name(nm); L.kW(p); L.kvar(q)
        for m, p, q in zip(self.mt, kw_mt, kvar_mt):
            L.Name(m["name"]); L.kW(p); L.kvar(q)
        G = self.d.Generators
        if pv_c is not None:
            for (g, t, c), p in zip(self.cons, pv_c):
                if c.get("pv_kw"): G.Name(f"pv_{c['name']}"); G.kW(p); G.kvar(0.0)
        if pv_mt is not None:
            for k, (p, q) in enumerate(pv_mt):
                G.Name(f"pvmt{k}"); G.kW(p); G.kvar(q)

    def solve(self):
        self.d.Solution.Solve()
        if not self.d.Solution.Converged():
            return None
        v = np.asarray(self.d.Circuit.YNodeVArray(), float)
        return v[0::2] + 1j * v[1::2]

    def taps(self, regs=("reg1a", "reg1b", "reg1c", "reg2a", "reg2b", "reg2c", "xfm1")):
        T = self.d.Transformers; out = []
        for r in regs:
            T.Name(r); T.Wdg(2); out.append(round(T.Tap(), 5))
        return tuple(out)


def base_loads(spec):
    C = consumers(spec)
    kwc = np.array([c["kw"] for _, _, c in C]); kvc = np.array([c["kvar"] for _, _, c in C])
    kwm = np.array([spec["s890_kw"] if m["name"] == "s890" else m["kw"] for m in spec["mt"]])
    kvm = np.array([spec["s890_kw"] * 0.5 if m["name"] == "s890" else m["kvar"] for m in spec["mt"]])
    cls = [c["cls"] for _, _, c in C] + [m["cls"] for m in spec["mt"]]
    return C, kwc, kvc, kwm, kvm, cls


def calibrate_taps(spec, seed=99, rounds=2, margin=2.0):
    """Planejamento de base: escolhe o tap de cada transformador para centrar a faixa de tensão dos seus
    consumidores numa semana típica (cargas do cadastro). Taps de 0,95 a 1,10 em passos de 2,5 %."""
    taps = np.arange(0.95, 1.1001, 0.025)
    for _ in range(rounds):
        ck = Circuit(spec); C, kwc, kvc, kwm, kvm, cls = base_loads(spec); nC = len(C)
        cur, typ, irr, _ = week_profiles(cls, np.random.default_rng(seed))
        V = []
        for k in range(0, STEPS, 2):
            f = cur[k]; ck.set_loads(kwc * f[:nC], kvc * f[:nC], kwm * f[nC:], kvm * f[nC:]); V.append(ck.solve())
        V = np.abs(np.array(V))
        for g in spec["groups"]:
            for t in g["trafos"]:
                vv = np.concatenate([V[:, ck.idx[consumer_node(c, t)]] for c in t["cons"]])
                lo, hi = np.percentile(vv, 0.5), np.percentile(vv, 99.5)
                best = max(taps, key=lambda tp: min(lo * tp / t["tap"] - 117 - margin, 133 - margin - hi * tp / t["tap"]))
                t["tap"] = float(round(best, 3))
    return spec


MT_CUSTOMERS = {"830": ["s830a", "s830b", "s830c"], "840": ["s840"], "844": ["s844"], "848": ["s848"], "860": ["s860"], "890": ["s890"]}
MT_BASE_LL = {"890": 4160.0}


def mt_reading(V, idx, bus):
    """Tensões entre fases (p.u. da nominal) na barra do consumidor de MT. V: (T, n) complexo."""
    ks = [idx[f"{bus}.{p}"] for p in (1, 2, 3)]
    vb = MT_BASE_LL.get(bus, 24900.0)
    return np.stack([np.abs(V[:, ks[a]] - V[:, ks[(a + 1) % 3]]) for a in range(3)], axis=1) / vb


def mt_wrapper(spec, path):
    """Arquivo .dss com o modelo de MT do cadastro (sem BT), para o estimador."""
    cmd = dss_commands(spec, bt=False)
    cmd[0] = f"redirect [{FEEDER}]"
    open(path, "w").write("\n".join(cmd) + "\n")
    return path