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_sites.csv")
l = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_lignes.csv")
t = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_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é (nouvelles)"), "v21": ("#16a085", "Barrages existants 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", "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, 5), color=col, alpha=0.6, lw=0, label=f"{lab} : {n:,.0f}".replace(",", " ") if lab else None)
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=3)
g1.set_title("Artesia 5, version 1, économie mondiale : les centrales choisies au coût du kWh livré, ligne comprise", fontsize=12)
g2 = fig.add_axes([0.2, 0.05, 0.76, 0.26])
cat = [("Centrales Hydroloop, barrages et flottant", "#2e8b57"), ("V15, tours, Pacha, Gold existantes, O Sky", "#7f8c8d"), ("Panneaux hybrides (20 % au plus par pays)", "#f1c40f"), ("Raccordement, lignes intérieures, renforcement", "#95a5a6"), ("Lignes d'export", "#5b2c6f"), ("Forfait réseaux et stockage", "#bdc3c7")]
lignes = [("Artesia 4, page (forfait)", [7685 - 932, 932, 0, 0, 0, 1000]), ("Artesia 4, réseau recalculé\n(62 % livrés)", [7685 - 932, 932, 0, 357 + 324 + 396 + 225, 1005, 0]), ("Artesia 5 version 1\n(95 % livrés partout)", [t.cent + t.sky, t.fixe, t.hyb, t.racc + t.hyb_res + t.res_fixe, t.exp, 0])]
for i, (nom, v) in enumerate(lignes):
    gauche = 0
    for (n, col), x in zip(cat, v):
        g2.barh([i], [x], left=gauche, color=col, label=n if (i == 2 and not n.startswith("Forfait")) or (i == 0 and n.startswith("Forfait")) 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, v in lignes])
g2.set_xlim(0, 14000)
g2.set_xlabel("Milliards d'euros sur 25 ans (scénario B, sans la Francis)")
g2.legend(loc="lower right", fontsize=8.5)
g2.grid(axis="x", alpha=0.3)
fig.savefig(f"{D}/schema_plan_mondial_artesia_5_v1.png", dpi=110)