import sys, gzip, struct, json, math
import numpy as np
import pandas as pd
D = sys.argv[1] if len(sys.argv) > 1 else "."
GWA = sys.argv[2] if len(sys.argv) > 2 else "gwa150_001deg.bin.gz"
RAYON, SEUIL, KM_PAR_V14, PRUDENCE = 15.0, 6.5, 12.0, 0.9
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, GWH_REF, MEUR_V14 = 885.0 / 899.4, 885.0, 364.0
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]
f = gzip.open(GWA, "rb")
n = struct.unpack("<I", f.read(4))[0]
h = json.loads(f.read(n))
tuiles = {k: np.frombuffer(f.read(1000000), dtype=np.uint8).reshape(1000, 1000) for k in h["keys"]}
def fenetre(lat, lon, km):
    dl = km / 111.32
    dg = km / (111.32 * max(math.cos(math.radians(lat)), 0.2))
    vals = []
    for la in np.arange(lat - dl, lat + dl, 0.01):
        R = int(math.floor((90 - la) * 100))
        ty, rr = R // 1000, R % 1000
        for lo in np.arange(lon - dg, lon + dg, 0.01):
            l2 = lo - 360 if lo >= 180 else lo + 360 if lo < -180 else lo
            C = int(math.floor((l2 + 180) * 100))
            tx, cc = C // 1000, C % 1000
            A = tuiles.get(ty * 100 + tx)
            if A is None:
                continue
            q = A[rr, cc]
            if q and ((la - lat) * 111.32) ** 2 + ((lo - lon) * 111.32 * math.cos(math.radians(lat))) ** 2 <= km * km:
                vals.append(q / 10.0)
    return np.array(vals)
v = pd.read_csv(f"{D}/v21_retenues_1mw.csv")
v = v[v.type != "V21 Soleil"].reset_index(drop=True)
out = []
for r in v.itertuples():
    w = fenetre(r.lat, r.lon, RAYON)
    nmax = r.gwh_v21 / GWH_REF
    if len(w) == 0:
        out.append(dict(vent_p90="", km2_venteux=0, v14=0.0, gwh=0.0, meur=0.0, classe="sans donnée"))
        continue
    p90 = float(np.percentile(w, 90))
    venteux = int((w >= SEUIL).sum())
    nv = min(nmax, venteux / KM_PAR_V14)
    vent = float(np.mean(w[w >= SEUIL])) if venteux else p90
    g = interp(VENT, vent * PRUDENCE) * CALAGE
    classe = "A" if vent >= 7 and venteux else "B" if venteux else "C" if p90 >= 6 else "D"
    out.append(dict(vent_p90=round(p90, 2), km2_venteux=venteux, v14=round(nv, 2), gwh=round(nv * g), meur=round(nv * MEUR_V14), classe=classe))
o = pd.concat([v[["pays", "nom", "lat", "lon", "poste_mw", "gwh_v21", "meur_v21"]], pd.DataFrame(out)], axis=1)
o.to_csv(f"{D}/vent_barrages_temperes_v21.csv", index=False, lineterminator="\n")
ok = o[o.v14 > 0]
print(len(o), "barrages ;", (o.classe == "sans donnée").sum(), "sans donnée ;", o.classe.value_counts().to_dict())
print("avec vent :", len(ok), "barrages,", round(ok.v14.sum()), "V14,", round(ok.gwh.sum() / 1000), "TWh,", round(ok.meur.sum() / 1000), "Md€ ; plafond de la page :", round(o.gwh_v21.sum() / 1000), "TWh")
print(ok.groupby("pays").gwh.sum().sort_values(ascending=False).head(12).div(1000).round().to_dict())