import json, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
INK="#14161a"; MUTED="#6b7078"; GRID="#d9d6cf"; OLIVE="#4d6e00"; LIME="#9ccf1f"; BLUE="#2f6fb0"; GRAY="#9aa0a8"; 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, DST = sys.argv[1], sys.argv[2]
F=["ieee13","ieee34","ieee123"]; LBL={"ieee13":"IEEE 13","ieee34":"IEEE 34","ieee123":"IEEE 123"}
D={f:json.load(open(f"{OUT}/p2_{f}.json")) for f in F}
MODES=[("ingenuo","Dados sem tratamento",RED),("auto","Sem tratamento + limpeza automática",BLUE),("tratado","Dados tratados",OLIVE)]
KS=[0,2,4]
def save(fig,name):
    fig.savefig(f"{DST}/{name}.svg",bbox_inches="tight",dpi=200); fig.savefig(f"{OUT}/{name}.png",dpi=140,bbox_inches="tight",facecolor="white")
def raw(f,key):
    r=D[f]["raw"][key]; st=np.array(r["status"]); ok=st=="ok"
    ev=np.full(len(st),np.nan); ev[ok]=r["err_v"]; rej=np.zeros(len(st),bool); rej[ok]=r["chi_rej"]
    ep=np.full(len(st),np.nan); ep[ok]=r["err_p"]
    return st,ok,ev,rej,ep
# Fig 1: estimativas aceitas
fig,axs=plt.subplots(1,3,figsize=(10.5,3.4),sharey=True)
w=0.26
for ax,f in zip(axs,F):
    for j,(m,l,c) in enumerate(MODES):
        vals=[]; divs=[]
        for K in KS:
            st,ok,ev,rej,ep=raw(f,f"{m}_K{K}"); vals.append(np.mean(ok&~rej)*100); divs.append(np.mean(~ok)*100)
        x=np.arange(3)+(j-1)*w
        ax.bar(x,vals,w,color=c,label=l)
        ax.bar(x,divs,w,bottom=vals,color="none",edgecolor=c,hatch="////",lw=0.6,label="Não convergiu" if (f=="ieee13" and j==1) else None)
    ax.set_xticks(range(3)); ax.set_xticklabels([f"K = {k}" for k in KS]); ax.set_title(LBL[f],fontsize=10,loc="left",color=INK); ax.set_ylim(0,105)
axs[0].set_ylabel("Estimativas aceitas pelo teste\nde consistência (%)")
axs[0].legend(frameon=False,fontsize=7.5,loc="lower left",bbox_to_anchor=(0,-0.42),ncol=2)
save(fig,"p2_fig2_aceitas"); plt.close(fig)
# Fig 2: erro de tensão (mediana e p90) e de carga
fig,axs=plt.subplots(2,3,figsize=(10.5,6.0))
for i,(metric,lab) in enumerate([("v","Erro máx. de |V| (%)"),("p","Erro de P por barra (%)")]):
    for jf,f in enumerate(F):
        ax=axs[i,jf]
        for j,(m,l,c) in enumerate(MODES):
            med=[];lo=[];hi=[]
            for K in KS:
                st,ok,ev,rej,ep=raw(f,f"{m}_K{K}"); a=(ev if metric=="v" else ep)[ok]
                med.append(np.median(a)); lo.append(np.percentile(a,10)); hi.append(np.percentile(a,90))
            x=np.arange(3)+(j-1)*0.22
            ax.errorbar(x,med,yerr=[np.array(med)-lo,np.array(hi)-med],fmt="o",color=c,ms=5,capsize=3,lw=1.3,label=l)
        ax.set_xticks(range(3)); ax.set_xticklabels([f"K = {k}" for k in KS]); ax.set_ylim(bottom=0)
        if i==0: ax.set_title(LBL[f],fontsize=10,loc="left",color=INK)
        if jf==0: ax.set_ylabel(lab)
h,l=axs[0,0].get_legend_handles_labels(); fig.tight_layout(rect=(0,0.05,1,1)); fig.legend(h,l,loc="lower center",ncol=3,frameon=False,fontsize=8)
save(fig,"p2_fig3_erros"); plt.close(fig)
# Fig 3: distribuição do erro de tensão, IEEE 34
fig,ax=plt.subplots(figsize=(6.4,3.5))
for key,l,c,ls in [("ingenuo_K0","Sem tratamento, sem sensores",RED,"-"),("tratado_K0","Tratados, sem sensores",OLIVE,"--"),
                   ("auto_K4","Sem tratamento + limpeza + 4 sensores",BLUE,"-"),("tratado_K4","Tratados + 4 sensores",OLIVE,"-")]:
    st,ok,ev,rej,ep=raw("ieee34",key); a=np.sort(ev[ok]); ax.plot(a,np.arange(1,len(a)+1)/len(a)*100,color=c,ls=ls,lw=2,label=l)
ax.set_xscale("log"); ax.set_xlabel("Erro máx. de |V| no cenário (%)"); ax.set_ylabel("Cenários com erro menor (%)")
ax.legend(frameon=False,fontsize=8,loc="upper left",bbox_to_anchor=(1.01,1.0)); fig.tight_layout(); save(fig,"p2_fig4_distribuicao"); plt.close(fig)
print("ok")