import sys,math
import numpy as np
import pandas as pd
import pycountry
D=sys.argv[1] if len(sys.argv)>1 else "."
PROFIL=sys.argv[2] if len(sys.argv)>2 else "soir"
HYDRO_DEPLACABLE=float(sys.argv[3]) if len(sys.argv)>3 else 0.3
REND,V15_TWH,FIXE_TWH=0.78,405.0,986.0
DEM={"soir":np.array([0.85]*6+[1.0]*11+[1.25]*6+[0.95]),"plat":np.ones(24),"souple":np.array([0.85]*6+[1.0]*4+[1.5]*6+[1.0]*1+[1.05]*6+[0.95])}[PROFIL]
DEM=DEM/DEM.sum()
SOL=np.array([max(0.0,math.sin(math.pi*(h+0.5-6)/12)) if 6<=h<18 else 0.0 for h in range(24)])
SOL=SOL/SOL.sum()
PLAT=np.full(24,1/24)
SOIR=np.array([1.0 if (h>=17 or h<5) else 0.0 for h in range(24)])
SOIR=SOIR/SOIR.sum()
def i3(c):
    p=pycountry.countries.get(alpha_2=c)
    return p.alpha_3 if p else c
s=pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_sites.csv",keep_default_na=False)
s["iso"]=[i3(c) for c in s.pays]
s["t"]=s.twh*s.pris
T=s.pivot_table(index="iso",columns="type",values="t",aggfunc="sum",fill_value=0.0)
p=pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_pays.csv",keep_default_na=False).set_index("iso")
r=pd.read_csv(f"{D}/remparts_v15.csv")
r["iso"]=[i3(c) for c in r.pays]
v15_tout=r.groupby("iso").stockage_mwh.sum()/1000
v15_plan=r[r.plan].groupby("iso").stockage_mwh.sum()/1000
h=pd.read_csv(f"{D}/owid-energy-data.csv",usecols=["iso_code","year","hydro_electricity"])
h=h[h.iso_code.notna()].dropna(subset=["hydro_electricity"]).sort_values("year").groupby("iso_code").tail(1).set_index("iso_code").hydro_electricity
prof={}
for iso,x in p.iterrows():
    g=lambda k:float(T.loc[iso,k]) if iso in T.index and k in T.columns else 0.0
    sol=g("soleil")+g("poste")
    mix=g("mixte")
    cre=g("crete")
    v21=g("v21")
    hyb=float(x.panneaux_hybrides_twh)+g("flottant")
    fix=float(x.lignes_fixes_twh)
    j=lambda twh:twh*1000/365
    o=j(sol*0.82+mix*0.8*0.43+hyb*0.8)*SOL+j(cre*0.81+mix*0.8*0.57+hyb*0.2+v21+fix)*PLAT+j(sol*0.18+cre*0.19+mix*0.2)*SOIR
    tot=o.sum()
    ex=j(float(x.exporte_twh))
    k=max(0.0,1-ex/tot) if tot>0 else 0.0
    prof[iso]=dict(x=x,o=o*k,exp=o*(1-k),sol=sol+mix*0.43+hyb*0.8,vent=cre+mix*0.57+hyb*0.2+v21,v14=j(sol*0.18+cre*0.19+mix*0.2))
EXP=sum(v["exp"] for v in prof.values())
EXP=EXP/EXP.sum() if EXP.sum()>0 else PLAT
rows=[]
for iso,v in prof.items():
    x=v["x"]
    if float(x.besoin_plan_twh)<=0: continue
    j=lambda twh:twh*1000/365
    offre=v["o"]+j(float(x.importe_twh))*EXP
    besoin=j(float(x.besoin_plan_twh))*DEM
    rr=besoin-offre
    manque=float(np.clip(rr,0,None).sum())
    surplus=float(np.clip(-rr,0,None).sum())
    stock=min(manque,surplus*REND)
    hyd=float(h.get(iso,0.0))*1000/365*HYDRO_DEPLACABLE
    hy=min(hyd,stock)
    rows.append(dict(iso=iso,pays=x.pays,besoin_twh=round(float(x.besoin_plan_twh),1),importe_twh=round(float(x.importe_twh),1),hydraulique_existante_deplacee_gwh_jour=round(hy,1),solaire_twh=round(v["sol"],1),vent_twh=round(v["vent"],1),stockage_v14_gwh_jour=round(v["v14"],1),stockage_a_ajouter_gwh_jour=round(stock-hy,1),pointe_du_soir_gw=round(float(np.clip(rr,0,None).max()),1),manque_non_stockable_gwh_jour=round(max(0.0,manque-surplus*REND),1),v15_possibles_gwh=round(float(v15_tout.get(iso,0.0)),1),v15_retenues_gwh=round(float(v15_plan.get(iso,0.0)),1)))
o=pd.DataFrame(rows).sort_values("stockage_a_ajouter_gwh_jour",ascending=False)
o["v15_qui_manquent_gwh"]=(o.stockage_a_ajouter_gwh_jour-o.v15_possibles_gwh).clip(lower=0).round(1)
o.to_csv(f"{D}/stockage_soir_pays.csv",index=False,lineterminator="\n")
print(o.head(25).to_string())
print(o[["hydraulique_existante_deplacee_gwh_jour","stockage_v14_gwh_jour","stockage_a_ajouter_gwh_jour","manque_non_stockable_gwh_jour","v15_possibles_gwh","v15_retenues_gwh","v15_qui_manquent_gwh","pointe_du_soir_gw"]].sum().round(0))