import json, sys, os, csv, re
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mt
FEEDERS = os.path.abspath(os.environ.get("G4W_FEEDERS", "feeders"))  # pasta IEEETestCases do electricdss-tst
INK="#14161a"; MUTED="#6b7078"; GRID="#d9d6cf"; OLIVE="#4d6e00"; LIME="#9ccf1f"; BLUE="#2f6fb0"; GRAY="#9aa0a8"
L1,L2,L3="#e8c98a","#c98a3e","#7a4312"
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]; DST=sys.argv[2]
F=["ieee13","ieee34","ieee123"]; LBL={"ieee13":"IEEE 13","ieee34":"IEEE 34","ieee123":"IEEE 123"}
COORD={"ieee13":FEEDERS + "/13Bus/IEEE13Node_BusXY.csv","ieee34":FEEDERS + "/34Bus/IEEE34_BusXY.csv","ieee123":FEEDERS + "/123Bus/BusCoords.dat"}
ROOT={"ieee13":"SOURCEBUS","ieee34":"SOURCEBUS","ieee123":"150"}
R={f:json.load(open(f"{OUT}/results_{f}.json")) for f in F}
def coords(p):
    xy={}
    for line in open(p):
        parts=[x for x in re.split(r"[,\s]+",line.strip()) if x]
        if len(parts)>=3:
            try: xy[parts[0].upper()]=(float(parts[1]),float(parts[2]))
            except ValueError: pass
    return xy
def get(xy,b):
    if b in xy: return xy[b]
    if b.endswith("R") and b[:-1] in xy: return xy[b[:-1]]
    return None
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")
# ---------- Fig. 1: redes (traçado ortogonal, marcadores sem sobreposição)
def _route(pa,pb,segs):
    """Trecho em L quando as coordenadas são diagonais; escolhe o canto que menos se sobrepõe ao traçado existente."""
    if abs(pa[0]-pb[0])<1e-6 or abs(pa[1]-pb[1])<1e-6: return [pa,pb]
    def ov(path):
        n=0
        for (p,q) in zip(path[:-1],path[1:]):
            for (u,v) in segs:
                if abs(p[1]-q[1])<1e-6 and abs(u[1]-v[1])<1e-6 and abs(p[1]-u[1])<1e-6:
                    lo,hi=max(min(p[0],q[0]),min(u[0],v[0])),min(max(p[0],q[0]),max(u[0],v[0])); n+=max(0,hi-lo)
                if abs(p[0]-q[0])<1e-6 and abs(u[0]-v[0])<1e-6 and abs(p[0]-u[0])<1e-6:
                    lo,hi=max(min(p[1],q[1]),min(u[1],v[1])),min(max(p[1],q[1]),max(u[1],v[1])); n+=max(0,hi-lo)
        return n
    c1=[pa,(pb[0],pa[1]),pb]; c2=[pa,(pa[0],pb[1]),pb]
    return c1 if ov(c1)<=ov(c2) else c2
for f in F:
    r=R[f]; xy=coords(COORD[f])
    if f=="ieee123":
        # 151 e 300 são só as pontas da chave de interligação (aberta, sem carga); no arquivo de coordenadas
        # ficam longe dos vizinhos e cruzam o desenho, então são desenhadas logo ao lado de 51 e 108
        x,y=xy["51"]; xy["151"]=(x+200,y)
        x,y=xy["108"]; xy["300"]=(x,y+200)
    w,h={"ieee13":(6.4,4.6),"ieee34":(8.2,4.2),"ieee123":(8.2,5.6)}[f]
    fig,ax=plt.subplots(figsize=(w,h))
    E=[(a,b,get(xy,a),get(xy,b)) for a,b in r["edges"]]; E=[e for e in E if e[2] and e[3]]
    segs=[(pa,pb) for a,b,pa,pb in E if abs(pa[0]-pb[0])<1e-6 or abs(pa[1]-pb[1])<1e-6]
    paths={}
    for a,b,pa,pb in E:
        path=_route(pa,pb,segs); paths[(a,b)]=path
        if len(path)==3: segs+= [(path[0],path[1]),(path[1],path[2])]
        ax.plot([q[0] for q in path],[q[1] for q in path],color=GRAY,lw=1.1,zorder=1,solid_capstyle="round")
    kw={}; ph={}
    for L in r["load_info"]:
        kw[L["bus"]]=kw.get(L["bus"],0)+L["kw"]; ph.setdefault(L["bus"],set()).update(L["phases"])
    mx=max(kw.values()); span=max(np.ptp([p[0] for p in xy.values()]),np.ptp([p[1] for p in xy.values()]))
    smax={"ieee13":160,"ieee34":110,"ieee123":55}[f]; smin={"ieee13":22,"ieee34":16,"ieee123":9}[f]
    loadpts=[]
    for b,v in kw.items():
        p=get(xy,b)
        if p:
            c={1:L1,2:L2,3:L3}[min(3,len(ph[b]))]; ax.scatter(*p,s=smin+smax*v/mx,color=c,edgecolor="white",lw=0.6,zorder=3); loadpts.append(p)
    seen=set()
    for a,b in r["regs"]:
        p=get(xy,a)
        if p and p not in seen: seen.add(p); ax.scatter(*p,marker="D",s=40,color=BLUE,edgecolor="white",lw=0.6,zorder=5)
    rt=get(xy,ROOT[f]); ax.scatter(*rt,marker="s",s=70,color=INK,zorder=6)
    LP=np.array(loadpts+[rt]+[get(xy,a) for a,b in r["regs"] if get(xy,a)])
    for k,s in enumerate(r["modes"]["VI"]["greedy_order"][:4]):
        a,b=r["lines_buses"][s]; path=paths.get((a,b)) or paths.get((b,a))
        if not path: continue
        # posição ao longo do trecho (só no primeiro segmento, perto da barra de origem) mais distante das cargas
        best=None
        for t in np.linspace(0.3,0.75,10):
            P=np.array(path[0])*(1-t)+np.array(path[1])*t; d=np.min(np.hypot(*(LP-P).T))
            if best is None or d>best[0]: best=(d,P)
        x,y=best[1]
        ax.scatter(x,y,s=150,color=LIME,edgecolor=OLIVE,lw=1.5,zorder=7); ax.annotate(str(k+1),(x,y),ha="center",va="center",fontsize=7.5,weight="bold",color=INK,zorder=8)
    if f=="ieee13":
        for b,p in xy.items():
            if b!="SOURCEBUS": ax.annotate(b,p,xytext=(4,4),textcoords="offset points",fontsize=6.5,color=MUTED)
    ax.scatter([],[],marker="s",s=60,color=INK,label="Subestação"); ax.scatter([],[],marker="D",s=40,color=BLUE,label="Regulador de tensão")
    ax.scatter([],[],s=120,color=LIME,edgecolor=OLIVE,lw=1.5,label="Ponto de medição (ordem da alocação, K = 4)")
    for c,l in [(L1,"Carga monofásica"),(L2,"Carga bifásica"),(L3,"Carga trifásica")]: ax.scatter([],[],s=50,color=c,edgecolor="white",lw=0.6,label=l)
    ax.scatter([],[],s=0,label="(área do círculo ∝ potência)")
    ax.set_axis_off(); ax.set_aspect("equal"); ax.legend(loc="upper left",bbox_to_anchor=(1.01,1.0),frameon=False,fontsize=7.5)
    save(fig,f"fig1_rede_{f}"); plt.close(fig)
# ---------- Fig. 2: erro x K
fig,axs=plt.subplots(3,3,figsize=(10,8.2))
tit=["Tensão — erro máx. de |V| (%)","Corrente nos trechos — erro (%)","Carga por barra — erro de P (%)"]
for i,f in enumerate(F):
    g=R[f]["modes"]["VI"]; k=np.arange(len(g["greedy"]["v"]))
    for j,key in enumerate(["v","i","p"]):
        ax=axs[i,j]; rnd=np.array(g["random"][key])
        ax.fill_between(k,rnd[:,1],rnd[:,2],color=GRAY,alpha=0.22,lw=0)
        ax.plot(k,rnd[:,0],color=GRAY,ls="--",marker="o",ms=2.5,lw=1.2,label="Aleatória (média e faixa 10–90 % de 30 sorteios)")
        ax.plot(k,g["trunk"][key],color=BLUE,marker="s",ms=3,lw=1.4,label="Uniforme no tronco")
        ax.plot(k,g["greedy"][key],color=OLIVE,marker="o",ms=3.5,lw=2.2,label="Otimizada (gulosa)")
        ax.set_ylim(bottom=0)
        if i==0: ax.set_title(tit[j],fontsize=9,loc="left",color=INK)
        if i==2: ax.set_xlabel("Pontos de medição, K")
        if j==0: ax.set_ylabel(LBL[f],fontsize=10,color=INK,weight="bold")
axs[0,0].legend(frameon=False,fontsize=7.5,loc="upper right")
fig.tight_layout(); save(fig,"fig2_erro_k"); plt.close(fig)
# ---------- Fig. 3: precisão da tensão do sensor
fig,axs=plt.subplots(1,2,figsize=(9,3.3))
cols={"ieee13":OLIVE,"ieee34":BLUE,"ieee123":"#a0522d"}
for f in F:
    s=R[f]["sigv"]; oi=s["only_i"]
    axs[0].plot(s["pct"],np.array(s["v"])/oi[0]*100,color=cols[f],marker="o",ms=4,label=LBL[f])
    axs[1].plot(s["pct"],np.array(s["p"])/oi[1]*100,color=cols[f],marker="o",ms=4,label=LBL[f])
for ax,t in zip(axs,["Erro máx. de |V| (% do caso só corrente)","Erro de P por barra (% do caso só corrente)"]):
    ax.axhline(100,color=GRAY,ls="--",lw=1); ax.set_xscale("log"); ax.set_xticks([0.1,0.2,0.5,1,2]); ax.xaxis.set_major_formatter(mt.FuncFormatter(lambda x,_: f"{x:g}".replace(".",",")))
    ax.set_xlabel("Erro da medição de tensão do sensor (%, desvio-padrão)"); ax.set_title(t,fontsize=9,loc="left")
axs[0].text(2,101,"só corrente",ha="right",va="bottom",fontsize=8,color=MUTED); axs[0].legend(frameon=False,fontsize=8)
fig.tight_layout(); save(fig,"fig3_tensao_sensor"); plt.close(fig)
# ---------- Fig. 4: sincronismo (fração do ganho mantida)
fig,ax=plt.subplots(figsize=(6.4,3.4))
for f in F:
    s=R[f]["sync"]; e0=R[f]["modes"]["VI"]["greedy"]["v"][0]; eb=s["v"][0]
    frac=(e0-np.array(s["v"]))/(e0-eb)*100
    ax.plot(s["us"],frac,color=cols[f],marker="o",ms=4,label=LBL[f])
ax.axvspan(0.05,1.5,color=LIME,alpha=0.15,lw=0); ax.text(0.07,52,"faixa típica\nGNSS (≤ 1 µs)",color=OLIVE,fontsize=8)
ax.set_xscale("log"); ax.set_xticks([0.1,1,10,100,1000]); ax.xaxis.set_major_formatter(mt.FuncFormatter(lambda x,_: f"{x:g}".replace(".",",")))
ax.set_ylim(0,105); ax.set_xlabel("Erro de sincronismo entre sensores (µs, desvio-padrão)"); ax.set_ylabel("Ganho dos 4 sensores mantido (%)")
ax.legend(frameon=False,fontsize=8,loc="lower left"); fig.tight_layout(); save(fig,"fig4_sincronismo"); plt.close(fig)
# ---------- Fig. 5: carga estimada x verdadeira (IEEE 34)
L=R["ieee34"]["loads_scatter"]; pt=np.array(L["p_true"]); p0=np.array(L["p_est0"]); p4=np.array(L["p_est4"])
fig,axs=plt.subplots(1,2,figsize=(8,3.6),sharex=True,sharey=True)
lim=[min(pt.min(),p0.min(),p4.min())*1.05, pt.max()*1.1]
for ax,pe,t,c in [(axs[0],p0,"Sem sensores (só pseudomedidas e subestação)",GRAY),(axs[1],p4,"Com 4 sensores",OLIVE)]:
    ax.scatter(pt.ravel(),pe.ravel(),s=5,alpha=0.3,color=c,lw=0,rasterized=True); ax.plot(lim,lim,color=INK,lw=0.8)
    ax.set_title(t,fontsize=9,loc="left"); ax.set_xlabel("P verdadeira por barra (kW)")
axs[0].set_ylabel("P estimada por barra (kW)"); fig.tight_layout(); save(fig,"fig5_cargas_ieee34"); plt.close(fig)
print("ok")
# ---------- Fig. extra: subestação como fasor x SCADA (alocação otimizada)
if all("scada" in R[f] for f in F):
    fig,axs=plt.subplots(1,3,figsize=(10,3.0))
    for ax,f in zip(axs,F):
        g=R[f]["modes"]["VI"]["greedy"]; s=R[f]["scada"]; k=np.arange(len(g["v"]))
        ax.plot(k,g["v"],color=OLIVE,marker="o",ms=3.5,lw=2,label="Saída medida como fasor de corrente")
        ax.plot(k,s["v"],color=BLUE,marker="s",ms=3.5,lw=1.6,ls="--",label="Saída medida como no SCADA (|I|, P, Q)")
        ax.set_title(LBL[f],fontsize=9.5,loc="left",color=INK); ax.set_xlabel("Pontos de medição, K"); ax.set_ylim(bottom=0)
    axs[0].set_ylabel("Erro máx. de |V| (%)"); axs[0].legend(frameon=False,fontsize=7.5)
    fig.tight_layout(); save(fig,"fig6_subestacao"); plt.close(fig)
print("ok2")