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 "."
l = pd.read_csv(f"{D}/reseaux_artesia_4_lignes.csv")
p = pd.read_csv(f"{D}/reseaux_artesia_4_pays.csv")
s = pd.read_csv(f"{D}/plan_mondial_artesia_4_sites.csv", keep_default_na=False, na_values=[""])
fig = plt.figure(figsize=(16, 13))
g1 = fig.add_axes([0.03, 0.36, 0.94, 0.58])
la, lo = np.meshgrid(np.arange(-60, 85, 0.5), np.arange(-180, 180, 0.5), indexing="ij")
terre = globe.is_land(la, lo)
g1.imshow(np.where(terre, 1.0, np.nan), extent=[-180, 180, -60, 85], origin="lower", cmap="Greys", vmin=0, vmax=4, aspect="auto")
s["twh"] = s.gwh_vendus * s.poids / 1000.0
g1.scatter(s.lon, s.lat, s=2, color="#2e8b57", alpha=0.35, label="Centrales placées par Artesia 4")
for _, r in l.iterrows():
    c = "#c0392b" if r.m_euros_par_gwh > 0.41 else "#1f5fa8"
    g1.plot([r.lon_de, r.lon_vers], [r.lat_de, r.lat_vers], color=c, lw=0.4 + r.twh_par_an / 60, alpha=0.75, solid_capstyle="round")
cent = s.groupby("pays").agg(lat=("lat", "mean"), lon=("lon", "mean"))
import pycountry
def a2(i):
    if i == "OWID_KOS":
        return "XK"
    c = pycountry.countries.get(alpha_3=i)
    return c.alpha_2 if c else None
p["a2"] = p.iso.map(a2)
sa = p[p.sans_acheteur_twh > 20].merge(cent, left_on="a2", right_index=True)
g1.scatter(sa.lon, sa.lat, s=sa.sans_acheteur_twh * 1.2, facecolors="none", edgecolors="#e67e22", lw=2, label="Production sans acheteur à portée (TWh/an, cercle orange)")
for _, r in sa[sa.sans_acheteur_twh > 180].iterrows():
    g1.annotate(f"{r.pays} {r.sans_acheteur_twh:,.0f}".replace(",", " "), (r.lon, r.lat), xytext=(-95, -4) if r.pays == "Nouvelle-Zélande" else (8, -14), textcoords="offset points", fontsize=8, color="#a04000")
nc = {"Chine": (34.5, 112.0), "Japon": (36.0, 138.5), "Corée du Sud": (36.5, 127.8), "Indonésie": (-3.0, 112.0), "Taïwan": (23.7, 121.0), "Inde": (21.0, 78.3), "Viêt Nam": (16.0, 106.5), "Malaisie": (3.5, 102.0), "Australie": (-31.0, 147.0), "Philippines": (13.0, 122.0), "Thaïlande": (15.0, 100.8)}
q = p[p.pays.isin(nc.keys())]
for _, r in q.iterrows():
    y, x = nc[r.pays]
    g1.scatter([x], [y], s=r.non_couvert_twh * 1.2, facecolors="none", edgecolors="#8e44ad", lw=2)
    if r.non_couvert_twh > 140:
        dx, dy = {"Chine": (-60, 40), "Japon": (12, 12), "Corée du Sud": (-20, 28), "Taïwan": (14, -4), "Inde": (-40, 22), "Indonésie": (10, 10)}.get(r.pays, (8, 8))
        g1.annotate(f"{r.pays} {r.non_couvert_twh:,.0f}".replace(",", " "), (x, y), xytext=(dx, dy), textcoords="offset points", fontsize=8, color="#5b2c6f")
g1.scatter([], [], s=80, facecolors="none", edgecolors="#8e44ad", lw=2, label="Électricité fossile sans centrale à portée (TWh/an, cercle violet)")
g1.plot([], [], color="#1f5fa8", lw=3, label="Ligne d'export moins chère qu'une V14 sur place")
g1.plot([], [], color="#c0392b", lw=3, label="Ligne d'export plus chère qu'une V14 sur place")
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)
g1.set_title("Artesia 4, réseaux recalculés : les lignes d'export choisies au moindre coût (5 000 km au plus, dont 1 500 km sous la mer)", fontsize=12)
g2 = fig.add_axes([0.17, 0.05, 0.78, 0.25])
post = ["Raccordement des centrales", "Lignes intérieures longues", "Lignes d'export", "Renforcement à l'arrivée", "Pertes (centrales en plus)"]
v = [float(x) for x in sys.argv[2].split(",")] if len(sys.argv) > 2 else [357, 324, 1005, 396, 225]
gauche = 0
cols = ["#7f8c8d", "#95a5a6", "#1f5fa8", "#5d6d7e", "#e67e22"]
for n, x, c in zip(post, v, cols):
    g2.barh([1], [x], left=gauche, color=c, label=f"{n} : {x:,.0f} Md€".replace(",", " "))
    gauche += x
g2.barh([0], [1000], color="#bdc3c7", label="Forfait d'Artesia 4 : 1 000 Md€ (réseaux et stockage)")
g2.set_yticks([0, 1])
g2.set_yticklabels(["Forfait Artesia 4", "Recalcul (10 714 TWh livrés)"])
g2.set_xlabel("Milliards d'euros sur 25 ans")
g2.legend(loc="lower right", fontsize=9)
g2.grid(axis="x", alpha=0.3)
fig.savefig(f"{D}/schema_reseaux_artesia_4.png", dpi=110)