import sys, csv, math, collections
import numpy as np
import pandas as pd
import pycountry
import geonamescache
from scipy.spatial import cKDTree
from scipy.optimize import linprog
from scipy.sparse import lil_matrix
from global_land_mask import globe
D = sys.argv[1] if len(sys.argv) > 1 else "."
MAX_KM = float(sys.argv[2]) if len(sys.argv) > 2 else 5000.0
MAX_MER = float(sys.argv[3]) if len(sys.argv) > 3 else 1500.0
DIST_MAX = float(sys.argv[4]) if len(sys.argv) > 4 else 400.0
ECRIRE = len(sys.argv) <= 2
SOLEIL = [(1800, 871.3), (1900, 917.4), (2000, 959.8), (2100, 998.5), (2200, 1033.7), (2300, 1065.7), (2400, 1094.8), (2500, 1121.5), (2600, 1145.6), (2700, 1167.1)]
MIXTE = [(1800, 940.1), (2100, 1002.3), (2400, 1054.8), (2700, 1096.3)]
VENT = [(5.5, 662.8), (5.75, 715.1), (6.0, 766.0), (6.25, 815.2), (6.5, 862.7), (6.75, 908.1), (7.0, 951.0), (7.5, 1030.3), (8.0, 1099.1), (8.5, 1155.5), (9.0, 1201.0)]
CALAGE_V14 = 885.0 / 899.4
COUT_SOLEIL, COUT_MIXTE, COUT_CRETE, EAU_APPOINT = 465.0, 461.0, 364.0, 15.0
PACHA, GOLD, FRAC_GOLD, PART_V21, PRUDENCE_VENT = (155.6, 31.1), (56.0, 6.4), 0.47, 2.0 / 3.0, 0.9
FIXE = {"V15 Canyon": (1000, 405.0, 747.0), "Titan Claudia": (1500, 43.0, 49.0), "Pacha existantes": (2700, 420.0, 84.0), "Gold existantes": (1126, 63.0, 7.2), "Sky existantes": (356, 0.0, 45.1), "O Sky sur les Sky": (583, 56.0, 29.3)}
HYBRIDE = float(sys.argv[5]) if len(sys.argv) > 5 else 810.0
CAP = float(sys.argv[6]) if len(sys.argv) > 6 else 0.2
TOTAL_PLAN = 17360.0
K_CHINE = float(sys.argv[7]) if len(sys.argv) > 7 else 0.6
FLOTTANT = (21.6, 85.2, 1100.0, 73.1)
LF = 0.5
GW = 1.0 / (8.76 * LF)
AC, POSTE, POSTE_EX, ARRIVEE, CONV, TERRE, MER, DETOUR, SEUIL_DOM, RAYON = 0.60, 50.0, 20.0, 100.0, 200.0, 0.30, 1.60, 1.15, 300.0, 50.0
def interp_bas(t, h):
    if h < t[0][0]:
        return t[0][1] - (t[1][1] - t[0][1]) * (t[0][0] - h) / (t[1][0] - t[0][0])
    return interp(t, h)
def interp(t, h):
    h = min(max(h, t[0][0]), t[-1][0])
    for (a, x), (b, y) in zip(t, t[1:]):
        if a <= h <= b:
            return x + (y - x) * (h - a) / (b - a)
    return t[-1][1]
def iso3(c):
    if c == "XK":
        return "OWID_KOS"
    if len(c) == 3:
        return c
    p = pycountry.countries.get(alpha_2=c)
    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(min(1.0, math.sin((c - a) / 2) ** 2 + math.cos(a) * math.cos(c) * math.sin((d - b) / 2) ** 2)))
def route(a, b, c, d):
    L = km(a, b, c, d)
    n = max(int(L / 25), 2)
    p = np.linspace(0, 1, n)
    v1, v2 = xyz(np.array([a]), np.array([b]))[0], xyz(np.array([c]), np.array([d]))[0]
    om = math.acos(max(-1.0, min(1.0, float(v1 @ v2))))
    if om < 1e-9:
        return 0.0, 0.0
    q = (np.sin((1 - p) * om)[:, None] * v1 + np.sin(p * om)[:, None] * v2) / math.sin(om)
    la, lo = np.degrees(np.arcsin(np.clip(q[:, 2], -1, 1))), np.degrees(np.arctan2(q[:, 1], q[:, 0]))
    mer = L * (1 - globe.is_land(np.clip(la, -89.99, 89.99), lo).mean())
    return (L - mer) * DETOUR, mer
def centre(lats, lons, w):
    v = (xyz(np.asarray(lats, float), np.asarray(lons, float)) * np.asarray(w, float)[:, None]).sum(0)
    return math.degrees(math.atan2(v[2], math.hypot(v[0], v[1]))), math.degrees(math.atan2(v[1], v[0]))
pays = pd.read_csv(f"{D}/electricite_pays_prix_part_fossile.csv").set_index("iso")
F = pays.fossile_twh
besoin = F * TOTAL_PLAN / F.sum()
nonloc_twh = sum(v[1] for v in FIXE.values())
nonloc = nonloc_twh * F / F.sum()
reb = {(round(float(r["lat"]), 3), round(float(r["lon"]), 3)): float(r["rebord_km_6ms"]) for r in csv.DictReader(open(f"{D}/plateaux_rebord_vent.csv", encoding="utf-8"))}
def pacha(p):
    return (PACHA[0] + FRAC_GOLD * GOLD[0], PACHA[1] + FRAC_GOLD * GOLD[1]) if p >= 1500 else (0.0, 0.0)
C = []
for r in csv.DictReader(open(f"{D}/plateaux_ceinture_soleil.csv", encoding="utf-8")):
    if r["classe"] == "C" or r["occupation"] == "occupe" or int(r["dist_grande_ville_km"]) > DIST_MAX:
        continue
    n = int(r["sites"]) if r["occupation"] == "libre" else int(r["sites"]) // 2
    if n == 0:
        continue
    h, p = int(r["heures_tracker"]), int(r["pluie_mm"])
    kmr = reb.get((round(float(r["lat"]), 3), round(float(r["lon"]), 3)), 0.0)
    nmix = min(n, int(kmr // 12) if p >= 200 else 0)
    gp, kp = pacha(p)
    gs, ks = interp(SOLEIL, h), COUT_SOLEIL + (EAU_APPOINT if p < 75 else 0.0)
    gm = interp(MIXTE, h)
    if nmix and COUT_MIXTE / gm < ks / gs:
        C.append(dict(type="mixte", n=nmix, g=gm + gp, k=COUT_MIXTE + kp, pacha=p >= 1500, lat=float(r["lat"]), lon=float(r["lon"]), pays=r["pays"], ville=r["ville"], dv=float(r["dist_grande_ville_km"])))
        n -= nmix
    if n:
        C.append(dict(type="soleil", n=n, g=gs + gp, k=ks + kp, pacha=p >= 1500, lat=float(r["lat"]), lon=float(r["lon"]), pays=r["pays"], ville=r["ville"], dv=float(r["dist_grande_ville_km"])))
for r in csv.DictReader(open(f"{D}/cretes_vent_v14.csv", encoding="utf-8")):
    if int(r["empl_AB"]) == 0 or r["occupation"] == "occupe" or int(r["dist_grande_ville_km"]) > DIST_MAX or r["vent_150m_moyen"] == "":
        continue
    n = int(r["empl_AB"]) if r["occupation"] == "libre" else int(r["empl_AB"]) // 2
    if n == 0:
        continue
    v, p = float(r["vent_150m_moyen"]) * PRUDENCE_VENT, int(r["pluie_mm"])
    gp, kp = pacha(p)
    C.append(dict(type="crete", n=n, g=interp(VENT, v) * CALAGE_V14 + gp, k=COUT_CRETE + kp, pacha=p >= 1500, lat=float(r["lat"]), lon=float(r["lon"]), pays=r["pays"], ville=r["ville"], dv=float(r["dist_grande_ville_km"])))
for r in csv.DictReader(open(f"{D}/v21_retenues_1mw.csv", encoding="utf-8")):
    if r["type"] != "V21 Soleil" or float(r["gwh_v21"]) <= 0:
        continue
    C.append(dict(type="v21", n=PART_V21, g=float(r["gwh_v21"]), k=float(r["meur_v21"]), pacha=False, lat=float(r["lat"]), lon=float(r["lon"]), pays=r["pays"], ville=r["nom"], dv=0.0))
for r in csv.DictReader(open(f"{D}/postes_fossiles_relief_v14.csv", encoding="utf-8")):
    n = int(r["v14_possibles"])
    if n == 0 or r["heures_tracker"] == "":
        continue
    gp, kp = pacha(int(r["pluie_mm"]))
    C.append(dict(type="poste", n=n, g=interp_bas(SOLEIL, float(r["heures_tracker"])) + gp, k=COUT_SOLEIL + 10.0 + kp, pacha=int(r["pluie_mm"]) >= 1500, lat=float(r["latitude"]), lon=float(r["longitude"]), pays=r["country"], ville=r["name"], dv=float(r["dist_relief_km"])))
for r in csv.DictReader(open(f"{D}/retenues_soleil_flottant.csv", encoding="utf-8")):
    if r["pvout_kwh_kwp"] == "" or float(r["poste_mw"]) < 20:
        continue
    n = max(1, min(5, int(float(r["poste_mw"]) // 200)))
    C.append(dict(type="flottant", n=n, g=FLOTTANT[0] + FLOTTANT[1] * float(r["pvout_kwh_kwp"]) / FLOTTANT[2], k=FLOTTANT[3], pacha=False, lat=float(r["lat"]), lon=float(r["lon"]), pays=r["pays"], ville=r["nom"], dv=5.0))
c = pd.DataFrame(C)
c["iso"] = c.pays.map(iso3)
c = c[c.iso.isin(pays.index)].reset_index(drop=True)
c["twh"] = c.g / 1000.0
c["gw"] = c.twh * GW
g = pd.read_csv(f"{D}/global_power_plant_database.csv", low_memory=False)
g["fossile"] = g.primary_fuel.isin({"Coal", "Gas", "Oil", "Petcoke"})
g = g[(g.fossile & (g.capacity_mw >= 50)) | (~g.fossile & (g.capacity_mw >= 100))].reset_index(drop=True)
lib = np.where(g.fossile, 0.95 * g.capacity_mw / 1000.0, 0.0)
libre = lib.copy()
arbre = cKDTree(xyz(g.latitude.values, g.longitude.values))
P = xyz(c.lat.values, c.lon.values)
vois = arbre.query_ball_point(P, 2 * math.sin(RAYON / 6371.0 / 2))
GX = xyz(g.latitude.values, g.longitude.values)
sorte, dpo = ["ville"] * len(c), [0.0] * len(c)
for i in np.argsort((c.k / c.g).values):
    if c.type[i] == "v21":
        sorte[i] = "barrage"
        continue
    if c.type[i] == "poste":
        sorte[i] = "poste_v14"
        continue
    if c.type[i] == "flottant":
        sorte[i] = "barrage_flottant"
        continue
    vs = vois[i]
    if not vs:
        continue
    ds = 6371.0 * 2 * np.arcsin(np.clip(np.linalg.norm(GX[vs] - P[i], axis=1) / 2, 0, 1))
    besoin_gw = c.gw[i] * c.n[i]
    pris = None
    for k in np.argsort(ds):
        j = vs[k]
        if g.fossile[j] and libre[j] >= besoin_gw:
            libre[j] -= besoin_gw
            pris = ("fossile", ds[k])
            break
    if pris is None:
        pris = ("autre", float(ds.min()))
    sorte[i], dpo[i] = pris[0], pris[1] * DETOUR
c["sorte"], c["dpo"] = sorte, dpo
villes = collections.defaultdict(list)
for v in geonamescache.GeonamesCache(min_city_population=15000).get_cities().values():
    k = iso3(v["countrycode"])
    if k:
        villes[k].append((v["latitude"], v["longitude"], v["population"]))
gros = {k: [(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 k, v in villes.items()}
charge = {k: centre([a for a, b, p in v], [b for a, b, p in v], [p for a, b, p in v]) for k, v in villes.items()}
c["dgros"] = [min((km(r.lat, r.lon, a, b) for a, b in gros.get(r.iso, [])), default=0.0) for r in c.itertuples()]
def raccord(r):
    if r.sorte == "barrage":
        return r.gw * ARRIVEE
    if r.type == "poste":
        return r.gw * (r.dv * DETOUR * AC + POSTE_EX)
    if r.type == "flottant":
        return r.gw * (r.dv * AC + POSTE_EX + ARRIVEE)
    if r.sorte == "fossile":
        return r.gw * (r.dpo * AC + POSTE_EX)
    d = r.dpo if r.sorte == "autre" else r.dv
    base = r.gw * (d * AC + (POSTE_EX if r.sorte == "autre" else POSTE)) if d <= SEUIL_DOM else r.gw * (CONV + d * DETOUR * TERRE + POSTE)
    dom = r.gw * (CONV + r.dgros * DETOUR * TERRE) if r.dgros > SEUIL_DOM else 0.0
    return base + dom + r.gw * ARRIVEE
c["racc"] = c.apply(raccord, axis=1) * np.where(c.iso == "CHN", K_CHINE, 1.0)
c["cout"] = c.k + c.racc
prod = {k: centre(t.lat, t.lon, t.g * t.n) for k, t in c.groupby("iso")}
isos = list(pays.index)
ix = {k: i for i, k in enumerate(isos)}
arcs = []
for a in isos:
    if a not in prod:
        continue
    for b in isos:
        if b == a or besoin[b] < 0.05 or b not in charge:
            continue
        t, m = route(prod[a][0], prod[a][1], charge[b][0], charge[b][1])
        if t + m <= MAX_KM and m <= MAX_MER:
            arcs.append((a, b, t, m, GW * (CONV + t * TERRE + m * MER) * (K_CHINE if "CHN" in (a, b) else 1.0), 0.015 + 0.03 * (t + m) / 1000.0))
nS, nA, nH = len(c), len(arcs), len(isos)
cout = np.r_[c.cout.values, [a[4] for a in arcs], np.full(nH, HYBRIDE + GW * (ARRIVEE + 30 * AC + POSTE)), np.full(nH, 5000.0)]
A = lil_matrix((len(isos), nS + nA + 2 * nH))
for s, k in enumerate(c.iso):
    A[ix[k], s] = -c.twh[s]
for q, a in enumerate(arcs):
    A[ix[a[0]], nS + q] += 1.0
    A[ix[a[1]], nS + q] -= 1.0 - a[5]
for h in range(nH):
    A[h, nS + nA + h] = -1.0
    A[h, nS + nA + nH + h] = -1.0
b = np.array([-(besoin[k] - nonloc[k]) for k in isos])
bornes = [(0, n) for n in c.n] + [(0, None)] * nA + [(0, CAP * besoin[k]) for k in isos] + [(0, None)] * nH
res = linprog(cout, A_ub=A.tocsr(), b_ub=b, bounds=bornes, method="highs")
y, f, hy, manque = res.x[:nS], res.x[nS:nS + nA], res.x[nS + nA:nS + nA + nH], res.x[nS + nA + nH:]
c["pris"] = y
u = c[c.pris > 1e-6].copy()
twh_sites = (u.twh * u.pris).sum()
md_cent = (u.k * u.pris).sum() / 1000
md_racc = (u.racc * u.pris).sum() / 1000
md_exp = sum(a[4] * x for a, x in zip(arcs, f)) / 1000
twh_exp = f.sum()
pertes = sum(a[5] * x for a, x in zip(arcs, f))
md_hyb = (hy * HYBRIDE).sum() / 1000
md_hyb_res = (hy * GW * (ARRIVEE + 30 * AC + POSTE)).sum() / 1000
n_pacha = (u.pris * u.pacha).sum()
n_gold = FRAC_GOLD * n_pacha
n_sky = round(n_gold * 225.0 / 1126)
md_sky = n_sky * 87.2 / 1000
md_fixe = sum(v[2] for v in FIXE.values())
md_res_fixe = nonloc_twh * GW * (ARRIVEE + 30 * AC + POSTE) / 1000
total = md_cent + md_racc + md_exp + md_hyb + md_hyb_res + md_sky + md_fixe + md_res_fixe
nb = u.groupby("type").pris.sum().round().to_dict()
tw = (u.twh * u.pris).groupby(u.type).sum().round().to_dict()
print(f"manque {manque.sum():.0f} TWh", {isos[i]: round(manque[i]) for i in np.argsort(-manque)[:8] if manque[i] > 1})
print(f"statut {res.status} candidats {nS} ({c.n.sum():.0f} emplacements) arcs {nA}")
print(f"centrales {nb} TWh {tw} ; sites {twh_sites:.0f} TWh ; hybrides {hy.sum():.0f} TWh ; export {twh_exp:.0f} TWh pertes {pertes:.0f}")
print(f"Md€ centrales {md_cent:.0f} pacha+ ; raccordement {md_racc:.0f} ; export {md_exp:.0f} ; hybrides {md_hyb:.0f} + reseau {md_hyb_res:.0f} ; sky {md_sky:.0f} ; fixes {md_fixe:.0f} + reseau fixes {md_res_fixe:.0f} ; TOTAL {total:.0f}")
hydro = twh_sites - (u.twh * u.pris)[u.type.isin(["v21", "flottant"])].sum() + nonloc_twh - FIXE["Titan Claudia"][1] - FIXE["O Sky sur les Sky"][1]
soir = 1042.0 / 17094.0 * hydro
print(f"postes : {(u.pris * (u.sorte == 'fossile')).sum():.0f} centrales sur poste fossile, {(u.pris * (u.sorte == 'autre')).sum():.0f} sur poste non fossile ; Pacha {n_pacha:.0f} Gold {n_gold:.0f} Sky {n_sky} ; pilotable le soir environ {soir:.0f} GW")
R = dict(soir=soir, manque=manque.sum(), total=total, cent=md_cent, racc=md_racc, exp=md_exp, hyb=md_hyb, hyb_res=md_hyb_res, sky=md_sky, fixe=md_fixe, res_fixe=md_res_fixe, twh_exp=twh_exp, pertes=pertes, hyb_twh=hy.sum())
if ECRIRE:
    pd.Series(R).round(1).to_csv(f"{D}/plan_mondial_artesia_5_v1_totaux.csv", header=["valeur"], lineterminator="\n")
    o = u[["type", "lat", "lon", "pays", "ville", "pris", "twh", "k", "racc", "sorte", "dpo", "dv"]].copy()
    o["twh_pris"] = o.twh * o.pris
    o.round(3).to_csv(f"{D}/plan_mondial_artesia_5_v1_sites.csv", index=False, lineterminator="\n")
    out = collections.defaultdict(float)
    inn = collections.defaultdict(float)
    with open(f"{D}/plan_mondial_artesia_5_v1_lignes.csv", "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["de", "vers", "twh_par_an", "gw", "km_terre", "km_mer", "md_euros", "lat_de", "lon_de", "lat_vers", "lon_vers"])
        for a, x in sorted(zip(arcs, f), key=lambda t: -t[1]):
            if x > 0.05:
                out[a[0]] += x
                inn[a[1]] += x * (1 - a[5])
                w.writerow([pays.pays[a[0]], pays.pays[a[1]], round(x, 1), round(x * GW, 1), round(a[2]), round(a[3]), round(a[4] * x / 1000, 1), round(prod[a[0]][0], 2), round(prod[a[0]][1], 2), round(charge[a[1]][0], 2), round(charge[a[1]][1], 2)])
    pp = (u.twh * u.pris).groupby(u.iso).sum()
    nn = u.pris.groupby(u.iso).sum()
    with open(f"{D}/plan_mondial_artesia_5_v1_pays.csv", "w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["iso", "pays", "fossile_twh", "besoin_plan_twh", "centrales_du_plan", "production_sites_twh", "lignes_fixes_twh", "panneaux_hybrides_twh", "exporte_twh", "importe_twh"])
        for k in pays.sort_values("fossile_twh", ascending=False).index:
            if besoin[k] < 0.05 and pp.get(k, 0) < 0.05:
                continue
            w.writerow([k, pays.pays[k], round(F[k], 1), round(besoin[k], 1), round(nn.get(k, 0)), round(pp.get(k, 0), 1), round(nonloc[k], 1), round(hy[ix[k]], 1), round(out[k], 1), round(inn[k], 1)])