import csv,math,sys,collections
import numpy as np
import rasterio
from scipy.spatial import cKDTree
PV=sys.argv[1] if len(sys.argv)>1 else "PVOUT.tif"
VILLES=sys.argv[2] if len(sys.argv)>2 else "cities15000.txt"
rows=list(csv.DictReader(open("plateaux_bruts.csv")))
coords=[(float(r["lon"]),float(r["lat"])) for r in rows]
with rasterio.open(PV) as src:
    pv=[v[0] for v in src.sample(coords)]
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])
def classe(h):
    return "A" if h>=2100 else "B" if h>=1800 else "C"
out=[]
for r,v in zip(rows,pv):
    la=math.radians(float(r["lat"]));lo=math.radians(float(r["lon"]))
    xyz=[math.cos(la)*math.cos(lo),math.cos(la)*math.sin(lo),math.sin(la)]
    d,i=tree.query(xyz);d2,i2=treeb.query(xyz)
    c=cities[i];cb=cities[big[i2]]
    dist=6371*2*math.asin(min(1,d/2));distb=6371*2*math.asin(min(1,d2/2))
    if not v==v: continue
    pva=float(v)*365.0
    trk=pva*1.25
    area=float(r["surface_km2"]);rim=float(r["rebord_km"])
    sites=int(min(area//7,max(1,rim//3)))
    out.append(dict(r,pvout_kwh_kwp=round(pva),heures_tracker=round(trk),classe=classe(trk),sites=sites,ville=c[2],pays=c[3],dist_km=round(dist),grande_ville=cb[2],dist_grande_ville_km=round(distb)))
with open("plateaux_ceinture_soleil.csv","w",newline="",encoding="utf-8") as f:
    w=csv.DictWriter(f,fieldnames=list(out[0].keys()));w.writeheader();w.writerows(out)
tot=collections.defaultdict(lambda:[0,0,0,0,0,0.0])
for o in out:
    if o["classe"]=="C": continue
    t=tot[o["pays"]];t[0]+=1;t[1]+=o["sites"];t[2]+=o["sites"] if o["classe"]=="A" else 0;t[3]+=o["sites"] if o["classe"]=="B" else 0;t[4]+=o["sites"] if o["dist_grande_ville_km"]<=150 else 0;t[5]+=float(o["surface_km2"])
with open("plateaux_ceinture_soleil_pays.csv","w",newline="",encoding="utf-8") as f:
    w=csv.writer(f);w.writerow(["pays","plateaux_A_B","sites_A_B","sites_A","sites_B","sites_a_moins_de_150_km_d_une_ville","surface_km2"])
    for k,t in sorted(tot.items(),key=lambda x:-x[1][1]): w.writerow([k,t[0],t[1],t[2],t[3],t[4],round(t[5])])
cand=[r for r in out if r["classe"]=="A" and r["dist_grande_ville_km"]<=150 and r["sites"]>=3]
cand.sort(key=lambda r:-(r["sites"]*r["heures_tracker"]))
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("plateaux_ceinture_soleil_top100.csv","w",newline="",encoding="utf-8") as f:
    w=csv.DictWriter(f,fieldnames=list(out[0].keys()));w.writeheader();w.writerows(top)
def S(l): return sum(r["sites"] for r in l)
AB=[r for r in out if r["classe"]!="C"]
print("plateaux",len(out),"sites",S(out),"surface",round(sum(float(r['surface_km2']) for r in out)),"km2")
for cl in "ABC": print("classe",cl,"plateaux",sum(1 for r in out if r["classe"]==cl),"sites",S([r for r in out if r["classe"]==cl]))
print("A+B : plateaux",len(AB),"sites",S(AB),"; a moins de 150 km d'une ville de 100 000 :",S([r for r in AB if r["dist_grande_ville_km"]<=150]),"dont A",S([r for r in AB if r["dist_grande_ville_km"]<=150 and r["classe"]=="A"]),"; a moins de 100 km :",S([r for r in AB if r["dist_grande_ville_km"]<=100]),"; a moins de 50 km :",S([r for r in AB if r["dist_grande_ville_km"]<=50]))
print("plateaux A+B de 20 km2 et plus",sum(1 for r in AB if float(r["surface_km2"])>=20),"; de 70 km2 et plus (10 sites)",sum(1 for r in AB if float(r["surface_km2"])>=70))
print("candidats du top 100 :",len(cand))