import sys
import numpy as np
import pandas as pd
import geonamescache
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from global_land_mask import globe
D = sys.argv[1] if len(sys.argv) > 1 else "."
N = "plan_mondial_artesia_6"
s = pd.read_csv(f"{D}/{N}_sites.csv", keep_default_na=False)
s = s[s.pris > 0]
p = pd.read_csv(f"{D}/{N}_pays.csv", keep_default_na=False).set_index("iso")
l = pd.read_csv(f"{D}/{N}_lignes.csv", keep_default_na=False)
v = pd.read_csv(f"{D}/{N}_v15.csv", keep_default_na=False)
v = v[v.x > 0.01]
gc = geonamescache.GeonamesCache()
pg = {c["iso3"]: c for c in gc.get_countries().values()}
cap = {}
for c in gc.get_cities().values():
    cap.setdefault((c["countrycode"], c["name"]), (c["latitude"], c["longitude"]))
def coord(k):
    c = pg.get(k)
    return cap.get((c["iso"], c["capital"])) if c else None
xy = {k: coord(k) for k in p.index}
p["lat"] = [xy[k][0] if xy[k] else np.nan for k in p.index]
p["lon"] = [xy[k][1] if xy[k] else np.nan for k in p.index]
la, lo = np.meshgrid(np.arange(-60, 85, 0.5), np.arange(-180, 180, 0.5), indexing="ij")
TERRE = np.where(globe.is_land(la, lo), 1.0, np.nan)
def fond(titre):
    fig = plt.figure(figsize=(16, 9))
    g = fig.add_axes([0.01, 0.02, 0.98, 0.91])
    g.imshow(TERRE, extent=[-180, 180, -60, 85], origin="lower", cmap="Greys", vmin=0, vmax=4, aspect="auto")
    x0, y0 = -42, -52
    k = 2000 / 111.32
    g.plot([x0, x0 + k], [y0, y0], color="black", lw=3)
    g.plot([x0, x0], [y0 - 1, y0 + 1], color="black", lw=2)
    g.plot([x0 + k, x0 + k], [y0 - 1, y0 + 1], color="black", lw=2)
    g.text(x0 + k / 2, y0 + 2, "2 000 km (à l'équateur)", ha="center", fontsize=9)
    g.set_xlim(-180, 180)
    g.set_ylim(-60, 85)
    g.set_xticks([])
    g.set_yticks([])
    g.set_title(titre, fontsize=13)
    return fig, g
DEC = {"Afghanistan": (-18, 3), "Pakistan": (-17, -2), "Inde": (2, -4), "Thaïlande": (3, -3), "Birmanie": (-18, 2), "Corée du Sud": (2, -4), "Chine": (2, 3), "Iran": (-4, 3)}
DEC_ECH = {"Arabie saoudite": (-24, -1), "Pakistan": (-16, -5), "Inde": (2, -6), "Birmanie": (3, 4), "Afghanistan": (-16, 3), "Thaïlande": (2, -4), "Corée du Sud": (2, -4), "Chine": (2, 3)}
def etiquette(g, r, txt, dec=None):
    dx, dy = (dec or DEC).get(r.pays, (2, 1.5))
    g.text(r.lon + dx, r.lat + dy, txt, fontsize=8, zorder=4)
def nombre(x):
    if 0 < abs(x) < 10 and x != int(x):
        return f"{x:.1f}".replace(".", ",")
    return f"{x:,.0f}".replace(",", " ")
pos = {}
for _, r in l.iterrows():
    pos.setdefault(r.de, []).append((r.lat_de, r.lon_de))
    pos.setdefault(r.vers, []).append((r.lat_vers, r.lon_vers))
TYPES = [("soleil", "V14 Soleil des plateaux", "#f39c12"), ("mixte", "V14 Soleil mixtes", "#d35400"), ("crete", "V14 des crêtes ventées", "#2e86c1"), ("poste", "V14 près des postes fossiles", "#8e44ad"), ("v21", "V21 sur barrage, au soleil", "#16a085"), ("v21v", "V21 sur barrage, au vent", "#1abc9c"), ("flottant", "Îles hybrides flottantes", "#27ae60"), ("gold4", "Gold 4 sur grandes sources", "#c0392b")]
fig, g = fond("Artesia 6 : où sont les centrales du plan (un point par site, taille selon le nombre de centrales)")
for t, nom, col in TYPES:
    q = s[s.type == t]
    g.scatter(q.lon, q.lat, s=3 + 4 * q.pris.clip(0, 6), color=col, alpha=0.7, lw=0, label=f"{nom} : {nombre(q.pris.sum())} centrales, {nombre(q.twh_pris.sum())} TWh par an")
g.scatter(v.lon, v.lat, s=22, marker="^", color="black", lw=0, label=f"V15 Canyon (stockage) : {nombre(v.x.sum())}")
g.legend(loc="lower left", fontsize=8.5, framealpha=0.92, markerscale=2.2)
fig.savefig(f"{D}/carte_{N}_centrales.png", dpi=100)
plt.close(fig)
fig, g = fond("Artesia 6 : qui exporte et qui importe l'électricité, et les lignes d'export entre pays")
for _, r in l.iterrows():
    g.plot([r.lon_de, r.lon_vers], [r.lat_de, r.lat_vers], color="#5b2c6f", lw=0.4 + r.twh_par_an / 70, alpha=0.5, solid_capstyle="round", zorder=2)
q = p.copy()
for k, r in q.iterrows():
    if r.pays in pos:
        a = np.array(pos[r.pays])
        q.loc[k, "lat"], q.loc[k, "lon"] = a[:, 0].mean(), a[:, 1].mean()
q = q.dropna(subset=["lat"])
q["net"] = q.exporte_twh - q.importe_twh
q = q[q.net.abs() > 5]
ex, im = q[q.net > 0], q[q.net < 0]
g.scatter(ex.lon, ex.lat, s=ex.net * 1.2 + 10, color="#27ae60", edgecolor="black", lw=0.5, alpha=0.8, zorder=3)
g.scatter(im.lon, im.lat, s=-im.net * 1.2 + 10, color="#c0392b", edgecolor="black", lw=0.5, alpha=0.8, zorder=3)
for k, r in q.reindex(q.net.abs().sort_values(ascending=False).index).head(14).iterrows():
    etiquette(g, r, f"{r.pays} {'+' if r.net > 0 else '-'}{nombre(abs(r.net))}", DEC_ECH)
h = [Line2D([], [], marker="o", ls="", color="#27ae60", markeredgecolor="black", markersize=10, label="Pays qui exporte plus qu'il n'importe (TWh par an)"), Line2D([], [], marker="o", ls="", color="#c0392b", markeredgecolor="black", markersize=10, label="Pays qui importe plus qu'il n'exporte (TWh par an)"), Line2D([], [], color="#5b2c6f", lw=3, label=f"Lignes d'export : {len(l)} lignes, {nombre(l.twh_par_an.sum())} TWh par an")]
g.legend(handles=h, loc="lower left", fontsize=9, framealpha=0.92)
fig.savefig(f"{D}/carte_{N}_echanges.png", dpi=100)
plt.close(fig)
fig, g = fond("Artesia 6 : le stockage du soir, pays par pays (GWh rendus chaque soir par les bassins agrandis et les barrages)")
q = p.dropna(subset=["lat"]).copy()
q["soir"] = q.bassins_agrandis_gwh_jour + q.barrages_deplaces_gwh_jour
q = q[q.soir > 5]
g.scatter(q.lon, q.lat, s=q.soir * 0.35 + 8, color="#2e86c1", edgecolor="black", lw=0.5, alpha=0.75, zorder=3)
for k, r in q.sort_values("soir", ascending=False).head(10).iterrows():
    etiquette(g, r, f"{r.pays} {nombre(r.soir)}")
g.scatter(v.lon, v.lat, s=26, marker="^", color="black", lw=0, zorder=4)
h = [Line2D([], [], marker="o", ls="", color="#2e86c1", markeredgecolor="black", markersize=10, label=f"Bassins des V14 agrandis et barrages qui turbinent le soir : {nombre((p.bassins_agrandis_gwh_jour + p.barrages_deplaces_gwh_jour).sum())} GWh par soir"), Line2D([], [], marker="^", ls="", color="black", markersize=8, label=f"V15 Canyon sur les remparts des villes : {nombre(v.x.sum())} centrales")]
g.legend(handles=h, loc="lower left", fontsize=9, framealpha=0.92)
fig.savefig(f"{D}/carte_{N}_soir.png", dpi=100)
plt.close(fig)
CAT = (("faible", 0.55, "#c0392b", "IDH faible"), ("moyen", 0.70, "#e67e22", "IDH moyen"), ("eleve", 0.80, "#2e86c1", "IDH élevé"), ("tres", 9.0, "#7f8c8d", "IDH très élevé ou inconnu"))
def couleur(x):
    try:
        x = float(x)
    except ValueError:
        return "#7f8c8d"
    for n, lim, c, t in CAT:
        if x < lim:
            return c
fig, g = fond("Artesia 6 : l'argent investi dans chaque pays, selon son indice de développement humain (IDH)")
q = p.dropna(subset=["lat"]).copy()
q = q[q.investissement_md_euros > 1]
g.scatter(q.lon, q.lat, s=q.investissement_md_euros * 0.9 + 8, c=[couleur(x) for x in q.idh_rapport_2025], edgecolor="black", lw=0.5, alpha=0.8, zorder=3)
for k, r in q.sort_values("investissement_md_euros", ascending=False).head(12).iterrows():
    etiquette(g, r, f"{r.pays} {nombre(r.investissement_md_euros)}")
h = [Line2D([], [], marker="o", ls="", color=c, markeredgecolor="black", markersize=10, label=f"{t} : {nombre(p[[couleur(x) == c for x in p.idh_rapport_2025]].investissement_md_euros.sum())} milliards d'euros") for n, lim, c, t in CAT]
g.legend(handles=h, loc="lower left", fontsize=9, framealpha=0.92, title="Investi dans le pays sur 25 ans")
fig.savefig(f"{D}/carte_{N}_investissement.png", dpi=100)
plt.close(fig)