"""Passo 2: alocação gulosa, comparações, modos de sensor e sincronismo num alimentador."""
import json, sys, os
import numpy as np
import opendssdirect as dss
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from se3f import Feeder, scenarios, Estimator, EstimatorSCADA, metrics
import itertools
FEEDERS = os.path.abspath(os.environ.get("G4W_FEEDERS", "feeders"))  # pasta IEEETestCases do electricdss-tst

MODES = {"VI": dict(mode="VI", sig_v=0.002), "I": dict(mode="I"), "VI_lowV": dict(mode="VI", sig_v=0.015)}


def evaluate(fd, sensors, V, rng, Vref, scada=False, floor_scale=1.0, **kw):
    est = EstimatorSCADA(fd, sensors, **kw) if scada else Estimator(fd, sensors, **kw)
    if floor_scale != 1.0: (est.lin if scada else est).floor_scale = floor_scale
    est.gain(Vref)
    return metrics(fd, V, est.estimate(est.measure(V, rng)))


def topology(fd):
    edges = [(L["bus1"], L["bus2"]) for L in fd.lines.values()]
    i = dss.Transformers.First()
    while i:
        b = [x.split(".")[0].upper() for x in dss.CktElement.BusNames()]
        edges.append((b[0], b[1])); i = dss.Transformers.Next()
    return edges


def hops(fd, edges, root):
    adj = {}
    for a, b in edges: adj.setdefault(a, set()).add(b); adj.setdefault(b, set()).add(a)
    d = {root: 0}; q = [root]
    while q:
        u = q.pop(0)
        for v in adj.get(u, ()):
            if v not in d: d[v] = d[u] + 1; q.append(v)
    return d


def run(name, path, sub_line, sub_v, root="SOURCEBUS", N_EVAL=400, N_GREEDY=120, KMAX=8, seed=2026, out="."):
    fd = Feeder(path, sub_meter_line=sub_line, sub_v_buses=sub_v, source_buses=(root,))
    # barra 610 do IEEE 123: secundário em delta, sem carga e sem referência de terra; a tensão fase-neutro
    # não é definida fisicamente (depende de admitâncias residuais), por isso fica fora da métrica de tensão
    fd.exclude_buses = ("610",) if name == "ieee123" else ()
    rng = np.random.default_rng(seed)
    _, Vg = scenarios(fd, N_GREEDY, rng)          # cenários para escolher
    _, Ve = scenarios(fd, N_EVAL, rng)            # cenários independentes para avaliar
    # referência a priori para dimensionar os desvios: fluxo de potência com as cargas previstas (nominais)
    Vref = fd.truth(np.ones(len(fd.loads))); fd.reset()
    cand = [l for l in fd.lines if l != fd.sub_meter_line and "_OPEN" not in fd.lines[l]["bus2"]]
    KMAX = min(KMAX, len(cand))
    edges = topology(fd); hp = hops(fd, edges, root)
    out_deg = {}
    for L in fd.lines.values(): out_deg[L["bus1"]] = out_deg.get(L["bus1"], 0) + 1
    leaves = sorted([l for l in cand if out_deg.get(fd.lines[l]["bus2"], 0) == 0], key=lambda l: -hp.get(fd.lines[l]["bus2"], 0))
    # heurística "uniforme no tronco": K sensores espaçados ao longo do caminho até a barra mais distante
    par = {root: None}; q = [root]; adj = {}
    for a, b in edges: adj.setdefault(a, set()).add(b); adj.setdefault(b, set()).add(a)
    while q:
        u = q.pop(0)
        for v in adj.get(u, ()):
            if v not in par: par[v] = u; q.append(v)
    far = max(hp, key=hp.get); path = []; u = far
    while par[u] is not None: path.append((par[u], u)); u = par[u]
    path = path[::-1]
    # trechos do tronco, em qualquer orientação de definição da linha
    trunk = [l for (a, b) in path for l in cand if {fd.lines[l]["bus1"], fd.lines[l]["bus2"]} == {a, b}]

    def trunk_pick(K):
        """K pontos distintos igualmente espaçados (em número de trechos) ao longo do tronco.
        Se K for maior que o número de trechos do tronco, usa todos (número efetivo menor que K)."""
        L = len(trunk)
        if K == 0: return []
        if K >= L: return list(trunk)
        idx = []
        for j in range(K):
            t = int(np.floor((j + 1) * L / (K + 1)))
            t = min(max(t, 0), L - 1)
            while t in idx: t = (t + 1) % L
            idx.append(t)
        return [trunk[i] for i in sorted(idx)]
    res = {"feeder": name, "excluded_buses": list(fd.exclude_buses), "nodes": len(fd.nodes), "loads": len(fd.loads), "candidates": cand, "modes": {}}

    m00 = evaluate(fd, [], Vg, np.random.default_rng(1), Vref)
    ref = (m00["v_max"].mean(), m00["i_max_rel"].mean(), m00["p_err_rel"].mean())

    def J(m):  # objetivo composto: tensão, corrente e carga, cada um relativo ao caso sem sensores
        return m["v_max"].mean() / ref[0] + m["i_max_rel"].mean() / ref[1] + m["p_err_rel"].mean() / ref[2]

    def greedy(kw):
        chosen = []
        for _ in range(KMAX):
            best = min((c for c in cand if c not in chosen),
                       key=lambda c: J(evaluate(fd, chosen + [c], Vg, np.random.default_rng(1), Vref, **kw)))
            chosen.append(best)
        return chosen

    def curve(order, kw):
        c = {"v": [], "i": [], "p": [], "ang": [], "iabs": [], "v_p90": [], "i_p90": [], "p_p90": []}
        for K in range(KMAX + 1):
            m = evaluate(fd, order[:K], Ve, np.random.default_rng(7), Vref, **kw)
            c["v_p90"].append(float(np.percentile(m["v_max"], 90))); c["i_p90"].append(float(np.percentile(m["i_max_rel"], 90)))
            c["p_p90"].append(float(np.percentile(m["p_err_rel"], 90)))
            c["v"].append(float(m["v_max"].mean())); c["i"].append(float(m["i_max_rel"].mean()))
            c["p"].append(float(m["p_err_rel"].mean())); c["ang"].append(float(m["ang_max"].mean()))
            c["iabs"].append(float(m["i_max_abs"].mean()))
        return c

    for mk, kw in MODES.items():
        g = greedy(kw)
        rnd = {"v": [], "p": [], "i": []}
        for K in range(KMAX + 1):
            vs, ps, is_ = [], [], []
            for t in range(30):
                pick = list(np.random.default_rng(100 + t).choice(cand, K, replace=False)) if K else []
                m = evaluate(fd, pick, Ve, np.random.default_rng(7), Vref, **kw)
                vs.append(m["v_max"].mean()); ps.append(m["p_err_rel"].mean()); is_.append(m["i_max_rel"].mean())
            for key, arr in (("v", vs), ("p", ps), ("i", is_)):
                rnd[key].append([float(np.mean(arr)), float(np.percentile(arr, 10)), float(np.percentile(arr, 90))])
        tr = {"v": [], "i": [], "p": [], "k_eff": [], "lines": []}
        for K in range(KMAX + 1):
            tp = trunk_pick(K); tr["k_eff"].append(len(set(tp))); tr["lines"].append(tp)
            m = evaluate(fd, tp, Ve, np.random.default_rng(7), Vref, **kw)
            tr["v"].append(float(m["v_max"].mean())); tr["i"].append(float(m["i_max_rel"].mean())); tr["p"].append(float(m["p_err_rel"].mean()))
        res["modes"][mk] = {"trunk": tr, "trunk_lines": trunk, "greedy_order": g, "greedy": curve(g, kw),
                             "leaf_order": leaves[:KMAX], "leaf": curve(leaves[:KMAX] + [c for c in cand if c not in leaves][:max(0, KMAX - len(leaves))], kw),
                             "random": rnd}
        print(name, mk, "ordem", g, "V:", np.round(res["modes"][mk]["greedy"]["v"], 3))
    # subestação como SCADA (|I| por fase e P, Q totais), mesma ordem de alocação (VI)
    gVI = res["modes"]["VI"]["greedy_order"]
    if os.environ.get("SKIP_SCADA") != "1":
        res["scada"] = curve(gVI, dict(scada=True))
        print(name, "SCADA V:", np.round(res["scada"]["v"], 3))
    # busca exaustiva (verificação do guloso) quando viável
    res["exhaustive"] = {}
    from math import comb
    for K in range(1, 5):
        if comb(len(cand), K) > int(os.environ.get("EXH_MAX", "40000")): continue
        best, bJ = None, 1e9
        for comb_ in itertools.combinations(cand, K):
            j = J(evaluate(fd, list(comb_), Vg, np.random.default_rng(1), Vref))
            if j < bJ: best, bJ = list(comb_), j
        gJ = J(evaluate(fd, gVI[:K], Vg, np.random.default_rng(1), Vref))
        me_b = evaluate(fd, best, Ve, np.random.default_rng(7), Vref); me_g = evaluate(fd, gVI[:K], Ve, np.random.default_rng(7), Vref)
        res["exhaustive"][K] = {"best": best, "J_best_sel": bJ, "J_greedy_sel": gJ, "n_comb": comb(len(cand), K),
                                "v_best": float(me_b["v_max"].mean()), "v_greedy": float(me_g["v_max"].mean()),
                                "i_best": float(me_b["i_max_rel"].mean()), "i_greedy": float(me_g["i_max_rel"].mean()),
                                "p_best": float(me_b["p_err_rel"].mean()), "p_greedy": float(me_g["p_err_rel"].mean())}
        print(name, "exaustiva K", K, res["exhaustive"][K])
    # sensibilidade ao piso da injeção nula (K = 0 e 4)
    res["floor"] = {}
    for fs in (0.01, 1.0, 100.0):
        res["floor"][fs] = [float(evaluate(fd, gVI[:K], Ve, np.random.default_rng(7), Vref, floor_scale=fs)["v_max"].mean()) for K in (0, 4)]
    # por barra / trecho: 0 vs 4 sensores (VI)
    g4 = res["modes"]["VI"]["greedy_order"][:4]
    m0 = evaluate(fd, [], Ve, np.random.default_rng(7), Vref); m4 = evaluate(fd, g4, Ve, np.random.default_rng(7), Vref)
    nodes = [fd.nodes[k] for k in range(len(fd.nodes)) if k not in fd.source]
    res["per_node"] = {"nodes": nodes, "v0": m0["v_node"].tolist(), "v4": m4["v_node"].tolist()}
    res["per_line"] = {"lines": m0["line_names"], "i0": m0["i_line_abs"].tolist(), "i4": m4["i_line_abs"].tolist(), "itrue": m0["i_line_true"].tolist()}
    res["loads_scatter"] = {"buses": m0["load_buses"], "p_true": m0["S_true"].real[:100].tolist(),
                            "p_est0": m0["S_est"].real[:100].tolist(), "p_est4": m4["S_est"].real[:100].tolist()}
    # precisão da tensão do sensor (mesma alocação, K = 4)
    res["sigv"] = {"pct": [0.1, 0.2, 0.5, 1.0, 2.0], "v": [], "p": []}
    for sv in res["sigv"]["pct"]:
        m = evaluate(fd, g4, Ve, np.random.default_rng(7), Vref, sig_v=sv / 100)
        res["sigv"]["v"].append(float(m["v_max"].mean())); res["sigv"]["p"].append(float(m["p_err_rel"].mean()))
    m = evaluate(fd, g4, Ve, np.random.default_rng(7), Vref, mode="I")
    res["sigv"]["only_i"] = [float(m["v_max"].mean()), float(m["p_err_rel"].mean())]
    # sincronismo
    res["sync"] = {"us": [0.1, 1, 10, 50, 100, 500, 1000], "v": [], "p": []}
    for s in res["sync"]["us"]:
        m = evaluate(fd, g4, Ve, np.random.default_rng(7), Vref, sync_us=s)
        res["sync"]["v"].append(float(m["v_max"].mean())); res["sync"]["p"].append(float(m["p_err_rel"].mean()))
    # geometria para o desenho
    res["edges"] = edges
    res["lines_buses"] = {k: [L["bus1"], L["bus2"]] for k, L in fd.lines.items()}
    res["load_info"] = [{"bus": L["bus"], "kw": L["kw"], "nph": len(L["branches"]), "phases": sorted({a.split('.')[1] for a, b in L["branches"]} | {b.split('.')[1] for a, b in L["branches"] if b})} for L in fd.loads]
    regs = []
    i = dss.RegControls.First()
    while i:
        dss.Circuit.SetActiveElement("Transformer." + dss.RegControls.Transformer())
        regs.append([x.split(".")[0].upper() for x in dss.CktElement.BusNames()]); i = dss.RegControls.Next()
    res["regs"] = regs
    json.dump(res, open(os.path.join(out, f"results_{name}.json"), "w"), indent=1)
    return res


CASES = {
    "ieee13": (FEEDERS + "/13Bus/IEEE13Nodeckt.dss", "650632", ("650",), "SOURCEBUS"),
    "ieee34": (FEEDERS + "/34Bus/ieee34Mod1.dss", "L1", ("800",), "SOURCEBUS"),
    "ieee123": (FEEDERS + "/123Bus/IEEE123Master.dss", "L115", ("149",), "150"),
}
if __name__ == "__main__":
    out = sys.argv[1]
    for nm in sys.argv[2:]:
        p, sl, sv, root = CASES[nm]
        run(nm, p, sl, sv, root=root, out=out)
