import glob,os,sys,math,csv,gzip,json,collections
import numpy as np
import rasterio
from scipy import ndimage
SRC=sys.argv[1] if len(sys.argv)>1 else "tiles"
VENT=sys.argv[2] if len(sys.argv)>2 else "gwa150_001deg.bin.gz"
PLAT=sys.argv[3] if len(sys.argv)>3 else "plateaux_ceinture_soleil.csv"
LATMAX=35.0
SLOPE_MAX=0.05
AREA_MIN_KM2=7.0
DROP_MIN=300.0
DROP_RADIUS_PX=3
SEUIL=6.0
def charge_vent(p):
    b=gzip.open(p).read()
    n=int.from_bytes(b[:4],"little");h=json.loads(b[4:4+n]);o=4+n;t={}
    for k in h["keys"]:
        t[k]=np.frombuffer(b,dtype=np.uint8,count=1000000,offset=o).reshape(1000,1000);o+=1000000
    return t
V=charge_vent(VENT)
def vent(lats,lons):
    R=np.floor((90.0-np.asarray(lats))*100).astype(int);C=np.floor((np.asarray(lons)+180.0)*100).astype(int)
    out=np.full(len(R),np.nan)
    for i,(r,c) in enumerate(zip(R,C)):
        a=V.get((r//1000)*100+c//1000)
        if a is not None:
            v=a[r%1000,c%1000]
            if v>0: out[i]=v/10.0
    return out
mix={(round(float(r["lat"]),3),round(float(r["lon"]),3)):1 for r in csv.DictReader(open(PLAT,encoding="utf-8")) if r["classe"]!="C"}
res={}
for f in sorted(glob.glob(os.path.join(SRC,"*gmted_mea300.tif"))):
    with rasterio.open(f) as r:
        z=r.read(1).astype(np.float32);tr=r.transform;nod=r.nodata
    z[z==nod]=np.nan
    H,W=z.shape
    lat0=tr[5];dlat=tr[4];lon0=tr[2];dlon=tr[0]
    lats=lat0+dlat*(np.arange(H)+0.5)
    keep=np.abs(lats)<=LATMAX
    if not keep.any(): continue
    pxy=111.32*abs(dlat)
    pxx=(111.32*abs(dlon)*np.cos(np.radians(lats)))[:,None]
    pxx2=np.broadcast_to(pxx,(H,W))
    zf=np.where(np.isnan(z),0.0,z)
    land=~np.isnan(z)&(z>-100)
    gy=(np.roll(zf,-1,0)-np.roll(zf,1,0))/(2*pxy*1000.0)
    gx=(np.roll(zf,-1,1)-np.roll(zf,1,1))/(2*pxx2*1000.0)
    slope=np.sqrt(gx*gx+gy*gy)
    flat=(slope<=SLOPE_MAX)&land&(z>0)
    flat[~keep,:]=False
    zmin=ndimage.minimum_filter(np.where(land,zf,1e6),size=2*DROP_RADIUS_PX+1)
    drop=zf-zmin
    rim=flat&(drop>=DROP_MIN)
    near=ndimage.binary_dilation(rim,structure=np.ones((3,3)),iterations=DROP_RADIUS_PX)
    flat=flat&near
    lab,n=ndimage.label(flat,structure=np.ones((3,3)))
    if n==0: continue
    idx=np.arange(1,n+1)
    cy,cx=zip(*ndimage.center_of_mass(flat,lab,index=idx))
    voulu={}
    for i in range(n):
        la=round(lat0+dlat*(cy[i]+0.5),3);lo=round(lon0+dlon*(cx[i]+0.5),3)
        if (la,lo) in mix: voulu[i+1]=(la,lo)
    if not voulu: continue
    rl=np.where(rim,lab,0)
    obj=ndimage.find_objects(rl)
    side=np.sqrt(pxy*pxx2)
    for k,key in voulu.items():
        sl=obj[k-1]
        if sl is None: res[key]=(np.nan,0.0,0.0,0);continue
        m=rl[sl]==k
        yy,xx=np.nonzero(m)
        yy=yy+sl[0].start;xx=xx+sl[1].start
        la=lat0+dlat*(yy+0.5);lo=lon0+dlon*(xx+0.5)
        v=vent(la,lo);s=side[yy,xx]
        ok=np.isfinite(v)
        km6=float(np.sum(s[ok&(v>=SEUIL)]))
        res[key]=(float(np.nanmean(v)) if ok.any() else np.nan,km6,float(np.sum(s)),int((~ok).sum()))
    print(os.path.basename(f)[:7],"plateaux trouves",len(voulu),flush=True)
with open("plateaux_rebord_vent.csv","w",newline="\n",encoding="utf-8") as f:
    w=csv.writer(f,lineterminator="\n");w.writerow(["lat","lon","vent_rebord_150m","rebord_km","rebord_km_6ms"])
    for (la,lo),(vm,km6,kmt,nan) in sorted(res.items()): w.writerow([la,lo,"" if vm!=vm else round(vm,2),round(kmt,1),round(km6,1)])
print("plateaux A et B",len(mix),"trouves",len(res),"avec 12 km a 6 m/s",sum(1 for v in res.values() if v[1]>=12))