import sys
import numpy as np
import pandas as pd
import pycountry
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 "."
N = "plan_mondial_artesia_5_v2_developpement_humain"
s = pd.read_csv(f"{D}/{N}_sites.csv", keep_default_na=False)
l = pd.read_csv(f"{D}/{N}_lignes.csv")
b = pd.read_csv(f"{D}/{N}_balayage.csv")
idh = pd.read_csv(f"{D}/idh_pnud_2025.csv").set_index("iso").idh_rapport_2025.to_dict()
def iso3(c):
    if len(c) == 3:
        return c
    p = pycountry.countries.get(alpha_2=c)
    return p.alpha_3 if p else None
COUL = (("faible", 0.55, "#c0392b", "IDH faible (moins de 0,55)"), ("moyen", 0.70, "#e67e22", "IDH moyen (0,55 à 0,70)"), ("eleve", 0.80, "#2e86c1", "IDH élevé (0,70 à 0,80)"), ("tres_eleve", 9.0, "#7f8c8d", "IDH très élevé (0,80 et plus) ou inconnu"))
def cat(k):
    i = idh.get(k)
    if i is None:
        return "tres_eleve"
    for n, lim, c, t in COUL:
        if i < lim:
            return n
s["cat"] = [cat(iso3(p)) for p in s.pays]
fig = plt.figure(figsize=(16, 13))
g1 = fig.add_axes([0.03, 0.40, 0.94, 0.55])
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")
for n, lim, col, txt in COUL:
    q = s[s.cat == n]
    g1.scatter(q.lon, q.lat, s=2 + 3 * q.pris.clip(0, 6), color=col, alpha=0.65, lw=0, label=f"Centrales dans un pays à {txt} : " + f"{q.twh_pris.sum():,.0f}".replace(",", " ") + " TWh par an")
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.55, solid_capstyle="round")
g1.plot([], [], color="#5b2c6f", lw=3, label=f"Lignes d'export : {l.twh_par_an.sum():,.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 2, développement humain : les centrales du plan selon l'indice de développement humain de leur pays", fontsize=12)
v1 = b[(b.alpha == 0) & (b.dep == 1)].iloc[0]
v2 = b[(b.alpha == 0.25) & (b.dep == 1)].iloc[0]
g2 = fig.add_axes([0.07, 0.06, 0.40, 0.27])
x = np.arange(4)
w = 0.38
for j, (r, lab, a) in enumerate(((v1, "Version 1 recalculée", 0.45), (v2, "Version 2", 1.0))):
    y = [r[f"inv_{n}"] for n, lim, col, t in COUL]
    y[3] += r.inv_sans_idh
    bars = g2.bar(x + (j - 0.5) * w, y, w, color=[c for n, lim, c, t in COUL], alpha=a, edgecolor="black", lw=0.5, label=lab)
    for xi, yi in zip(x + (j - 0.5) * w, y):
        g2.text(xi, yi + 40, f"{yi:,.0f}".replace(",", " "), ha="center", fontsize=8.5)
g2.set_xticks(x)
g2.set_xticklabels(["IDH faible", "IDH moyen", "IDH élevé", "IDH très élevé"])
g2.set_ylabel("Milliards d'euros investis dans le pays")
g2.set_title("Où va l'argent : version 1 (pâle) et version 2 (vive)", fontsize=10.5)
g2.grid(axis="y", alpha=0.3)
g2.set_ylim(0, 4500)
g3 = fig.add_axes([0.56, 0.06, 0.40, 0.27])
q = b[b.dep == 1].sort_values("alpha")
g3.plot(q.alpha * 100, q.total, "o-", color="black", label="Coût réel du plan")
g3.set_xlabel("Remise accordée dans le choix à un pays à IDH faible (%)")
g3.set_ylabel("Milliards d'euros sur 25 ans")
g3.axhline(10000, color="#c0392b", ls="--", lw=1)
g3.text(39.5, 9993, "10 000 milliards", color="#c0392b", fontsize=8.5, ha="right")
g4 = g3.twinx()
g4.plot(q.alpha * 100, q.inv_faible + q.inv_moyen, "s-", color="#e67e22", label="Investi dans les pays à IDH faible et moyen")
g4.plot(q.alpha * 100, q.china_russie, "^-", color="#5b2c6f", label="Russie vers Chine, TWh par an")
g4.set_ylabel("Milliards d'euros, ou TWh par an")
g3.axvline(25, color="#2e86c1", lw=1)
g3.text(25.3, q.total.min() + 5, "version 2", color="#2e86c1", fontsize=9)
h1, l1 = g3.get_legend_handles_labels()
h2, l2 = g4.get_legend_handles_labels()
g3.legend(h1 + h2, l1 + l2, loc="upper left", fontsize=8)
g3.set_title("Ce que coûte la remise au développement", fontsize=10.5)
g3.grid(alpha=0.3)
fig.savefig(f"{D}/schema_{N}.png", dpi=110)