import csv,math,sys,collections
import numpy as np
import rasterio
from rasterio.warp import transform
from scipy.spatial import cKDTree
VILLES=sys.argv[1] if len(sys.argv)>1 else "cities15000.txt"
WC=sys.argv[2] if len(sys.argv)>2 else "wc"
SMOD=sys.argv[3] if len(sys.argv)>3 else "ghs/GHS_SMOD_E2020_GLOBE_R2023A_54009_1000_V1_0.tif"
BUILT=sys.argv[4] if len(sys.argv)>4 else "ghs/GHS_BUILT_S_E2020_GLOBE_R2023A_54009_1000_V1_0.tif"
GWH=[(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)]
def interp(t,x):
    x=min(max(x,t[0][0]),t[-1][0])
    for (a,u),(b,v) in zip(t,t[1:]):
        if a<=x<=b: return u+(v-u)*(x-a)/(b-a)
    return t[-1][1]
rows=list(csv.DictReader(open("cretes_brutes.csv",encoding="utf-8")))
cities=[]
for l in open(VILLES,encoding="utf-8"):
    p=l.rstrip("\n").split("\t")
    cities.append((float(p[4]),float(p[5]),p[1],p[8],int(p[14]) if p[14] else 0))
clat=np.radians(np.array([c[0] for c in cities]));clon=np.radians(np.array([c[1] for c in cities]))
cxyz=np.c_[np.cos(clat)*np.cos(clon),np.cos(clat)*np.sin(clon),np.sin(clat)]
tree=cKDTree(cxyz)
big=[i for i,c in enumerate(cities) if c[4]>=100000]
treeb=cKDTree(cxyz[big])
lats=np.array([float(r["lat"]) for r in rows]);lons=np.array([float(r["lon"]) for r in rows])
la=np.radians(lats);lo=np.radians(lons)
xyz=np.c_[np.cos(la)*np.cos(lo),np.cos(la)*np.sin(lo),np.sin(la)]
d,ii=tree.query(xyz);d2,i2=treeb.query(xyz)
pluie=np.zeros(len(rows))
for m in range(1,13):
    with rasterio.open(f"{WC}/wc2.1_5m_prec_{m:02d}.tif") as src:
        a=src.read(1).astype(np.float32);nod=src.nodata;tr=src.transform
    a[a==nod]=np.nan
    for i,(x,y) in enumerate(zip(lons,lats)):
        c=int((x-tr[2])/tr[0]);r=int((y-tr[5])/tr[4])
        w=a[max(r-1,0):r+2,max(c-1,0):c+2]
        v=np.nanmean(w) if np.isfinite(w).any() else np.nan
        pluie[i]+=v if v==v else 0.0
xs,ys=transform("EPSG:4326","ESRI:54009",list(lons),list(lats))
def fenetre(src,x,y,demi):
    c=int((x-src.transform[2])/src.transform[0]);r=int((y-src.transform[5])/src.transform[4])
    return src.read(1,window=((max(r-demi,0),r+demi+1),(max(c-demi,0),c+demi+1)))
hab=np.zeros(len(rows));bati=np.zeros(len(rows))
try:
    bsrc=rasterio.open(BUILT)
except Exception:
    bsrc=None
with rasterio.open(SMOD) as s:
    for i,r in enumerate(rows):
        demi=max(3,int(math.ceil(min(float(r["longueur_km"]),60.0)/4)))
        w=fenetre(s,xs[i],ys[i],demi).astype(np.int16)
        terre=w>10;n=max(int(terre.sum()),1)
        hab[i]=100.0*((w>=12)&terre).sum()/n
        if bsrc is not None:
            wb=fenetre(bsrc,xs[i],ys[i],demi).astype(np.float32);wb[wb<0]=np.nan
            bati[i]=100.0*np.nanmean(wb)/1e6 if np.isfinite(wb).any() else 0.0
def occupation(h,b):
    if h<=20 and b<1: return "libre"
    if h<=60 and b<3: return "partage"
    return "occupe"
def classe(v):
    if v=="" : return "sans_vent"
    v=float(v)
    return "A" if v>=7 else "B" if v>=6.5 else "C" if v>=6 else "D"
def pluie_classe(p):
    return "sans_pluie" if p<75 else "sec" if p<200 else "tempere" if p<1000 else "arrose" if p<1500 else "pacha"
out=[]
for i,r in enumerate(rows):
    c=cities[ii[i]];cb=cities[big[i2[i]]]
    dist=6371*2*math.asin(min(1,d[i]/2));distb=6371*2*math.asin(min(1,d2[i]/2))
    eA=int(r["empl_A"]);eAB=int(r["empl_AB"]);eABC=int(r["empl_ABC"])
    vref=float(r["vent_150m_moyen"]) if r["vent_150m_moyen"]!="" else 0.0
    out.append(dict(r,classe=classe(r["vent_150m_moyen"]),empl_B=eAB-eA,empl_C=eABC-eAB,gwh_par_v14=round(interp(GWH,max(vref,6.5))) if eAB>0 else "",ville=c[2],pays=c[3],dist_km=round(dist),grande_ville=cb[2],dist_grande_ville_km=round(distb),pluie_mm=round(float(pluie[i])),habite_pct=round(float(hab[i])),bati_pct=round(float(bati[i]),1),occupation=occupation(hab[i],bati[i]),pluie_classe=pluie_classe(pluie[i]),pacha_possible="oui" if pluie[i]>=1500 else "non"))
with open("cretes_vent_v14.csv","w",newline="",encoding="utf-8") as f:
    w=csv.DictWriter(f,fieldnames=list(out[0].keys()));w.writeheader();w.writerows(out)
def S(l,k="empl_AB"): return sum(int(r[k]) for r in l)
def pres(r): return int(r["dist_grande_ville_km"])<=150
def prud(r,k="empl_AB"):
    n=int(r[k])
    return n if r["occupation"]=="libre" else n//2 if r["occupation"]=="partage" else 0
print("crêtes",len(out),"emplacements bruts",S(out,"emplacements"),"A",S(out,"empl_A"),"B",S(out,"empl_B"),"C",S(out,"empl_C"),"A+B",S(out))
print("crêtes par classe de vent moyen",collections.Counter(r["classe"] for r in out))
AB=[r for r in out if int(r["empl_AB"])>0]
print("A+B près d'une grande ville",S([r for r in AB if pres(r)]),"dont A",S([r for r in AB if pres(r)],"empl_A"),"à moins de 50 km",S([r for r in AB if int(r["dist_grande_ville_km"])<=50]))
for o in ("libre","partage","occupe"):
    l=[r for r in AB if r["occupation"]==o]
    print("occupation",o,"crêtes",len(l),"A+B",S(l),"près d'une ville",S([r for r in l if pres(r)]))
for p in ("sans_pluie","sec","tempere","arrose","pacha"):
    l=[r for r in AB if r["pluie_classe"]==p]
    print("pluie",p,"crêtes",len(l),"A+B",S(l),"près d'une ville",S([r for r in l if pres(r)]))
P=sum(prud(r) for r in AB);PV=sum(prud(r) for r in AB if pres(r));PVA=sum(prud(r,"empl_A") for r in AB if pres(r));PVP=sum(prud(r) for r in AB if pres(r) and int(r["pluie_mm"])>=200)
print("compte prudent A+B",P,"près d'une ville",PV,"dont A",PVA,"près d'une ville avec 200 mm et plus",PVP,"fois les 5 000 :",round(PV/5000,1))
gw=sum(prud(r)*float(r["gwh_par_v14"]) for r in AB if pres(r))
print("production des emplacements prudents près d'une ville",round(gw/1000),"TWh, moyenne",round(gw/max(PV,1)),"GWh")
ISO={}
tot=collections.defaultdict(lambda:[0,0,0,0,0,0,0])
for r in out:
    t=tot[r["pays"]];t[0]+=1 if int(r["empl_AB"])>0 else 0;t[1]+=int(r["empl_A"]);t[2]+=int(r["empl_B"]);t[3]+=int(r["empl_C"])
    if pres(r): t[4]+=int(r["empl_AB"])
    t[5]+=prud(r)
    if pres(r): t[6]+=prud(r)
with open("cretes_vent_v14_pays.csv","w",newline="",encoding="utf-8") as f:
    w=csv.writer(f);w.writerow(["pays","cretes_A_B","empl_A","empl_B","empl_C","empl_A_B_a_moins_de_150_km_d_une_ville","empl_A_B_prudents","empl_A_B_prudents_pres_d_une_ville"])
    for k,t in sorted(tot.items(),key=lambda x:-x[1][6]): w.writerow([k]+t)
cand=[r for r in AB if pres(r) and int(r["empl_AB"])>=3 and r["occupation"]!="occupe"]
cand.sort(key=lambda r:-(int(r["empl_AB"])*float(r["vent_150m_moyen"] or 0)*(1.0 if r["occupation"]=="libre" else 0.7)))
top=[];cnt=collections.Counter()
for r in cand:
    if cnt[r["pays"]]>=6: continue
    cnt[r["pays"]]+=1;top.append(r)
    if len(top)==100: break
with open("cretes_vent_v14_top100.csv","w",newline="",encoding="utf-8") as f:
    w=csv.DictWriter(f,fieldnames=list(out[0].keys()));w.writeheader();w.writerows(top)
print("candidats",len(cand),"top 100 : emplacements",S(top),"pays",len(cnt))