import csv,sys,math,json,unicodedata
import numpy as np
HYDROWASTE=sys.argv[1] if len(sys.argv)>1 else "HydroWASTE_v10.csv"
GTOPO=sys.argv[2] if len(sys.argv)>2 else "gtopo.bin"
RAYON_KM=float(sys.argv[3]) if len(sys.argv)>3 else 10.0
SEUIL_M=float(sys.argv[4]) if len(sys.argv)>4 else 100.0
RENDEMENT=0.88
CHUTE_MAX=1500.0
dem=np.memmap(GTOPO,dtype="<i2",mode="r",shape=(21600,43200))
def maille(lat,lon):
    r=min(max(int((90-lat)*120),0),21599)
    c=int(((lon+180)%360)*120)%43200
    return r,c
def energie_gwh(q_m3_j,chute_m):
    return q_m3_j*365*9.81*min(chute_m,CHUTE_MAX)*RENDEMENT/3.6e6/1e3
stations=[]
for x in csv.DictReader(open(HYDROWASTE,encoding="utf-8",errors="replace")):
    try:
        lat=float(x["LAT_WWTP"]);lon=float(x["LON_WWTP"]);q=float(x["WASTE_DIS"] or 0)
    except ValueError:
        continue
    if q<=0:
        continue
    r,c=maille(lat,lon)
    autour=np.array(dem[max(r-1,0):r+2,max(c-1,0):c+2]);terre=autour[autour>-500]
    if terre.size==0:
        continue
    alt_maille=int(dem[r,c]) if dem[r,c]>-500 else int(terre.max())
    alt_prudente=int(terre.min())
    n=int(round(RAYON_KM*1.2));nlon=int(round(n/max(math.cos(math.radians(lat)),0.2)))
    fen=np.array(dem[max(r-n,0):r+n+1,max(c-nlon,0):c+nlon+1]);terre=fen[fen>-500]
    bas=int(terre.min()) if terre.size else 0
    if (fen<=-500).any():
        bas=min(bas,0)
    chute_opt=max(alt_maille-bas,0);chute_prud=max(alt_prudente-bas,0)
    if chute_opt<SEUIL_M:
        continue
    stations.append(dict(id=x["WASTE_ID"],nom=x["WWTP_NAME"].strip(),pays=x["COUNTRY"],lat=lat,lon=lon,debit_m3_j=q,population=float(x["POP_SERVED"] or 0),alt_station=alt_maille,point_bas=bas,chute_prudente=chute_prud,chute_optimiste=chute_opt,gwh_base=energie_gwh(q,chute_prud),gwh_optimiste=energie_gwh(q,chute_opt)))
stations.sort(key=lambda s:-s["gwh_base"])
with open("stations_epuration_monde_chute_%dm_plus.csv"%int(SEUIL_M),"w",newline="",encoding="utf-8-sig") as f:
    w=csv.writer(f,delimiter=";")
    w.writerow(["id_hydrowaste","nom","pays","lat","lon","debit_m3_j","population","altitude_station_m","point_bas_%dkm_m"%int(RAYON_KM),"chute_prudente_m","chute_optimiste_m","GWh_an_base","GWh_an_optimiste"])
    for s in stations:
        w.writerow([s["id"],s["nom"],s["pays"],s["lat"],s["lon"],round(s["debit_m3_j"]),round(s["population"]),s["alt_station"],s["point_bas"],s["chute_prudente"],s["chute_optimiste"],round(s["gwh_base"],1),round(s["gwh_optimiste"],1)])
tot_base=sum(s["gwh_base"] for s in stations);tot_opt=sum(s["gwh_optimiste"] for s in stations)
tot_prud=sum(s["gwh_base"] for s in stations if s["chute_prudente"]>=150)
print("%d stations avec au moins %.0f m de chute a moins de %.0f km"%(len(stations),SEUIL_M,RAYON_KM))
print("Energie par an : prudente (chute >= 150 m) %.1f TWh, base %.1f TWh, optimiste %.1f TWh"%(tot_prud/1000,tot_base/1000,tot_opt/1000))
print("Debit cumule : %.1f km3 par an"%(sum(s["debit_m3_j"] for s in stations)*365/1e9))
pays={}
for s in stations:
    p=pays.setdefault(s["pays"],[0,0.0]);p[0]+=1;p[1]+=s["gwh_base"]
for k,v in sorted(pays.items(),key=lambda kv:-kv[1][1])[:20]:
    print("%-25s %5d stations %7.0f GWh par an"%(k,v[0],v[1]))