import sys
import numpy as np
import pandas as pd
import geonamescache
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_6"
p = pd.read_csv(f"{D}/{N}_pays.csv", keep_default_na=False).set_index("iso")
b = pd.read_csv(f"{D}/{N}_balayage.csv")
s = pd.read_csv(f"{D}/{N}_sites.csv", keep_default_na=False)
gc = geonamescache.GeonamesCache()
pays_gc = {v["iso3"]: v for v in gc.get_countries().values()}
villes = gc.get_cities()
cap = {}
for v in villes.values():
    cap.setdefault((v["countrycode"], v["name"]), (v["latitude"], v["longitude"]))
def coord(k):
    c = pays_gc.get(k)
    if not c:
        return None
    return cap.get((c["iso"], c["capital"]))
p["kwh"] = pd.to_numeric(p.kwh_par_habitant_avant, errors="coerce")
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")
g1.scatter(s.lon, s.lat, s=1.5, color="#7f8c8d", alpha=0.35, lw=0)
g1.scatter([], [], s=12, color="#7f8c8d", label="Centrales du plan")
q = p[p.ajout_500_kwh_twh > 0.05].copy()
xy = [coord(k) for k in q.index]
q["lat"] = [c[0] if c else np.nan for c in xy]
q["lon"] = [c[1] if c else np.nan for c in xy]
q = q.dropna(subset=["lat"])
sc = g1.scatter(q.lon, q.lat, s=q.ajout_500_kwh_twh * 18 + 8, c=q.sans_electricite_millions, cmap="YlOrRd", vmin=0, vmax=60, edgecolor="black", lw=0.5, alpha=0.85, zorder=3)
for k, r in q.sort_values("ajout_500_kwh_twh", ascending=False).head(9).iterrows():
    g1.text(r.lon + 2.5, r.lat + 1.5, r.pays.replace("République démocratique du Congo", "RD Congo"), fontsize=8.5, zorder=4)
cb = fig.colorbar(sc, ax=g1, fraction=0.025, pad=0.01)
cb.set_label("Millions de personnes sans électricité")
for t, v in ((10, "10 TWh"), (40, "40 TWh"), (80, "80 TWh")):
    g1.scatter([], [], s=t * 18 + 8, color="#f39c12", edgecolor="black", lw=0.5, label=f"Électricité ajoutée par an : {v}")
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, labelspacing=2.6, borderpad=1.6, handletextpad=2.6)
g1.set_title("Artesia 6 : au moins 500 kWh par habitant et par an, et l'électricité pour les 696 millions de personnes qui ne l'ont pas", fontsize=12)
g2 = fig.add_axes([0.07, 0.06, 0.40, 0.27])
top = p.sort_values("ajout_500_kwh_twh", ascending=False).head(12)[::-1]
noms = [n.replace("République démocratique du Congo", "RD Congo") for n in top.pays]
g2.barh(noms, top.ajout_500_kwh_twh, color="#f39c12", edgecolor="black", lw=0.5)
for i, (v, w) in enumerate(zip(top.ajout_500_kwh_twh, top.sans_electricite_millions)):
    g2.text(v + 1, i, f"{v:.0f} TWh, {w:.0f} M sans électricité", va="center", fontsize=8)
g2.set_xlim(0, 125)
g2.set_xlabel("Électricité ajoutée par an (TWh)")
g2.set_title("Les douze pays qui reçoivent le plus", fontsize=10.5)
g2.grid(axis="x", alpha=0.3)
g3 = fig.add_axes([0.56, 0.06, 0.40, 0.27])
g3.plot(b.seuil, b.total, "o-", color="black", label="Coût du plan (Md€ sur 25 ans)")
g3.axhline(10000, color="#c0392b", ls="--", lw=1)
g3.text(1000, 9990, "10 000 milliards", color="#c0392b", fontsize=8.5, ha="right", va="top")
g3.axvline(500, color="#2e86c1", lw=1)
g3.text(510, b.total.max() - 20, "Artesia 6", color="#2e86c1", fontsize=9, va="top")
g3.set_xlabel("Minimum garanti par habitant et par an (kWh)")
g3.set_ylabel("Milliards d'euros sur 25 ans")
g4 = g3.twinx()
g4.plot(b.seuil, b.ajout_twh, "s-", color="#f39c12", label="Électricité ajoutée (TWh par an)")
g4.set_ylabel("TWh par an")
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 le minimum garanti", fontsize=10.5)
g3.grid(alpha=0.3)
fig.savefig(f"{D}/schema_{N}.png", dpi=110)