import sys
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from global_land_mask import globe
D = sys.argv[1] if len(sys.argv) > 1 else "."
s = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_recalcul_sites.csv")
v = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_recalcul_v15.csv")
l = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_recalcul_lignes.csv")
t = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_recalcul_totaux.csv", index_col=0).valeur
u = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_recalcul_souple_totaux.csv", index_col=0).valeur
fig = plt.figure(figsize=(16, 13))
g1 = fig.add_axes([0.03, 0.38, 0.94, 0.57])
la, lo = np.meshgrid(np.arange(-60, 85, 0.5), np.arange(-180, 180, 0.5), indexing="ij")
g1.imshow(np.where(globe.is_land(la, lo), 1.0, np.nan), extent=[-180, 180, -60, 85], origin="lower", cmap="Greys", vmin=0, vmax=4, aspect="auto")
st = {"crete": ("#1f5fa8", "V14 de crête"), "soleil": ("#e67e22", "V14 Soleil et mixtes des plateaux"), "mixte": ("#e67e22", None), "poste": ("#c0392b", "V14 Soleil près d'un poste fossile libéré"), "v21": ("#16a085", "Barrages au soleil en V21"), "v21v": ("#0b5345", "Barrages tempérés au vent en V21"), "flottant": ("#8e44ad", "Barrages avec îles hybrides flottantes")}
for k, (col, lab) in st.items():
    q = s[s.type == k]
    n = len(q) if k in ("v21", "v21v", "flottant") else q.pris.sum() + (s[s.type == "mixte"].pris.sum() if k == "soleil" else 0)
    g1.scatter(q.lon, q.lat, s=2 + 3 * q.pris.clip(0, 6), color=col, alpha=0.6, lw=0, label=f"{lab} : {n:,.0f}".replace(",", " ") if lab else None)
m = s[s.v14_en_max >= 6]
g1.scatter(m.lon, m.lat, s=40, facecolors="none", edgecolors="black", lw=0.8, label=f"Sites avec au moins une MAX 3 : {len(m):,.0f}".replace(",", " "))
g1.scatter(v.lon, v.lat, s=30, marker="^", color="black", label=f"V15 Canyon retenues : {v.x.sum():,.0f}".replace(",", " "))
for _, r in l.iterrows():
    g1.plot([r.lon_de, r.lon_vers], [r.lat_de, r.lat_vers], color="#5b2c6f", lw=0.4 + r.twh_par_an / 80, alpha=0.6, solid_capstyle="round")
g1.plot([], [], color="#5b2c6f", lw=3, label=f"Lignes d'export : {t.twh_exp:,.0f} TWh par an".replace(",", " "))
x0, y0 = -42, -52
k = 2000 / 111.32
g1.plot([x0, x0 + k], [y0, y0], color="black", lw=3)
g1.plot([x0, x0], [y0 - 1, y0 + 1], color="black", lw=2)
g1.plot([x0 + k, x0 + k], [y0 - 1, y0 + 1], color="black", lw=2)
g1.text(x0 + k / 2, y0 + 2, "2 000 km (à l'équateur)", ha="center", fontsize=9)
g1.set_xlim(-180, 180)
g1.set_ylim(-60, 85)
g1.set_xticks([])
g1.set_yticks([])
g1.legend(loc="lower left", fontsize=8.5, framealpha=0.9, markerscale=2)
g1.set_title("Artesia 5, version 1 recalculée : prix de chaque pays, versions MAX, V15 placées et stockage du soir", fontsize=12)
g2 = fig.add_axes([0.2, 0.05, 0.76, 0.26])
cat = [("Centrales Hydroloop, barrages et flottant", "#2e8b57"), ("Tours, Pacha, Gold existantes, O Sky (et V15 dans la 1re version)", "#7f8c8d"), ("Panneaux hybrides", "#f1c40f"), ("Raccordement, lignes intérieures, renforcement", "#95a5a6"), ("Lignes d'export", "#5b2c6f"), ("Stockage du soir : bassins des V14 agrandis et V15", "#1a5276")]
lignes = [("Version 1 du 29/10, matin\n(sans le stockage du soir)", [6691.0, 961.6, 1336.8, 1389.6 + 63.3 + 37.9, 823.5, 0.0]), ("Version 1 recalculée,\ndemande d'aujourd'hui", [t.cent + t.sky, t.fixe, t.hyb, t.racc + t.hyb_res + t.res_fixe, t.exp, t.ext + t.v15]), ("Version 1 recalculée,\ndemande souple", [u.cent + u.sky, u.fixe, u.hyb, u.racc + u.hyb_res + u.res_fixe, u.exp, u.ext + u.v15])]
lignes = lignes[::-1]
for i, (nom, w) in enumerate(lignes):
    gauche = 0
    for (n, col), x in zip(cat, w):
        g2.barh([i], [x], left=gauche, color=col, label=n if i == 1 else None)
        gauche += x
    g2.text(gauche + 60, i, f"{gauche:,.0f} Md€".replace(",", " "), va="center", fontsize=10)
g2.set_yticks(range(len(lignes)))
g2.set_yticklabels([n for n, w in lignes])
g2.set_xlim(0, 15800)
g2.set_xlabel("Milliards d'euros sur 25 ans (scénario B, sans la Francis)")
g2.legend(loc="lower right", fontsize=8)
g2.grid(axis="x", alpha=0.3)
fig.savefig(f"{D}/schema_plan_mondial_artesia_5_v1_recalcul.png", dpi=110)