"""
Estudo estatístico do Paper 3: muitos casos sorteados por tipo de problema e conserto.
Uso: python estudo.py <saida> <tipo|all> <n_casos_por_tipo> <seed0>
"""
import os, sys, json, time, math, copy
import numpy as np
import opendssdirect as dss
from rede import (build_spec, calibrate_taps, Circuit, consumers, consumer_node, base_loads, week_profiles, drp_drc,
                  mt_reading, MT_CUSTOMERS, STEPS, FEEDER)
from casos import make_case, apply_wrong_fix, background_errors, TYPES, MT_TYPES
from estima import estimate_week, estimate_week_v2

G_SD = 0.20            # erro do cadastro de carga por consumidor (lognormal)
PSEUDO_REL = 0.30      # incerteza atribuída às cargas previstas


def line_lengths():
    dss.Text.Command("clear"); os.chdir(os.path.dirname(FEEDER)); dss.Text.Command(f"compile [{FEEDER}]")
    out = {}
    for nm in dss.Lines.AllNames():
        dss.Lines.Name(nm)
        out[nm] = (dss.Lines.Bus1().split(".")[0].upper(), dss.Lines.Bus2().split(".")[0].upper(), dss.Lines.Length())
    return out


def run_week(spec, gvec, cur, irr):
    """Semana verdadeira. Retorna V (STEPS x n), nós, taps por passo."""
    ck = Circuit(spec)
    C, kwc, kvc, kwm, kvm, cls = base_loads(spec); nC = len(C)
    pvc = np.array([c.get("pv_kw", 0.0) for _, _, c in C])
    V = np.empty((STEPS, len(ck.nodes)), complex); taps = []
    for k in range(STEPS):
        f = cur[k] * gvec
        pv_mt = []
        for pv in spec["pv_mt"]:
            p = pv["kw"] * irr[k]; q = -p * math.tan(math.acos(min(1.0, abs(pv["pf"])))) if pv["pf"] < 0 else 0.0
            pv_mt.append((p, q))
        ck.set_loads(kwc * f[:nC], kvc * f[:nC], kwm * f[nC:], kvm * f[nC:], pv_c=pvc * irr[k], pv_mt=pv_mt)
        v = ck.solve()
        if v is None: return None
        V[k] = v; taps.append(ck.taps())
    return V, ck.nodes, ck.idx, taps


def readings(spec, V, idx):
    C = consumers(spec)
    bt = np.array([np.abs(V[:, idx[consumer_node(c, t)]]) for _, t, c in C]).T
    mt = {b: mt_reading(V, idx, b) for b in MT_CUSTOMERS}
    return bt, mt


def indicators(bt, mt):
    out = [drp_drc(bt[:, j], "bt") for j in range(bt.shape[1])]
    out += [drp_drc(mt[b], "mt") for b in MT_CUSTOMERS]
    return np.array(out)


def viol(d):
    return (d[:, 0] > 3.0) | (d[:, 1] > 0.5)


def run_case(ctype, spec0, rng, line_len, case_id, Ks=(0, 2, 4)):
    """v2: rede verdadeira com erros de cadastro (impedância, fase); taps dos reguladores desconhecidos no estimador;
    estimação também na semana ANTES (diagnóstico MT x BT) e DEPOIS (o conserto resolveu?)."""
    specT0, moved = background_errors(spec0, rng, line_len)
    case = make_case(ctype, specT0, rng, line_len, spec_cad0=spec0)
    C = consumers(spec0); nC = len(C); ids = [c["name"] for _, _, c in C] + [f"MT{b}" for b in MT_CUSTOMERS]
    cls = [c["cls"] for _, _, c in C] + [m["cls"] for m in spec0["mt"]]
    g = np.exp(G_SD * rng.standard_normal(len(cls)))
    s1, s2 = int(rng.integers(1e9)), int(rng.integers(1e9))
    cur1, typ1, irr1, irr_typ1 = week_profiles(cls, np.random.default_rng(s1), start_dow=int(rng.integers(7)))
    r = run_week(case["before"], g, cur1, irr1)
    if r is None: return dict(status="nao_convergiu", tipo=ctype)
    V1, nodes1, idx1, taps1 = r
    d1 = indicators(*readings(case["before"], V1, idx1))
    # semana de referência: mesma carga e clima, rede verdadeira sem o problema -> separa violações causadas pelo problema
    r0 = run_week(specT0, g, cur1, irr1)
    if r0 is None: return dict(status="nao_convergiu", tipo=ctype)
    d0 = indicators(*readings(specT0, r0[0], r0[2]))
    comp = np.where(viol(d1) & ~viol(d0))[0]
    if len(comp) == 0:
        return dict(status="sem_violacao", tipo=ctype, info=case["info"])
    if case["conserto"] == "errado" and ctype in MT_TYPES:
        tnames = {C[j][1]["name"] for j in comp if j < nC}
        apply_wrong_fix(case, tnames)
    cur2, typ2, irr2, irr_typ2 = week_profiles(cls, np.random.default_rng(s2), start_dow=int(rng.integers(7)))
    r = run_week(case["after"], g, cur2, irr2)
    if r is None: return dict(status="nao_convergiu", tipo=ctype)
    V2, nodes2, idx2, taps2 = r
    d2 = indicators(*readings(case["after"], V2, idx2))
    r0b = run_week(specT0, g, cur2, irr2)
    d0b = indicators(*readings(specT0, r0b[0], r0b[2])) if r0b is not None else np.zeros_like(d2)
    t0 = time.time()
    # ---- diagnóstico (semana antes): cadastro antes do conserto (usina FV conhecida com fp 1)
    cad1 = copy.deepcopy(spec0); cad1["pv_mt"] = [dict(p, pf=1.0) for p in case["cad"]["pv_mt"]]
    _, kwc, kvc, kwm, kvm, _ = base_loads(cad1)
    ps1 = dict(kwc=typ1[:, :nC] * kwc, kvc=typ1[:, :nC] * kvc, kwm=typ1[:, nC:] * kwm, kvm=typ1[:, nC:] * kvm)
    e1 = estimate_week_v2(cad1, V1, nodes1, ps1, [p["kw"] * irr_typ1 for p in cad1["pv_mt"]], Ks=Ks, rel=PSEUDO_REL,
                          seed=s1 % 100000, stride=2)
    da = {K: indicators(e1[K]["bt"], e1[K]["mt"]) for K in Ks}
    # ---- depois do conserto
    _, kwc, kvc, kwm, kvm, _ = base_loads(case["cad"])
    ps2 = dict(kwc=typ2[:, :nC] * kwc, kvc=typ2[:, :nC] * kvc, kwm=typ2[:, nC:] * kwm, kvm=typ2[:, nC:] * kvm)
    e2 = estimate_week_v2(case["cad"], V2, nodes2, ps2, [p["kw"] * irr_typ2 for p in case["cad"]["pv_mt"]], Ks=Ks,
                          rel=PSEUDO_REL, seed=s2 % 100000, stride=2)
    de = {K: indicators(e2[K]["bt"], e2[K]["mt"]) for K in Ks}
    T2 = np.array([t[:6] for t in taps2])
    tap_err = {K: float(np.nanmean(np.abs(e2[K]["tap"][:, 3:6] - T2[:, 3:6]))) for K in Ks}
    rec = dict(status="ok", versao=2, id=case_id, tipo=ctype, conserto=case["conserto"], info=case["info"],
               fase_errada=len(moved), n_comp=int(len(comp)), t_est=round(time.time() - t0, 1),
               fail={K: [e1[K]["fail"], e2[K]["fail"]] for K in Ks}, tap_err_reg2=tap_err,
               comp=[dict(id=ids[j], nivel="bt" if j < nC else "mt", antes=d1[j].round(2).tolist(), depois=d2[j].round(2).tolist(),
                          est_antes={K: da[K][j].round(2).tolist() for K in Ks},
                          est={K: de[K][j].round(2).tolist() for K in Ks}) for j in comp],
               novas=[dict(id=ids[j], nivel="bt" if j < nC else "mt", depois=d2[j].round(2).tolist(),
                           est={K: de[K][j].round(2).tolist() for K in Ks}) for j in np.where(viol(d2) & ~viol(d1) & ~viol(d0b))[0]],
               fundo_antes=int(viol(d0).sum()), fundo_depois=int(viol(d0b).sum()),
               falsas={K: int(np.sum(viol(de[K]) & ~viol(d2) & ~viol(d1))) for K in Ks})
    return rec


def main():
    out, which, n, seed0 = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4])
    os.makedirs(out, exist_ok=True); out = os.path.abspath(out)
    spec0 = calibrate_taps(build_spec(1))
    json.dump(spec0, open(os.path.join(out, "spec_base.json"), "w"))
    line_len = line_lengths()
    types = TYPES if which == "all" else which.split(",")
    fn = os.path.join(out, f"casos_{which.replace(',', '+')}_{seed0}.jsonl")
    for ctype in types:
        rng = np.random.default_rng(seed0 * 1000 + TYPES.index(ctype))
        done = tries = 0
        while done < n and tries < 6 * n:
            tries += 1
            rec = run_case(ctype, spec0, rng, line_len, f"{ctype}_{seed0}_{tries}")
            with open(fn, "a") as f: f.write(json.dumps(rec, default=float) + "\n")
            if rec["status"] == "ok": done += 1
            print(ctype, tries, rec["status"], rec.get("n_comp"), rec.get("conserto"), rec.get("t_est"), flush=True)


if __name__ == "__main__":
    main()