"""Análise do estudo v2: (1) diagnóstico MT x BT antes do conserto; (2) o conserto resolveu? (3) economia."""
import json, sys, glob, collections
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from casos import TYPES, MT_TYPES, LABEL
INK="#14161a"; MUTED="#6b7078"; GRID="#d9d6cf"; OLIVE="#4d6e00"; BLUE="#2f6fb0"; RED="#b5452e"
plt.rcParams.update({"font.family":"DejaVu Sans","font.size":9.5,"axes.edgecolor":MUTED,"axes.labelcolor":INK,"xtick.color":MUTED,"ytick.color":MUTED,
 "axes.grid":True,"grid.color":GRID,"grid.linewidth":0.6,"axes.spines.top":False,"axes.spines.right":False,"svg.fonttype":"none",
 "figure.facecolor":"none","axes.facecolor":"none"})
OUT = sys.argv[1]
_all = [json.loads(l) for f in sorted(glob.glob(OUT + "/casos_*.jsonl")) for l in open(f)]
_seen = set(); recs = []
for r in _all:                      # remove duplicados (mesmo caso rodado por dois processos)
    k = r.get("id") or (r["tipo"], json.dumps(r.get("info"), sort_keys=True))
    if r["status"] == "ok" and k in _seen: continue
    _seen.add(k); recs.append(r)
V = lambda d: d[0] > 3 or d[1] > 0.5
Ks = ("0", "2", "4")
COL = {"0": RED, "2": BLUE, "4": OLIVE}; LAB = {"0": "Sem sensores", "2": "2 sensores na MT", "4": "4 sensores na MT"}

# ---------------- economia (cenários; premissas documentadas)
EUSD_BT = 75.0
MT_KW = {"MT830": 45, "MT840": 27, "MT844": 405, "MT848": 60, "MT860": 60, "MT890": 250}
eusd_mt = lambda kw: kw * 20.82 + kw * 0.5 * 730 / 1000 * 258.0
def comp_month(d, nivel):
    eu = EUSD_BT if nivel == "bt" else None
    return ((d[0] - 3) / 100 * 3 if d[0] > 3 else 0) + ((d[1] - 0.5) / 100 * (7 if nivel == "bt" else 5) if d[1] > 0.5 else 0)
def cmonth(c):
    eu = EUSD_BT if c["nivel"] == "bt" else eusd_mt(MT_KW[c["id"]])
    return comp_month(c["antes"], c["nivel"]) * eu
D, D0, CM, CYCLE, CREW = 4, 1, 700.0, 3, 3000.0     # meses até a comprobatória; logo; custo de uma medição; meses de um ciclo a mais; custo de uma intervenção

res = {"tipos": {}, "geral": {}}
rows = []
for t in TYPES:
    R = [r for r in recs if r["tipo"] == t]
    ok = [r for r in R if r["status"] == "ok"]
    st = collections.Counter(r["status"] for r in R)
    if not ok: res["tipos"][t] = dict(status=dict(st)); continue
    rows.append(t)
    cause_mt = t in MT_TYPES
    d = dict(status=dict(st), casos=len(ok), consumidores=sum(len(r["comp"]) for r in ok), diag={}, depois={}, econ={})
    for K in Ks:
        # diagnóstico por caso: maioria dos compensados explicada pela tensão estimada (=> causa na MT)
        dec = []
        for r in ok:
            frac = np.mean([V(c["est_antes"][K]) for c in r["comp"]])
            dec.append((frac >= 0.5) == cause_mt)
        d["diag"][K] = float(np.mean(dec))
        cons = [(r, c) for r in ok for c in r["comp"]]
        tr = np.array([not V(c["depois"]) for _, c in cons]); pr = np.array([not V(c["est"][K]) for _, c in cons])
        d["depois"][K] = dict(acerto=float((pr == tr).mean()), reprovada=float((pr & ~tr).mean()), atraso=float((~pr & tr).mean()))
        # economia por caso (R$): (a) comprobatória antecipada nos que de fato regularizaram e o estimador viu;
        # (b) medição poupada nos que não regularizaram e o estimador viu; (c) diagnóstico — ver abaixo
        ganho = []
        for r in ok:
            g = 0.0
            for c in r["comp"]:
                p_, t_ = (not V(c["est"][K])), (not V(c["depois"]))
                if p_ and t_: g += (D - D0) * cmonth(c)
                if (not p_) and (not t_): g += CM
            ganho.append(g)
        d["econ"][K] = dict(pos_conserto=float(np.mean(ganho)),
                            custo_ciclo_extra=float(np.mean([sum(cmonth(c) for c in r["comp"]) * CYCLE + CREW for r in ok])))
    res["tipos"][t] = d
    print(f'{LABEL[t]:38s} casos={len(ok):2d} | diag: ' + " ".join(f'K{K}={d["diag"][K]:.2f}' for K in Ks) +
          ' | reprovada: ' + " ".join(f'K{K}={d["depois"][K]["reprovada"]:.2f}' for K in Ks) +
          ' | atraso: ' + " ".join(f'K{K}={d["depois"][K]["atraso"]:.2f}' for K in Ks))
for grp, types in (("MT", MT_TYPES), ("BT", [t for t in TYPES if t not in MT_TYPES])):
    tt = [t for t in types if t in rows]
    if not tt: continue
    res["geral"][grp] = {K: dict(diag=float(np.mean([res["tipos"][t]["diag"][K] for t in tt])),
                                 reprovada=float(np.mean([res["tipos"][t]["depois"][K]["reprovada"] for t in tt])),
                                 atraso=float(np.mean([res["tipos"][t]["depois"][K]["atraso"] for t in tt]))) for K in Ks}
    print(grp, json.dumps(res["geral"][grp]))
json.dump(res, open(OUT + "/resumo_v2.json", "w"), indent=1)

if rows:
    y = np.arange(len(rows))[::-1]; h = 0.26
    ylab = [("MT · " if t in MT_TYPES else "BT · ") + LABEL[t] for t in rows]
    fig, axs = plt.subplots(1, 2, figsize=(11.5, 0.42 * len(rows) + 1.8), sharey=True)
    for q, K in enumerate(Ks):
        axs[0].barh(y + (1 - q) * h, [res["tipos"][t]["diag"][K] * 100 for t in rows], h, color=COL[K], label=LAB[K])
        axs[1].barh(y + (1 - q) * h, [res["tipos"][t]["depois"][K]["reprovada"] * 100 for t in rows], h, color=COL[K])
    axs[0].set_title("Antes do conserto: acerta se a causa está\nna média ou na baixa tensão (% dos casos)", loc="left", fontsize=9.5, color=INK)
    axs[1].set_title("Depois do conserto: diz que regularizou, mas não\n(% dos consumidores; comprobatória reprovada)", loc="left", fontsize=9.5, color=INK)
    for ax in axs: ax.set_xlim(0, 100)
    axs[0].set_yticks(y); axs[0].set_yticklabels(ylab, fontsize=8.5)
    hh, ll = axs[0].get_legend_handles_labels(); fig.tight_layout(rect=(0, 0.05, 1, 1))
    fig.legend(hh, ll, loc="lower center", ncol=3, frameon=False, fontsize=8.5)
    fig.savefig(OUT + "/p3_fig_v2_diag_depois.svg", bbox_inches="tight"); fig.savefig(OUT + "/p3_fig_v2_diag_depois.png", dpi=150, bbox_inches="tight", facecolor="white")