import sys,glob,os,math
import numpy as np
import pandas as pd
import rasterio
from rasterio.windows import Window
from scipy.ndimage import minimum_filter,maximum_filter
from scipy.spatial import cKDTree
from global_land_mask import globe
import geonamescache
import pycountry
D=sys.argv[1] if len(sys.argv)>1 else "."
TUILES=sys.argv[2] if len(sys.argv)>2 else "."
CHUTE_MIN,CHUTE_MAX,HORIZ_KM,ALT_MAX,REPLAT_M,VILLE_KM,POP_MIN,ESPACE_KM,GRAPPE_KM,POSTE_MW,POSTE_KM=1000.0,2000.0,10.0,3500.0,200.0,50.0,300000,3.0,6.0,500.0,30.0
GALERIE_INCLUSE,GALERIE_MEUR,STOCK_PLAN=3.0,15.0,2943000.0
EXCLUS=[("Grand Canyon",35.6,37.0,-114.2,-111.5),("Little Colorado",35.6,36.3,-111.9,-111.0),("Yosemite",37.5,38.2,-119.9,-119.2),("Uluru",-25.6,-25.0,130.7,131.3)]
P=120
tu={}
for f in sorted(glob.glob(os.path.join(TUILES,"*gmted*.tif"))):
    try:
        with rasterio.open(f) as s: s.read(1,window=Window(0,s.height-1,s.width,1))
        tu[os.path.basename(f)[:7]]=f
    except Exception:
        tu.setdefault(os.path.basename(f)[:7],f)
def nom(r,c):
    return f"{abs(r):02d}{'N' if r>=0 else 'S'}{abs(c):03d}{'E' if c>=0 else 'W'}"
def ouvre(n):
    if n not in tu: return None
    try:
        s=rasterio.open(tu[n]);s.read(1,window=Window(0,0,1,1));return s
    except Exception:
        return None
gc=geonamescache.GeonamesCache().get_cities()
V=pd.DataFrame([dict(ville=c["name"],pays=c["countrycode"],pop=c["population"],lat=c["latitude"],lon=c["longitude"]) for c in gc.values() if c["population"]>=POP_MIN])
def xyz(lat,lon):
    la,lo=np.radians(lat),np.radians(lon)
    return np.c_[np.cos(la)*np.cos(lo),np.cos(la)*np.sin(lo),np.sin(la)]
arbre=cKDTree(xyz(V.lat.values,V.lon.values))
G=pd.read_csv(f"{D}/global_power_plant_database.csv",low_memory=False)
G=G[G.capacity_mw>=POSTE_MW].reset_index(drop=True)
arbre_p=cKDTree(xyz(G.latitude.values,G.longitude.values))
cand=[]
etat={}
for n in sorted(tu):
    r0=int(n[:2])*(1 if n[2]=="N" else -1)
    c0=int(n[3:6])*(1 if n[6]=="E" else -1)
    src=ouvre(n)
    if src is None:
        etat[n]="illisible";continue
    H,W=src.shape
    res=src.res[0]
    top=src.bounds.top
    left=src.bounds.left
    Z=np.full((H+2*P,W+2*P),np.nan,dtype=np.float32)
    z=src.read(1).astype(np.float32)
    z[z<=-1000]=np.nan
    Z[P:P+H,P:P+W]=z
    for dr in (-1,0,1):
        for dc in (-1,0,1):
            if dr==0 and dc==0: continue
            nr,nc=r0+20*dr,c0+30*dc
            if nc>=180: nc-=360
            if nc<-180: nc+=360
            s2=ouvre(nom(nr,nc))
            if s2 is None: continue
            rows=slice(0,P) if dr==1 else (slice(P+H,P+H+P) if dr==-1 else slice(P,P+H))
            cols=slice(0,P) if dc==-1 else (slice(P+W,P+W+P) if dc==1 else slice(P,P+W))
            wr=(s2.shape[0]-P if dr==1 else (0 if dr==-1 else 0))
            wc=(s2.shape[1]-P if dc==-1 else (0 if dc==1 else 0))
            hh=rows.stop-rows.start;ww=cols.stop-cols.start
            a=s2.read(1,window=Window(wc,wr,ww,hh),boundless=True,fill_value=-32768).astype(np.float32)
            a[a<=-1000]=np.nan
            Z[rows,cols]=a
    latc=(top-(np.arange(H+2*P)-P+0.5)*res)
    zz=np.where(np.isfinite(Z),Z,9999.0)
    zm=np.where(np.isfinite(Z),Z,-9999.0)
    k=int(round(HORIZ_KM/111.32/res))*2+1
    bas=minimum_filter(zz,size=(k,k))
    loc=maximum_filter(zm,size=3)-minimum_filter(zz,size=3)
    chute=Z-bas
    ok=np.isfinite(Z)&(chute>=CHUTE_MIN)&(Z<=ALT_MAX)&(loc<=REPLAT_M)
    ok[:P,:]=False;ok[-P:,:]=False;ok[:,:P]=False;ok[:,-P:]=False
    iy,ix=np.nonzero(ok)
    etat[n]=f"{len(iy)} cellules"
    if len(iy)==0: continue
    lat=top-(iy-P+0.5)*res
    lon=left+(ix-P+0.5)*res
    d,j=arbre.query(xyz(lat,lon))
    km=6371.0*2*np.arcsin(np.minimum(1.0,d/2))
    dp,jp=arbre_p.query(xyz(lat,lon))
    kp=6371.0*2*np.arcsin(np.minimum(1.0,dp/2))
    m=(km<=VILLE_KM)|(kp<=POSTE_KM)
    for y,x,la,lo,jj,kk,jpp,kpp in zip(iy[m],ix[m],lat[m],lon[m],j[m],km[m],jp[m],kp[m]):
        if any(a<=la<=b and c<=lo<=e for _,a,b,c,e in EXCLUS): continue
        h=k//2
        cx=math.cos(math.radians(la))
        win=Z[y-h:y+h+1,x-h:x+h+1]
        yy,xx=np.mgrid[-h:h+1,-h:h+1]
        dist=np.hypot(yy*res*111.32,xx*res*111.32*cx)
        lw=loc[y-h:y+h+1,x-h:x+h+1]
        zt=float(Z[y,x])
        okb=np.isfinite(win)&(dist<=HORIZ_KM)&(zt-win>=CHUTE_MIN)&((lw<=REPLAT_M)|(win<=20))
        if not okb.any(): continue
        dd=np.where(okb,dist,1e9)
        b=np.unravel_index(np.argmin(dd),dd.shape)
        zb=float(win[b])
        lab=la-(b[0]-h)*res;lob=lo+(b[1]-h)*res
        mer=bool(zb<=20 and (globe.is_ocean(lab,lob) or globe.is_ocean(lab+res,lob) or globe.is_ocean(lab-res,lob) or globe.is_ocean(lab,lob+res) or globe.is_ocean(lab,lob-res)))
        if zb<=20 and not mer and float(lw[b])>REPLAT_M: continue
        cand.append(dict(lat=round(la,4),lon=round(lo,4),alt_haut=round(float(Z[y,x])),alt_bas=round(zb),chute=round(min(float(Z[y,x])-zb,CHUTE_MAX)),dist_bas_km=round(float(dist[b]),1),cote=mer,ville=V.ville.iat[jj],pays=V.pays.iat[jj],pop_ville=int(V["pop"].iat[jj]),dist_ville_km=round(float(kk),1),poste=G.name.iat[jpp],poste_mw=round(float(G.capacity_mw.iat[jpp])),pays_poste=G.country.iat[jpp],dist_poste_km=round(float(kpp),1)))
    print(n,etat[n],len(cand),flush=True)
c=pd.DataFrame(cand)
c=c.sort_values(["dist_bas_km","dist_ville_km"],ascending=[True,True]).reset_index(drop=True)
t=cKDTree(xyz(c.lat.values,c.lon.values))
pris=np.zeros(len(c),bool);bloque=np.zeros(len(c),bool)
rr=ESPACE_KM/6371.0
for i in range(len(c)):
    if bloque[i]: continue
    pris[i]=True
    for q in t.query_ball_point(xyz(np.array([c.lat[i]]),np.array([c.lon[i]]))[0],rr):
        bloque[q]=True
e=c[pris].reset_index(drop=True)
t2=cKDTree(xyz(e.lat.values,e.lon.values))
parent=list(range(len(e)))
def f(a):
    while parent[a]!=a:
        parent[a]=parent[parent[a]];a=parent[a]
    return a
for a,b in t2.query_pairs(GRAPPE_KM/6371.0):
    parent[f(a)]=f(b)
e["rempart"]=[f(i) for i in range(len(e))]
def iso2(r):
    if r.dist_ville_km<=VILLE_KM: return r.pays
    c=pycountry.countries.get(alpha_3=r.pays_poste)
    return c.alpha_2 if c else r.pays
e["pays"]=[iso2(r) for r in e.itertuples()]
CH=[(1000.0,473.0,657.0),(1500.0,650.0,986.0),(2000.0,815.0,1314.0)]
def interp(x,i):
    if x<=CH[0][0]: return CH[0][i]
    for a,b in zip(CH,CH[1:]):
        if x<=b[0]: return a[i]+(b[i]-a[i])*(x-a[0])/(b[0]-a[0])
    return CH[-1][i]
e["galerie_km"]=np.maximum(0.0,e.dist_bas_km-GALERIE_INCLUSE).round(1)
e["meur"]=[round(interp(x,1)+GALERIE_MEUR*g,1) for x,g in zip(e.chute,e.galerie_km)]
e["gwh_pointe"]=[round(interp(x,2)) for x in e.chute]
e["stockage_mwh"]=(800000*e.chute*9.81*0.9/3.6e6).round()
g=e.groupby("rempart")
e["v15_rempart"]=e.rempart.map(g.size())
e["lieu"]=[r.ville if r.dist_ville_km<=VILLE_KM else r.poste for r in e.itertuples()]
o=(e.meur/e.gwh_pointe).sort_values().index
e["plan"]=False
e.loc[o[:int((e.stockage_mwh[o].cumsum()<STOCK_PLAN).sum())+1],"plan"]=True
e.to_csv(f"{D}/remparts_v15.csv",index=False,lineterminator="\n")
R=g.agg(lat=("lat","mean"),lon=("lon","mean"),v15=("lat","size"),chute_moy=("chute","mean"),chute_max=("chute","max"),cote=("cote","max"),lieu=("lieu","first"),pays=("pays","first"),plan=("plan","sum"),stockage_mwh=("stockage_mwh","sum"),meur=("meur","sum"),pop_ville=("pop_ville","max"),dist_ville_km=("dist_ville_km","min")).reset_index(drop=True)
R=R.sort_values(["v15","chute_moy"],ascending=False)
R.round(2).to_csv(f"{D}/remparts_v15_grappes.csv",index=False,lineterminator="\n")
pd.Series(etat).to_csv(f"{D}/remparts_v15_tuiles.csv",header=["etat"],lineterminator="\n")
print(len(c),"cellules,",len(e),"emplacements,",len(R),"remparts")