import sys, csv, math, collections
import numpy as np
import pandas as pd
import pycountry
import geonamescache
from global_land_mask import globe
from scipy.optimize import linprog
D = sys.argv[1] if len(sys.argv) > 1 else "."
K = float(sys.argv[2]) if len(sys.argv) > 2 else 1.0
LF = float(sys.argv[3]) if len(sys.argv) > 3 else 0.5
MAX_KM = float(sys.argv[4]) if len(sys.argv) > 4 else 5000.0
MAX_MER = float(sys.argv[5]) if len(sys.argv) > 5 else 1500.0
ECRIRE = len(sys.argv) <= 2
CONV, TERRE, MER, AC, POSTE, ARRIVEE = 200.0 * K, 0.30 * K, 1.60 * K, 0.60 * K, 50.0 * K, 100.0 * K
DETOUR, SEUIL_DOM, V14 = 1.15, 300.0, 0.41
TOTAL_PLAN, NONLOC = 17360.0, 1328.0 + 405.0 + 43.0 + 420.0 + 63.0
GW = 1.0 / (8.76 * LF)
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 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))
def route(a, b, c, d):
    L = km(a, b, c, d)
    n = max(int(L / 25), 2)
    p1 = np.array([math.cos(math.radians(a)) * math.cos(math.radians(b)), math.cos(math.radians(a)) * math.sin(math.radians(b)), math.sin(math.radians(a))])
    p2 = np.array([math.cos(math.radians(c)) * math.cos(math.radians(d)), math.cos(math.radians(c)) * math.sin(math.radians(d)), math.sin(math.radians(c))])
    om = math.acos(max(-1.0, min(1.0, float(p1 @ p2))))
    if om < 1e-9:
        return 0.0, 0.0
    t = np.linspace(0, 1, n)
    p = (np.sin((1 - t) * om)[:, None] * p1 + np.sin(t * om)[:, None] * p2) / math.sin(om)
    la = np.degrees(np.arcsin(np.clip(p[:, 2], -1, 1)))
    lo = np.degrees(np.arctan2(p[:, 1], p[:, 0]))
    terre = globe.is_land(np.clip(la, -89.99, 89.99), lo)
    mer = L * (1 - terre.mean())
    return (L - mer) * DETOUR, mer
def centre(lats, lons, w):
    la, lo, w = np.radians(lats), np.radians(lons), np.asarray(w, float)
    x, y, z = (w * np.cos(la) * np.cos(lo)).sum(), (w * np.cos(la) * np.sin(lo)).sum(), (w * np.sin(la)).sum()
    return math.degrees(math.atan2(z, math.hypot(x, y))), math.degrees(math.atan2(y, x))
villes = collections.defaultdict(list)
for v in geonamescache.GeonamesCache(min_city_population=15000).get_cities().values():
    c = iso3(v["countrycode"]) if v["countrycode"] != "XK" else "OWID_KOS"
    if c:
        villes[c].append((v["latitude"], v["longitude"], v["population"]))
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()
s["twh"] = s.gwh_vendus * s.poids / 1000.0
F = pays.fossile_twh
nonloc = NONLOC * F / F.sum()
loc = s.groupby("iso").twh.sum().reindex(pays.index, fill_value=0.0)
offre = loc + nonloc
besoin = F * TOTAL_PLAN / F.sum()
local = np.minimum(offre, besoin)
surplus, deficit = offre - local, besoin - local
charge = {}
for c in pays.index:
    v = villes.get(c, [])
    if v:
        charge[c] = centre([a for a, b, p in v], [b for a, b, p in v], [p for a, b, p in v])
    else:
        charge[c] = None
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["part_locale"] = s.iso.map(lambda c: local[c] / offre[c] if offre[c] > 0 else 0.0)
s["gw_racc"] = s.twh * GW
s["dist_racc"] = s.dist_grande_ville_km.fillna(0.0)
s.loc[s.centrale == "v21", "dist_racc"] = 0.0
s["cout_racc"] = s.gw_racc * (s.dist_racc * AC + POSTE)
s.loc[s.centrale == "v21", "cout_racc"] = 0.0
def dom(r):
    g = gros.get(r.iso)
    if not g:
        return 0.0
    return min(km(r.lat, r.lon, a, b) for a, b in g)
s["dist_dom"] = s.apply(dom, axis=1)
s["dom_km"] = np.where(s.dist_dom > SEUIL_DOM, s.dist_dom * DETOUR, 0.0)
s["cout_dom"] = np.where(s.dom_km > 0, s.twh * s.part_locale * GW * (CONV + s.dom_km * TERRE), 0.0)
prod = {c: centre(g.lat, g.lon, g.twh) for c, g in s.groupby("iso")}
src = [c for c in pays.index if surplus[c] > 0.05]
dst = [c for c in pays.index if deficit[c] > 0.05 and charge[c]]
arcs = []
for i in src:
    a = prod.get(i) or charge[i]
    for j in dst:
        b = charge[j]
        t, m = route(a[0], a[1], b[0], b[1])
        if t + m <= MAX_KM and m <= MAX_MER:
            arcs.append((i, j, t, m, GW * (CONV + t * TERRE + m * MER), 0.015 + 0.03 * (t + m) / 1000.0))
P = 10000.0
n, m = len(arcs), len(src)
cvec = np.array([a[4] for a in arcs] + [P] * m)
Aeq = np.zeros((m, n + m))
for k, i in enumerate(src):
    for q, a in enumerate(arcs):
        if a[0] == i:
            Aeq[k, q] = 1.0
    Aeq[k, n + k] = 1.0
beq = np.array([surplus[i] for i in src])
Aub = np.zeros((len(dst), n + m))
for k, j in enumerate(dst):
    for q, a in enumerate(arcs):
        if a[1] == j:
            Aub[k, q] = 1.0
bub = np.array([deficit[j] for j in dst])
res = linprog(cvec, A_ub=Aub, b_ub=bub, A_eq=Aeq, b_eq=beq, bounds=(0, None), method="highs")
x = res.x[:n]
flux = [(a, x[q]) for q, a in enumerate(arcs) if x[q] > 0.01]
exp_twh = sum(f for a, f in flux)
exp_cout = sum(a[4] * f for a, f in flux)
pertes = sum(a[5] * f for a, f in flux) + (s.twh * s.dist_racc * 0.005 / 100).sum() + (s.twh * s.part_locale * np.where(s.dom_km > 0, 0.015 + 0.03 * s.dom_km / 1000, 0)).sum()
bloque = sum(surplus[i] for i in src) - exp_twh
cher = sum(f for a, f in flux if a[4] / 1000 > V14)
cher_cout = sum(a[4] * f for a, f in flux if a[4] / 1000 > V14)
arrivee = TOTAL_PLAN * GW * ARRIVEE
racc = s.cout_racc.sum() + (NONLOC - 1328.0 - 43.0) * GW * (30 * AC + POSTE)
domc = s.cout_dom.sum()
total = racc + domc + exp_cout + arrivee
pertes_cout = pertes * V14 * 1000
R = dict(racc=racc / 1000, dom=domc / 1000, exp=exp_cout / 1000, arr=arrivee / 1000, total=total / 1000, pertes_twh=pertes, pertes_md=pertes_cout / 1000, exp_twh=exp_twh, bloque=bloque, local=local.sum(), surplus=surplus.sum(), deficit=deficit.sum(), cher_twh=cher, cher_md=cher_cout / 1000, dom_twh=float((s.twh * s.part_locale)[s.dom_km > 0].sum()), arcs=len(arcs), flux=len(flux))
print(" ".join(f"{k}={v:.1f}" if isinstance(v, float) else f"{k}={v}" for k, v in R.items()))
if ECRIRE:
    with open(f"{D}/reseaux_artesia_4_lignes.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f, lineterminator="\n")
        w.writerow(["de", "vers", "twh_par_an", "gw", "km_terre", "km_mer", "md_euros", "m_euros_par_gwh", "pertes_pct", "lat_de", "lon_de", "lat_vers", "lon_vers"])
        for a, fl in sorted(flux, key=lambda t: -t[1]):
            p, q = prod.get(a[0]) or charge[a[0]], charge[a[1]]
            w.writerow([pays.pays[a[0]], pays.pays[a[1]], round(fl, 1), round(fl * GW, 1), round(a[2]), round(a[3]), round(a[4] * fl / 1000, 1), round(a[4] / 1000, 3), round(100 * a[5], 1), round(p[0], 2), round(p[1], 2), round(q[0], 2), round(q[1], 2)])
    sortant = collections.defaultdict(float)
    entrant = collections.defaultdict(float)
    for a, fl in flux:
        sortant[a[0]] += fl
        entrant[a[1]] += fl
    rc = s.groupby("iso").cout_racc.sum() / 1000
    dc = s.groupby("iso").cout_dom.sum() / 1000
    with open(f"{D}/reseaux_artesia_4_pays.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f, lineterminator="\n")
        w.writerow(["iso", "pays", "fossile_twh", "besoin_plan_twh", "offre_plan_twh", "consomme_sur_place_twh", "exporte_twh", "importe_twh", "sans_acheteur_twh", "non_couvert_twh", "raccordement_md", "lignes_interieures_md"])
        for c in pays.sort_values("fossile_twh", ascending=False).index:
            if offre[c] < 0.05 and besoin[c] < 0.05:
                continue
            w.writerow([c, pays.pays[c], round(F[c], 1), round(besoin[c], 1), round(offre[c], 1), round(local[c], 1), round(sortant[c], 1), round(entrant[c], 1), round(surplus[c] - sortant[c], 1), round(deficit[c] - entrant[c], 1), round(rc.get(c, 0.0), 1), round(dc.get(c, 0.0), 1)])