"""Modelo multicriterio y generación de candidatos con valoración 1-5 estrellas.

Seis criterios independientes, cada uno 0-100:

  sonido    nivel estimado que llega a la casa más cercana, en dB. No es una
            puntuación abstracta: es propagación real (divergencia esférica +
            absorción atmosférica + apantallamiento del relieve + absorción del
            arbolado) partiendo de un equipo de 130 dB a 1 m. Es el criterio que
            de verdad decide si te dan el aviso.
  soledad   lejanía de edificios, núcleos y densidad de construcción alrededor.
  agua      río, arroyo o fuente cerca, pero no encima.
  arbolado  cobertura arbórea alrededor: sombra y pantalla visual sin que sea
            selva impenetrable.
  acceso    pista en condiciones cerca, llano y sin pedregal.
  clima     máxima del mes más cálido y mínima del más frío.

Las estrellas salen de la media ponderada penalizada por el eslabón más débil:
un sitio perfecto salvo que no se puede llegar no es un sitio de cuatro
estrellas. Los criterios legales no puntúan, eliminan.

Los candidatos se eligen por bloques de 20 km cogiendo el mejor punto elegible
de cada bloque. Así salen cientos de opciones repartidas por todo el país y con
un abanico natural de estrellas, en vez de sesenta variantes del mismo barranco.
"""
import csv
import pickle
import sys

import numpy as np
from shapely.geometry import Point
from shapely.strtree import STRtree

sys.path.insert(0, str(__file__.rsplit("/", 1)[0]))
import config as C

# --- eliminatorias ---
MIN_D_BUILD = 2000
MIN_D_ROAD = 400        # a menos de esto te ve cada coche que pasa
MAX_D_ACCESS = 800
MAX_SLOPE = 10.0
MIN_D_PROT = 300        # margen al borde del espacio protegido oficial
MAX_SPL = 45.0          # dB estimados en la casa más cercana. 45 es el límite
                        # nocturno típico en suelo rural: por encima, el sitio
                        # no es viable por mucho que puntúe en lo demás.
MIN_SEP = 10000         # separación mínima entre candidatos

# --- modelo acústico ---
L0 = 130.0              # dB a 1 m de un equipo grande
ALPHA_KM = 0.5          # dB/km de absorción atmosférica en frecuencias bajas
MAX_FOREST_ATT = 6.0
MAX_TERRAIN_ATT = 12.0

BLOCK = 200             # celdas = 20 km
MADRID = C.to_grid(-3.7038, 40.4168)


def load(n):
    return np.load(C.INTERIM / f"{n}.npy")


def ramp(x, lo, hi):
    """0 en lo, 100 en hi (o al revés si hi<lo)."""
    return np.clip((x - lo) / (hi - lo), 0, 1) * 100


def band(x, bad_lo, good_lo, good_hi, bad_hi):
    """100 dentro de [good_lo, good_hi], cayendo a 0 en bad_lo / bad_hi."""
    up = np.clip((x - bad_lo) / max(good_lo - bad_lo, 1e-6), 0, 1)
    dn = np.clip((bad_hi - x) / max(bad_hi - good_hi, 1e-6), 0, 1)
    return np.minimum(up, dn) * 100


def main():
    print("cargando…", flush=True)
    d_build = load("d_build")
    d_major = load("d_major")
    d_minor = load("d_minor")
    d_place = load("d_place_any")
    dens = load("dens_build_5km")
    slope = load("slope_deg")
    tpi = load("tpi_2km")
    d_water = load("d_water")
    d_spring = load("d_spring")
    d_tg = load("d_track_good")
    d_tb = load("d_track_bad")
    frac_forest = load("frac_forest")
    frac_rock = load("frac_rock")
    tmax = load("tmax")
    tmin = load("tmin")
    dem = load("dem")

    spain = load("mask_spain").astype(bool)
    prot = load("mask_protected_official").astype(bool)
    d_prot = load("d_protected_official")
    military = load("mask_military").astype(bool)
    water_a = load("mask_water").astype(bool)

    d_access = np.minimum(np.minimum(d_major, d_minor), np.minimum(d_tg, d_tb))

    # ---------- eliminatorias ----------
    d_road = np.minimum(d_major, d_minor)
    ok = (spain & ~prot & ~military & ~water_a &
          (d_build >= MIN_D_BUILD) &
          (d_road >= MIN_D_ROAD) &
          (d_access <= MAX_D_ACCESS) &
          (slope <= MAX_SLOPE) &
          (d_prot >= MIN_D_PROT) &
          np.isfinite(slope) & np.isfinite(tmax))
    # El umbral acústico se aplica después de calcular el nivel, más abajo.
    print(f"celdas elegibles: {ok.sum():,} ({ok.sum()*0.01:,.0f} km2, "
          f"{ok.sum()/spain.sum()*100:.1f}% del país)", flush=True)

    # ---------- criterio sonido ----------
    d = np.maximum(d_build, 1.0)
    att_forest = frac_forest * MAX_FOREST_ATT
    att_terrain = np.clip(-tpi / 4.0, 0, MAX_TERRAIN_ATT)
    spl = (L0 - 20.0 * np.log10(d) - ALPHA_KM * d / 1000.0
           - att_forest - att_terrain).astype(np.float32)
    # 45 dB es el límite nocturno típico en suelo rural; 30 dB ya es ruido de
    # fondo del campo de noche, o sea inaudible en la práctica.
    c_sonido = ramp(-spl, -48, -30)
    ok_before_spl = int(ok.sum())
    ok &= (spl <= MAX_SPL)
    print(f"  tras el umbral acústico de {MAX_SPL:.0f} dB: {ok.sum():,} celdas "
          f"({ok.sum()*0.01:,.0f} km2)", flush=True)

    # ---------- resto de criterios ----------
    # La distancia a carretera pesa aquí de forma explícita: sin ella el modelo
    # colocaba puntos de cinco estrellas sobre el propio asfalto, porque miraba
    # solo lo lejos que estaba la casa más cercana. Da igual que no te oigan si
    # te ve cada coche que pasa.
    c_soledad = (0.40 * ramp(d_build, 1200, 6000) +
                 0.25 * ramp(d_road, MIN_D_ROAD, 4000) +
                 0.20 * ramp(d_place, 1500, 9000) +
                 0.15 * (100 - np.clip(dens / 150.0, 0, 1) * 100))
    d_any_water = np.minimum(d_water, d_spring)
    c_agua = band(d_any_water, 0, 120, 900, 4000)
    c_arbolado = band(frac_forest * 100, 0, 20, 65, 100)
    c_acceso = (0.55 * (100 - ramp(d_tg, 0, MAX_D_ACCESS)) +
                0.25 * (100 - ramp(slope, 1.5, MAX_SLOPE)) +
                0.20 * (100 - np.clip(frac_rock * 400, 0, 100)))
    c_clima = (0.6 * (100 - ramp(tmax, 30, 40)) +
               0.4 * (100 - ramp(-tmin, 1, 8)))

    crit = {"sonido": c_sonido, "soledad": c_soledad, "agua": c_agua,
            "arbolado": c_arbolado, "acceso": c_acceso, "clima": c_clima}
    W = {"sonido": .30, "soledad": .22, "acceso": .18,
         "agua": .12, "arbolado": .10, "clima": .08}

    total = sum(W[k] * crit[k] for k in W).astype(np.float32)
    # El eslabón más débil: sin esto salen sitios de 5 estrellas a los que no se
    # puede llegar o donde te oyen desde el pueblo.
    weakest = np.minimum.reduce([crit["sonido"], crit["acceso"], crit["soledad"]])
    total = (0.75 * total + 0.25 * weakest).astype(np.float32)
    total[~ok] = -1

    # ---------- muestreo por bloques ----------
    print(f"muestreando el mejor punto de cada bloque de {BLOCK*C.RES/1000:.0f} km…",
          flush=True)
    H = (C.HEIGHT // BLOCK) * BLOCK
    Wd = (C.WIDTH // BLOCK) * BLOCK
    view = total[:H, :Wd].reshape(H // BLOCK, BLOCK, Wd // BLOCK, BLOCK)
    view = view.transpose(0, 2, 1, 3).reshape(-1, BLOCK * BLOCK)
    best = view.argmax(axis=1)
    bestval = view.max(axis=1)
    keep = np.flatnonzero(bestval > 0)
    nb_c = Wd // BLOCK
    rows_out = []
    for bi in keep:
        br, bc = divmod(int(bi), nb_c)
        lr, lc = divmod(int(best[bi]), BLOCK)
        r, c = br * BLOCK + lr, bc * BLOCK + lc
        rows_out.append((r, c, float(bestval[bi])))
    rows_out.sort(key=lambda t: -t[2])
    print(f"  {len(rows_out)} tras el muestreo por bloques", flush=True)

    # Un punto por bloque no garantiza separación: dos bloques vecinos pueden
    # elegir cada uno una celda pegada a su frontera común y quedar a 100 m.
    # Filtro voraz por score: nos quedamos con el mejor y descartamos todo lo
    # que caiga a menos de MIN_SEP.
    from scipy.spatial import cKDTree
    xs = np.array([C.X_MIN + (c + 0.5) * C.RES for _, c, _ in rows_out])
    ys = np.array([C.Y_MAX - (r + 0.5) * C.RES for r, _, _ in rows_out])
    tree = cKDTree(np.c_[xs, ys])
    dead = np.zeros(len(rows_out), dtype=bool)
    kept = []
    for i in range(len(rows_out)):        # ya vienen ordenados por score
        if dead[i]:
            continue
        kept.append(i)
        for j in tree.query_ball_point([xs[i], ys[i]], MIN_SEP):
            if j != i:
                dead[j] = True
    rows_out = [rows_out[i] for i in kept]
    print(f"  {len(rows_out)} tras exigir {MIN_SEP/1000:.0f} km de separación",
          flush=True)

    # ---------- etiquetado ----------
    with open(C.INTERIM / "area_admin_municipality.pkl", "rb") as fh:
        munis = pickle.load(fh)
    with open(C.INTERIM / "area_admin_region.pkl", "rb") as fh:
        regs = [x for x in pickle.load(fh) if x["geom"].area > 1e6]
    with open(C.INTERIM / "prot_official.pkl", "rb") as fh:
        prots = pickle.load(fh)
    t_mu, t_rg = STRtree([m["geom"] for m in munis]), STRtree([x["geom"] for x in regs])
    t_pr = STRtree([p["geom"] for p in prots])

    def stars_of(v):
        return 5 if v >= 72 else 4 if v >= 62 else 3 if v >= 52 else 2 if v >= 42 else 1

    out = []
    for i, (r, c, val) in enumerate(rows_out, 1):
        x = C.X_MIN + (c + 0.5) * C.RES
        y = C.Y_MAX - (r + 0.5) * C.RES
        lon, lat = C.to_wgs(x, y)
        p = Point(x, y)

        def look(tree, recs):
            for j in tree.query(p):
                if recs[j]["geom"].contains(p):
                    return recs[j]["name"]
            # Los límites municipales de OSM dejan alguna costura sin cubrir;
            # en ese caso nos quedamos con el municipio más cercano.
            near, nd_ = "", 1e9
            for j in tree.query(p.buffer(4000)):
                dd_ = recs[j]["geom"].distance(p)
                if dd_ < nd_:
                    nd_, near = dd_, recs[j]["name"]
            return near

        near, nd = "", 1e9
        for j in t_pr.query(p.buffer(12000)):
            dd = prots[j]["geom"].distance(p)
            if dd < nd:
                nd, near = dd, f'{prots[j]["name"]} ({prots[j]["kind"]})'

        out.append({
            "id": i, "lat": round(float(lat), 5), "lon": round(float(lon), 5),
            "estrellas": stars_of(val), "score": round(val, 1),
            "municipio": look(t_mu, munis), "comunidad": look(t_rg, regs),
            "c_sonido": round(float(c_sonido[r, c])),
            "c_soledad": round(float(c_soledad[r, c])),
            "c_agua": round(float(c_agua[r, c])),
            "c_arbolado": round(float(c_arbolado[r, c])),
            "c_acceso": round(float(c_acceso[r, c])),
            "c_clima": round(float(c_clima[r, c])),
            "db_en_casa": round(float(spl[r, c]), 1),
            "d_edificio": int(d_build[r, c]),
            "d_carretera": int(min(d_major[r, c], d_minor[r, c])),
            "d_pista": int(d_tg[r, c]),
            "d_agua": int(d_any_water[r, c]),
            "d_protegido": int(d_prot[r, c]),
            "protegido_cerca": near if nd < 12000 else "",
            "km_protegido": round(nd / 1000, 1) if nd < 1e8 else None,
            "arbolado_pct": round(float(frac_forest[r, c]) * 100),
            "roca_pct": round(float(frac_rock[r, c]) * 100),
            "pendiente": round(float(slope[r, c]), 1),
            "tpi": round(float(tpi[r, c]), 1),
            "cota": int(dem[r, c]) if np.isfinite(dem[r, c]) else 0,
            "tmax": round(float(tmax[r, c]), 1),
            "tmin": round(float(tmin[r, c]), 1),
            "km_madrid": round(float(np.hypot(x - MADRID[0], y - MADRID[1]) / 1000), 1),
            "row": r, "col": c,
        })

    # Las cifras del informe y del visor se leen de aquí, para que no se queden
    # obsoletas al cambiar un umbral del modelo.
    import json
    stats = {
        "n": len(out),
        "elegibles_km2": round(float(ok_before_spl * 0.01)),
        "elegibles_pct": round(float(ok_before_spl / spain.sum() * 100), 1),
        "acustico_km2": round(float(ok.sum() * 0.01)),
        "estrellas": {str(s): sum(1 for o in out if o["estrellas"] == s)
                      for s in range(1, 6)},
        "mejor_db": round(min(o["db_en_casa"] for o in out), 1),
        "max_d_build_km": round(float(np.nanmax(np.where(spain, d_build, np.nan))) / 1000, 1),
        "prot_pct": round(float((prot & spain).sum() / spain.sum() * 100), 1),
        "prot_km2": round(float((prot & spain).sum() * 0.01)),
        "min_d_build": MIN_D_BUILD, "min_d_road": MIN_D_ROAD,
        "max_spl": MAX_SPL, "min_sep_km": MIN_SEP // 1000,
        "max_d_access": MAX_D_ACCESS, "max_slope": MAX_SLOPE,
        "min_d_prot": MIN_D_PROT,
    }
    with open(C.OUT / "stats.json", "w") as fh:
        json.dump(stats, fh, ensure_ascii=False, indent=1)

    with open(C.OUT / "candidatos_estrellas.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(out[0].keys()))
        w.writeheader()
        w.writerows(out)
    with open(C.INTERIM / "cands.pkl", "wb") as fh:
        pickle.dump(out, fh)

    dist = {s: sum(1 for o in out if o["estrellas"] == s) for s in (5, 4, 3, 2, 1)}
    print("\nreparto de estrellas:")
    for s in (5, 4, 3, 2, 1):
        print(f"  {'★'*s:<5} {dist[s]:>4}")
    print(f"\ntop 12:")
    for o in out[:12]:
        print(f"  {'★'*o['estrellas']:<5} {o['score']:5.1f} {o['lat']:8.4f},{o['lon']:9.4f} "
              f"{o['municipio'][:22]:22s} {o['db_en_casa']:5.1f}dB "
              f"pista={o['d_pista']:>4}m agua={o['d_agua']:>5}m")


if __name__ == "__main__":
    main()