import sys, glob, math, os
import numpy as np
import pandas as pd
import rasterio
from rasterio.windows import Window
from scipy.ndimage import minimum_filter
D = sys.argv[1] if len(sys.argv) > 1 else "."
TUILES = sys.argv[2] if len(sys.argv) > 2 else "."
PV = sys.argv[3] if len(sys.argv) > 3 else "World_PVOUT_GISdata_LTAy_AvgDailyTotals_GlobalSolarAtlas-v2_GEOTIFF.zip"
PLUIE = sys.argv[4] if len(sys.argv) > 4 else "wc2.1_5m_prec.zip"
RAYON, CHUTE, VERSANT_KM, KM2_PAR_SITE, GW_PAR_V14, PRUDENCE = 50.0, 300.0, 3.0, 25.0, 0.25, 0.5
g = pd.read_csv(f"{D}/global_power_plant_database.csv", low_memory=False)
g = g[g.primary_fuel.isin(["Coal", "Gas", "Oil", "Petcoke"]) & (g.capacity_mw >= 50)].reset_index(drop=True)
tuiles = {}
for f in glob.glob(os.path.join(TUILES, "*gmted*.tif")):
    tuiles[os.path.basename(f)[:7]] = f
def nom(la, lo):
    r = int(math.floor((la + 10) / 20)) * 20 - 10
    c = int(math.floor(lo / 30)) * 30
    return f"{abs(r):02d}{'N' if r >= 0 else 'S'}{abs(c):03d}{'E' if c >= 0 else 'W'}"
ouverts = {}
pv = rasterio.open(f"/vsizip/{PV}/World_PVOUT_GISdata_LTAy_AvgDailyTotals_GlobalSolarAtlas-v2_GEOTIFF/PVOUT.tif")
pluie = [rasterio.open(f"/vsizip/{PLUIE}/wc2.1_5m_prec_{m:02d}.tif") for m in range(1, 13)]
def cumul(lon, lat):
    v = [next(b.sample([(lon, lat)]))[0] for b in pluie]
    return round(float(sum(x for x in v if x > -1000)))
def soleil(lon, lat):
    v = next(pv.sample([(lon, lat)]))[0]
    return round(float(v) * 365) if v and v > 0 else ""
rows = []
for r in g.itertuples():
    t = nom(r.latitude, r.longitude)
    out = dict(name=r.name, country=r.country, primary_fuel=r.primary_fuel, capacity_mw=r.capacity_mw, latitude=r.latitude, longitude=r.longitude, cellules_300m=0, versants=0, dist_relief_km="", pvout_kwh_kwp="", heures_tracker="", tuile=t if t in tuiles else "absente")
    if t in tuiles and t not in ouverts:
        try:
            ouverts[t] = rasterio.open(tuiles[t])
        except Exception:
            ouverts[t] = None
    if t in tuiles and ouverts[t] is None:
        out["tuile"] = "illisible"
    if t in tuiles and ouverts[t] is not None:
        src = ouverts[t]
        res = src.res[0]
        ky = RAYON / 111.32 / res
        kx = RAYON / (111.32 * max(math.cos(math.radians(r.latitude)), 0.1)) / res
        m = int(math.ceil(max(kx, ky))) + 4
        row, col = src.index(r.longitude, r.latitude)
        w = Window(col - m, row - m, 2 * m + 1, 2 * m + 1)
        z = src.read(1, window=w, boundless=True, fill_value=-32768).astype(float)
        z[z <= -1000] = np.nan
        if np.isfinite(z).sum() > 10:
            zz = np.where(np.isfinite(z), z, 9999.0)
            k = int(round(VERSANT_KM / 111.32 / res)) * 2 + 1
            drop = np.where(np.isfinite(z), z - minimum_filter(zz, size=k), 0.0)
            yy, xx = np.mgrid[-m:m + 1, -m:m + 1]
            d = np.hypot(yy * res * 111.32, xx * res * 111.32 * math.cos(math.radians(r.latitude)))
            ok = (drop >= CHUTE) & (d <= RAYON)
            n = int(ok.sum())
            out["cellules_300m"] = n
            if n:
                surf = n * (res * 111.32) ** 2 * math.cos(math.radians(r.latitude))
                out["versants"] = int(surf / KM2_PAR_SITE * PRUDENCE)
                out["dist_relief_km"] = round(float(d[ok].min()), 1)
    try:
        v = next(pv.sample([(r.longitude, r.latitude)]))[0]
        if v and v > 0:
            out["pvout_kwh_kwp"] = round(float(v) * 365, 0)
            out["heures_tracker"] = round(float(v) * 365 * 1.25, 0)
    except Exception:
        pass
    out["pluie_mm"] = cumul(r.longitude, r.latitude)
    rows.append(out)
o = pd.DataFrame(rows)
o["libere_gw"] = (0.95 * o.capacity_mw / 1000).round(3)
o["v14_possibles"] = np.minimum(o.versants, np.floor(o.libere_gw / GW_PAR_V14)).astype(int)
o.to_csv(f"{D}/postes_fossiles_relief_v14.csv", index=False, lineterminator="\n")
v = pd.read_csv(f"{D}/v21_retenues_1mw.csv")
v["pvout_kwh_kwp"] = [soleil(a, b) for a, b in zip(v.lon, v.lat)]
v[["pays", "nom", "lat", "lon", "poste_mw", "type", "pvout_kwh_kwp"]].to_csv(f"{D}/retenues_soleil_flottant.csv", index=False, lineterminator="\n")
print(len(o), "postes ;", o.tuile.isin(["absente", "illisible"]).sum(), "sans tuile ;", (o.versants > 0).sum(), "avec versant de 300 m a 50 km ;", int(o.v14_possibles.sum()), "V14 possibles")
print(o.groupby("country").v14_possibles.sum().sort_values(ascending=False).head(20).to_string())