import sys, csv, math, collections
import numpy as np
import pandas as pd
import pycountry
import geonamescache
from scipy.spatial import cKDTree
D = sys.argv[1] if len(sys.argv) > 1 else "."
RAYON = float(sys.argv[2]) if len(sys.argv) > 2 else 50.0
LF = 0.5
GW = 1.0 / (8.76 * LF)
AC, POSTE, POSTE_EXISTANT, ARRIVEE, CONV, TERRE, DETOUR, SEUIL_DOM = 0.60, 50.0, 20.0, 100.0, 200.0, 0.30, 1.15, 300.0
TOTAL_PLAN, V14F, NONLOC = 17360.0, 1328.0, 1328.0 + 405.0 + 43.0 + 420.0 + 63.0
FOSSILE = {"Coal", "Gas", "Oil", "Petcoke"}
def iso3(a2):
    if a2 == "XK":
        return "OWID_KOS"
    p = pycountry.countries.get(alpha_2=a2)
    return p.alpha_3 if p else None
def xyz(lat, lon):
    la, lo = np.radians(lat), np.radians(lon)
    return np.c_[np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)]
def km(a, b, c, d):
    a, b, c, d = map(math.radians, (a, b, c, d))
    return 6371.0 * 2 * math.asin(math.sqrt(math.sin((c - a) / 2) ** 2 + math.cos(a) * math.cos(c) * math.sin((d - b) / 2) ** 2))
g = pd.read_csv(f"{D}/global_power_plant_database.csv", low_memory=False)
g["fossile"] = g.primary_fuel.isin(FOSSILE)
g = g[(g.fossile & (g.capacity_mw >= 50)) | (~g.fossile & (g.capacity_mw >= 100))].reset_index(drop=True)
g["libere_gw"] = np.where(g.fossile, 0.95 * g.capacity_mw / 1000.0, 0.0)
lib_tot = g.libere_gw.sum()
g["libre_gw"] = g.libere_gw * (1 - V14F * GW / lib_tot)
pays = pd.read_csv(f"{D}/electricite_pays_prix_part_fossile.csv").set_index("iso")
s = pd.read_csv(f"{D}/plan_mondial_artesia_4_sites.csv", keep_default_na=False, na_values=[""])
s["iso"] = s.pays.map(lambda x: iso3(x) if isinstance(x, str) else None)
s = s[s.iso.isin(pays.index)].copy().reset_index(drop=True)
s["twh"] = s.gwh_vendus * s.poids / 1000.0
s["gw"] = s.twh * GW
F = pays.fossile_twh
loc = s.groupby("iso").twh.sum().reindex(pays.index, fill_value=0.0)
offre = loc + NONLOC * F / F.sum()
besoin = F * TOTAL_PLAN / F.sum()
local = np.minimum(offre, besoin)
s["part_locale"] = s.iso.map(lambda c: local[c] / offre[c] if offre[c] > 0 else 0.0)
villes = collections.defaultdict(list)
for v in geonamescache.GeonamesCache(min_city_population=15000).get_cities().values():
    c = iso3(v["countrycode"])
    if c:
        villes[c].append((v["latitude"], v["longitude"], v["population"]))
gros = {c: [(a, b) for a, b, p in v if p >= 1000000] or [(a, b) for a, b, p in sorted(v, key=lambda t: -t[2])[:1]] for c, v in villes.items()}
s["dist_dom"] = [min((km(r.lat, r.lon, a, b) for a, b in gros.get(r.iso, [])), default=0.0) for r in s.itertuples()]
s["dom_km"] = np.where(s.dist_dom > SEUIL_DOM, s.dist_dom * DETOUR, 0.0)
s["d_ville"] = np.where(s.centrale == "v21", 0.0, s.dist_grande_ville_km.fillna(0.0))
s["racc_avant"] = np.where(s.centrale == "v21", 0.0, s.gw * (s.d_ville * AC + POSTE))
s["dom_avant"] = np.where(s.dom_km > 0, s.gw * s.part_locale * (CONV + s.dom_km * TERRE), 0.0)
s["arr_avant"] = s.gw * ARRIVEE
arbre = cKDTree(xyz(g.latitude.values, g.longitude.values))
corde = 2 * math.sin(RAYON / 6371.0 / 2)
P = xyz(s.lat.values, s.lon.values)
voisins = arbre.query_ball_point(P, corde)
libre = g.libre_gw.values.copy()
choix, dist, sorte = [-1] * len(s), [0.0] * len(s), ["aucun"] * len(s)
ordre = np.argsort(-s.gw.values)
for i in ordre:
    if s.centrale[i] == "v21":
        sorte[i] = "barrage"
        continue
    vs = voisins[i]
    if not vs:
        continue
    ds = 6371.0 * 2 * np.arcsin(np.clip(np.linalg.norm(g.loc[vs, ["latitude", "longitude"]].pipe(lambda t: xyz(t.latitude.values, t.longitude.values)) - P[i], axis=1) / 2, 0, 1))
    best = None
    for k in np.argsort(ds):
        j = vs[k]
        if g.fossile[j] and libre[j] >= s.gw[i]:
            best = (j, ds[k], "fossile")
            break
    if best is None:
        k = int(np.argmin(ds))
        best = (vs[k], ds[k], "autre")
    j, d, t = best
    if t == "fossile":
        libre[j] -= s.gw[i]
    choix[i], dist[i], sorte[i] = j, float(d) * DETOUR, t
s["poste"], s["d_poste"], s["sorte"] = choix, dist, sorte
s["d_apres"] = np.where(s.sorte.isin(["fossile", "autre"]), np.minimum(s.d_poste, s.d_ville.where(s.d_ville > 0, s.d_poste)), s.d_ville)
s["racc_apres"] = np.where(s.sorte == "barrage", 0.0, np.where(s.sorte.isin(["fossile", "autre"]), s.gw * (s.d_apres * AC + POSTE_EXISTANT), s.racc_avant))
s["dom_apres"] = np.where(s.sorte == "fossile", 0.0, s.dom_avant)
s["arr_apres"] = np.where(s.sorte == "fossile", 0.0, s.arr_avant)
nl_gw = (NONLOC - V14F) * GW
racc_nl = nl_gw * (30 * AC + POSTE)
arr_nl_avant = NONLOC * GW * ARRIVEE
arr_nl_apres = nl_gw * ARRIVEE
EXPORT, PERTES_EXP = 1005.0, None
av = dict(racc=(s.racc_avant.sum() + racc_nl) / 1000, dom=s.dom_avant.sum() / 1000, arr=(s.arr_avant.sum() + arr_nl_avant) / 1000)
ap = dict(racc=(s.racc_apres.sum() + racc_nl) / 1000, dom=s.dom_apres.sum() / 1000, arr=(s.arr_apres.sum() + arr_nl_apres) / 1000)
pertes_ac_av = (s.twh * s.d_ville * 0.005 / 100).sum()
pertes_ac_ap = (s.twh * np.where(s.sorte.isin(["fossile", "autre"]), s.d_apres, s.d_ville) * 0.005 / 100).sum()
pertes_dom_av = (s.twh * s.part_locale * np.where(s.dom_km > 0, 0.015 + 0.03 * s.dom_km / 1000, 0)).sum()
pertes_dom_ap = (s.twh * s.part_locale * np.where((s.dom_km > 0) & (s.sorte != "fossile"), 0.015 + 0.03 * s.dom_km / 1000, 0)).sum()
gain_pertes = (pertes_ac_av - pertes_ac_ap + pertes_dom_av - pertes_dom_ap) * 0.41
tot_av = sum(av.values()) + EXPORT
tot_ap = sum(ap.values()) + EXPORT
n = s.sorte.value_counts().to_dict()
tw = s.groupby("sorte").twh.sum().to_dict()
gw_used = (g.libre_gw.values - libre).sum()
print(f"noeuds {len(g)} fossiles {int(g.fossile.sum())} libere {lib_tot:.0f} GW apres V14F {g.libre_gw.sum():.0f} GW utilise {gw_used:.0f} GW reste {libre.sum():.0f} GW")
print("centrales", n, "twh", {k: round(v) for k, v in tw.items()})
print("avant", {k: round(v) for k, v in av.items()}, round(tot_av), "apres", {k: round(v) for k, v in ap.items()}, round(tot_ap), "gain", round(tot_av - tot_ap), "pertes TWh", round(pertes_ac_av - pertes_ac_ap + pertes_dom_av - pertes_dom_ap), "Md", round(gain_pertes))
print("gain par sorte", {k: round((s[s.sorte == k].racc_avant + s[s.sorte == k].dom_avant + s[s.sorte == k].arr_avant - s[s.sorte == k].racc_apres - s[s.sorte == k].dom_apres - s[s.sorte == k].arr_apres).sum() / 1000) for k in ["fossile", "autre", "barrage", "aucun"]}, "gain V14F arrivee", round((arr_nl_avant - arr_nl_apres) / 1000))
print("d_poste moyen fossile", round(s[s.sorte == "fossile"].d_apres.mean()), "autre", round(s[s.sorte == "autre"].d_apres.mean()), "d_ville moyen", round(s[s.sorte != "barrage"].d_ville.mean()))
g["utilise_gw"] = g.libre_gw - libre
g["reste_gw"] = libre
if len(sys.argv) <= 2:
    g["iso"] = g.country
    r = g[g.fossile].groupby("iso").agg(libere_gw=("libere_gw", "sum"), utilise_gw=("utilise_gw", "sum"), reste_gw=("reste_gw", "sum"), postes=("name", "count")).sort_values("libere_gw", ascending=False)
    r["pays"] = [pays.pays.get(i, i) for i in r.index]
    r.round(1).to_csv(f"{D}/postes_fossiles_liberes_pays.csv", lineterminator="\n")
    o = s[["rang", "centrale", "lat", "lon", "pays", "twh", "gw", "sorte", "d_apres", "d_ville"]].copy()
    o["poste"] = [g.name[j] if j >= 0 else "" for j in s.poste]
    o["combustible_poste"] = [g.primary_fuel[j] if j >= 0 else "" for j in s.poste]
    o["poste_lat"] = [round(g.latitude[j], 4) if j >= 0 else "" for j in s.poste]
    o["poste_lon"] = [round(g.longitude[j], 4) if j >= 0 else "" for j in s.poste]
    o.round(2).to_csv(f"{D}/postes_artesia_4_centrales.csv", index=False, lineterminator="\n")
    g[["name", "country", "primary_fuel", "capacity_mw", "latitude", "longitude", "libere_gw", "utilise_gw", "reste_gw"]].round(3).to_csv(f"{D}/postes_noeuds.csv", index=False, lineterminator="\n")
    print(r.head(15).to_string())