import sys, glob, os, math
import numpy as np
import pandas as pd
import rasterio
from scipy.spatial import cKDTree
D = sys.argv[1] if len(sys.argv) > 1 else "."
TUILES = sys.argv[2] if len(sys.argv) > 2 else "."
RAYON, CHUTE, Q_MIN, Q_MAX_FIABLE, MEUR_CAPTAGE, MEUR_KM = 25.0, 300.0, 1.0, 50.0, 2.0, 0.5
tu = {os.path.basename(f)[:7]: f for f in glob.glob(os.path.join(TUILES, "*gmted*.tif"))}
ouv = {}
def alt(la, lo):
    r = int(math.floor((la + 10) / 20)) * 20 - 10
    c = int(math.floor(lo / 30)) * 30
    t = f"{abs(r):02d}{'N' if r >= 0 else 'S'}{abs(c):03d}{'E' if c >= 0 else 'W'}"
    if t not in tu:
        return None
    if t not in ouv:
        try:
            ouv[t] = rasterio.open(tu[t])
        except Exception:
            ouv[t] = None
    if ouv[t] is None:
        return None
    v = next(ouv[t].sample([(lo, la)]))[0]
    return float(v) if v > -1000 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)]
s = pd.read_csv(f"{D}/sources_karstiques_wokas.csv")
s = s[(s.debit_moyen_m3s >= Q_MIN) & s.Latitude.notna()].reset_index(drop=True)
s["fiable"] = s.debit_moyen_m3s <= Q_MAX_FIABLE
s["q_utile"] = np.minimum(s.debit_etiage_q10_m3s.fillna(0) * 0.5 + s.debit_moyen_m3s * 0.5, 10.0) * 0.5
p = pd.read_csv(f"{D}/plan_mondial_artesia_5_v1_sites.csv")
p = p[p.type.isin(["crete", "soleil", "mixte", "poste"])].reset_index(drop=True)
arbre = cKDTree(xyz(p.lat.values, p.lon.values))
rows = []
for r in s.itertuples():
    d, j = arbre.query(xyz(np.array([r.Latitude]), np.array([r.Longitude]))[0])
    km = 6371.0 * 2 * math.asin(min(1.0, d / 2))
    a_s = alt(r.Latitude, r.Longitude)
    if km > RAYON or a_s is None:
        rows.append(dict(source=r.Name, pays=r.ISO, debit_moyen_m3s=r.debit_moyen_m3s, debit_etiage_m3s=r.debit_etiage_q10_m3s, fiable=r.fiable, site_km=round(km, 1), site_type="", alt_source=a_s, alt_site="", chute_utile_m=0, gwh=0.0, meur=0.0))
        continue
    q = p.iloc[j]
    a_p = alt(q.lat, q.lon)
    bas = (a_p - CHUTE) if a_p is not None else None
    h = min(CHUTE, max(0.0, a_s - bas)) if bas is not None else 0.0
    gwh = r.q_utile * 9.81 * h * 0.88 * 8760 / 1e6 if r.fiable else 0.0
    rows.append(dict(source=r.Name, pays=r.ISO, debit_moyen_m3s=r.debit_moyen_m3s, debit_etiage_m3s=r.debit_etiage_q10_m3s, fiable=r.fiable, site_km=round(km, 1), site_type=q.type, alt_source=round(a_s), alt_site=round(a_p) if a_p else "", chute_utile_m=round(h), gwh=round(gwh, 1), meur=round(MEUR_CAPTAGE + MEUR_KM * km, 1) if gwh > 0 else 0.0))
o = pd.DataFrame(rows).sort_values("gwh", ascending=False)
o.to_csv(f"{D}/gold4_sources_artesia_5.csv", index=False, lineterminator="\n")
g = o[o.gwh > 0]
print(len(s), "sources de 1 m3/s et plus ;", (o.site_km <= RAYON).sum(), "à 25 km d'une centrale du plan ;", len(g), "utiles ;", round(g.gwh.sum()), "GWh,", round(g.meur.sum(), 1), "M€")
print(g.head(15).to_string())