"""
Catálogo de problemas de tensão e consertos. Cada caso gera três especificações:
verdade antes (com o problema), verdade depois (com o conserto) e cadastro depois (o que a distribuidora acredita).
"""
import copy
import numpy as np
from rede import consumers

MT_TYPES = ["reg_travado_3f", "reg_travado_1f", "reg_ajuste", "capacitor_desligado", "crescimento_mt", "desequilibrio_mt", "fv_mt"]
BT_TYPES = ["trafo_sobrecarregado", "tap_errado", "rede_bt_ruim", "desequilibrio_bt", "conexao_ruim", "fv_bt"]
TYPES = MT_TYPES + BT_TYPES
LABEL = {"reg_travado_3f": "Regulador travado (3 fases)", "reg_travado_1f": "Regulador com uma fase travada",
         "reg_ajuste": "Regulador mal ajustado", "capacitor_desligado": "Banco de capacitores desligado",
         "crescimento_mt": "Carga acima da capacidade da MT", "desequilibrio_mt": "Desequilíbrio de fases na MT",
         "fv_mt": "Usina FV na MT (sobretensão)", "trafo_sobrecarregado": "Transformador sobrecarregado",
         "tap_errado": "Tap do transformador errado", "rede_bt_ruim": "Rede secundária longa/fina",
         "desequilibrio_bt": "Desequilíbrio na BT", "conexao_ruim": "Conexão/ramal ruim", "fv_bt": "FV na BT (sobretensão)"}
FIX_P = {"completo": 0.5, "parcial": 0.25, "errado": 0.25}

# barras a jusante do regulador 2 (onde problemas do regulador aparecem) e regiões para crescimento de carga
DOWN_REG2 = {"832", "858", "834", "860", "836", "840", "842", "844", "846", "848", "862", "838", "864", "888", "890"}
REGIONS = {"cliente_844": (["s844"], [("852R", "832"), ("832", "858"), ("858", "834"), ("834", "842"), ("842", "844")]),   # (consumidores de MT que crescem, trechos recondutorados)
           "cliente_848": (["s848"], [("852R", "832"), ("832", "858"), ("858", "834"), ("834", "842"), ("842", "844"), ("844", "846"), ("846", "848")]),
           "clientes_840_860": (["s840", "s860"], [("852R", "832"), ("832", "858"), ("858", "834"), ("834", "860"), ("860", "836"), ("836", "840")]),
           "cliente_830": (["s830a", "s830b", "s830c"], [("824", "828"), ("828", "830")])}
THREE_PH = {"802", "806", "808", "812", "814", "850", "816", "824", "828", "830", "854", "852", "832", "858", "834", "842",
            "844", "846", "848", "860", "836", "840"}


def _trafos(spec):
    return [(g, t) for g in spec["groups"] for t in g["trafos"]]


def _find(spec, name):
    for g, t in _trafos(spec):
        if t["name"] == name: return g, t


def _lines_from(path_buses):
    """Linhas do IEEE 34 entre barras do caminho (por nome de barra; resolvido no circuito)."""
    return path_buses


def make_case(ctype, spec0, rng, line_len, spec_cad0=None):
    """Retorna dict(tipo, conserto, before, after, cad, info). line_len: comprimentos originais das linhas (km/kft)."""
    B = copy.deepcopy(spec0)                    # verdade antes
    info = {}
    fix = rng.choice(list(FIX_P), p=list(FIX_P.values()))
    downstream = lambda g: g["bus"] in DOWN_REG2
    tr = _trafos(B)

    # ------------------------------------------------ problema (verdade antes)
    if ctype == "reg_travado_3f":
        tap = float(rng.choice([0.95, 0.9625, 0.975])); info["tap"] = tap
        for ph in "abc": B["regs"][f"stuck2{ph}"] = tap
    elif ctype == "reg_travado_1f":
        ph = str(rng.choice(list("abc"))); tap = float(rng.choice([0.95, 0.9625, 0.975])); info.update(fase=ph, tap=tap)
        B["regs"][f"stuck2{ph}"] = tap
    elif ctype == "reg_ajuste":
        v = float(rng.choice([114.0, 115.0, 116.0, 129.0, 130.0])); info["vreg2"] = v; B["regs"]["vreg2"] = v
    elif ctype == "capacitor_desligado":
        caps = [["c844"], ["c848"], ["c844", "c848"]][rng.integers(3)]; info["caps"] = caps
        for c in caps: B["caps"][c] = False
        _grow_all(B, rng.uniform(1.15, 1.4))            # crescimento geral moderado (desconhecido)
    elif ctype == "crescimento_mt":
        reg = str(rng.choice(list(REGIONS))); info["regiao"] = reg; f = rng.uniform(3.0, 6.0); info["fator"] = f
        for m in B["mt"]:
            if m["name"] in REGIONS[reg][0]: m["kw"] *= f; m["kvar"] *= f
    elif ctype == "desequilibrio_mt":
        cand = [(g, t) for g, t in tr if g["bus"] in THREE_PH and not g["ll"]]
        ph = int(rng.integers(1, 4)); info["fase"] = ph
        pick = [cand[i] for i in rng.choice(len(cand), size=min(len(cand), int(rng.integers(10, 18))), replace=False)]
        info["trafos"] = [t["name"] for _, t in pick]
        for g, t in pick:
            t["prim"] = f"{g['bus']}.{ph}.0"
    elif ctype == "fv_mt":
        bus = str(rng.choice(["840", "848", "860", "844"])); cap = float(rng.uniform(2500, 5000))
        B["pv_mt"].append(dict(bus=bus, kw=cap, pf=1.0)); info.update(bus=bus, kw=cap)
    else:
        cand = [(g, t) for g, t in tr if len(t["cons"]) >= 3] if ctype != "conexao_ruim" else tr
        g, t = cand[rng.integers(len(cand))]; info["trafo"] = t["name"]; info["barra"] = g["bus"]
        if ctype == "trafo_sobrecarregado":
            f = rng.uniform(2.0, 3.0); info["fator"] = f
            for c in t["cons"]: c["kw"] *= f; c["kvar"] *= f
        elif ctype == "tap_errado":
            d = float(rng.choice([-0.05, -0.05, -0.075, +0.05])); info["delta_tap"] = d; t["tap"] = round(t["tap"] + d, 4)
        elif ctype == "rede_bt_ruim":
            m = rng.uniform(1.8, 2.8); info["len_mult"] = m; t["len_mult"] = m; t["cond"] = "4"
            for c in t["cons"]: c["kw"] *= 1.2; c["kvar"] *= 1.2
        elif ctype == "desequilibrio_bt":
            for c in t["cons"]:
                if rng.random() < 0.9: c["leg"] = 1
                c["kw"] *= 1.5; c["kvar"] *= 1.5
        elif ctype == "conexao_ruim":
            c = t["cons"][rng.integers(len(t["cons"]))]; c["bad_r"] = float(rng.uniform(0.3, 0.9)); info["consumidor"] = c["name"]
        elif ctype == "fv_bt":
            for c in t["cons"]:
                if rng.random() < 0.8: c["pv_kw"] = float(rng.uniform(5, 9))
            if t["tap"] < 1.0: pass
            info["pv_total"] = sum(c.get("pv_kw", 0) for c in t["cons"])

    # ------------------------------------------------ conserto (verdade depois) e cadastro depois
    A = copy.deepcopy(B)          # verdade depois
    Cd = copy.deepcopy(spec_cad0 if spec_cad0 is not None else spec0)     # cadastro: rede de base + o que a distribuidora sabe que fez
    info["conserto"] = str(fix)

    def bump_taps(specs, names, d):
        for s in specs:
            for g, t in _trafos(s):
                if t["name"] in names: t["tap"] = round(t["tap"] + d, 4)

    def vreg_delta(specs, d, which="vreg2"):
        for s in specs: s["regs"][which] = s["regs"][which] + d

    if ctype in ("reg_travado_3f", "reg_travado_1f"):
        if fix == "completo":
            for ph in "abc": A["regs"].pop(f"stuck2{ph}", None)
        elif fix == "parcial":
            for ph in "abc":
                if A["regs"].get(f"stuck2{ph}") is not None: A["regs"][f"stuck2{ph}"] = round(A["regs"][f"stuck2{ph}"] + 0.025, 4)
            for ph in "abc":
                if A["regs"].get(f"stuck2{ph}") is not None: Cd["regs"][f"stuck2{ph}"] = A["regs"][f"stuck2{ph}"]
        info["wrong_target"] = "trafos_reclamantes"
    elif ctype == "reg_ajuste":
        if fix == "completo": A["regs"]["vreg2"] = 122.0
        elif fix == "parcial": A["regs"]["vreg2"] = (A["regs"]["vreg2"] + 122.0) / 2
        Cd["regs"]["vreg2"] = A["regs"]["vreg2"]
        info["wrong_target"] = "trafos_reclamantes"
    elif ctype == "capacitor_desligado":
        if fix == "completo":
            for c in info["caps"]: A["caps"][c] = True
        elif fix == "parcial":
            A["caps"][info["caps"][0]] = True
        info["wrong_target"] = "trafos_reclamantes"
    elif ctype == "crescimento_mt":
        if fix == "completo":
            for s in (A, Cd):
                for ln, (b1, b2, L) in line_len.items():
                    if (b1, b2) in REGIONS[info["regiao"]][1] or (b2, b1) in REGIONS[info["regiao"]][1]:
                        s["line_scale"][ln] = dict(len=L * 0.45)
        elif fix == "parcial":
            cb = {"cliente_844": "844", "cliente_848": "848", "clientes_840_860": "860", "cliente_830": "830"}[info["regiao"]]
            for s in (A, Cd): s["extra_caps"].append(dict(bus=cb, kvar=float(rng.choice([300, 450, 600]))))
        info["wrong_target"] = "trafos_reclamantes"
    elif ctype == "desequilibrio_mt":
        names = info["trafos"]
        keep = names if fix == "completo" else names[: len(names) // 2] if fix == "parcial" else []
        for g, t in _trafos(A):
            if t["name"] in keep: t.pop("prim", None)
        if fix == "errado": vreg_delta([A, Cd], +2.0)
    elif ctype == "fv_mt":
        if fix == "completo":
            A["pv_mt"][0]["pf"] = -0.90
        elif fix == "parcial":
            vreg_delta([A, Cd], -2.0)
        info["wrong_target"] = "trafos_reclamantes_baixa"
        Cd["pv_mt"] = copy.deepcopy(A["pv_mt"])        # usina conhecida no cadastro (capacidade e fp)
    else:
        name = info["trafo"]
        gA, tA = _find(A, name); gC, tC = _find(Cd, name)
        if ctype == "trafo_sobrecarregado":
            if fix == "completo":
                for t in (tA, tC): t["kva"] = min(150, t["kva"] * 2)
                for t in (tA, tC):            # novo transformador maior + divisão do circuito (metade dos consumidores no vão 1)
                    for c in t["cons"]: c["span"] = 1
            elif fix == "parcial": bump_taps([A, Cd], [name], +0.025)
        elif ctype == "tap_errado":
            if fix == "completo": tA["tap"] = tC["tap"]
            elif fix == "parcial": tA["tap"] = round(tA["tap"] - info["delta_tap"] / 2, 4); tC["tap"] = round(tC["tap"] - info["delta_tap"] / 2, 4)
        elif ctype == "rede_bt_ruim":
            if fix == "completo":
                for t in (tA, tC): t["cond"] = "1/0"
                bump_taps([A, Cd], [name], +0.025)
            elif fix == "parcial": bump_taps([A, Cd], [name], +0.025)
        elif ctype == "desequilibrio_bt":
            orig = {c["name"]: c["leg"] for c in tC["cons"]}
            for i, c in enumerate(tA["cons"]):
                if fix == "completo" or (fix == "parcial" and i % 2 == 0): c["leg"] = orig[c["name"]]
        elif ctype == "conexao_ruim":
            c = next(c for c in tA["cons"] if c["name"] == info["consumidor"])
            if fix == "completo": c.pop("bad_r")
            elif fix == "parcial": c["bad_r"] *= 0.4
        elif ctype == "fv_bt":
            if fix == "completo":
                for c in tA["cons"]:
                    if c.get("pv_kw"): c["pv_kw"] *= 0.3             # limitação de injeção nos inversores
            elif fix == "parcial": bump_taps([A, Cd], [name], -0.025)
        info["wrong_target"] = "vreg_baixa" if ctype == "fv_bt" else "vreg"
        if fix == "errado":
            if info["wrong_target"] == "vreg_baixa": vreg_delta([A, Cd], -2.0, "vreg2" if gA["bus"] in DOWN_REG2 else "vreg1")
            else: vreg_delta([A, Cd], +2.0, "vreg2" if gA["bus"] in DOWN_REG2 else "vreg1")
    return dict(tipo=ctype, conserto=str(fix), before=B, after=A, cad=Cd, info=info)


def _grow_all(spec, f):
    for g in spec["groups"]:
        for t in g["trafos"]:
            for c in t["cons"]: c["kw"] *= f; c["kvar"] *= f
    for m in spec["mt"]: m["kw"] *= f; m["kvar"] *= f


def apply_wrong_fix(case, compensated_trafos):
    """Conserto no lugar errado para problemas de MT: troca dos transformadores dos reclamantes por outros maiores."""
    for s in (case["after"], case["cad"]):
        for g in s["groups"]:
            for t in g["trafos"]:
                if t["name"] in compensated_trafos:
                    t["kva"] = min(150, t["kva"] * 2)


def background_errors(spec0, rng, line_len, imp_sd=0.03, phase_frac=0.10):
    """Rede verdadeira de base: o cadastro (spec0) com erros que a distribuidora não conhece — impedância das linhas
    de MT (fator lognormal por linha) e fase de ligação de parte dos transformadores monofásicos em barras trifásicas."""
    T = copy.deepcopy(spec0)
    T["line_base"] = {ln: L for ln, (_, _, L) in line_len.items()}
    T["line_err"] = {ln: float(np.exp(imp_sd * rng.standard_normal())) for ln in line_len}
    moved = []
    for g in T["groups"]:
        if g["bus"] in THREE_PH and not g["ll"]:
            for t in g["trafos"]:
                if rng.random() < phase_frac:
                    ph = int(rng.choice([p for p in (1, 2, 3) if p != g["nodes"][0]]))
                    t["prim"] = f"{g['bus']}.{ph}.0"; moved.append(t["name"])
    return T, moved
