"""
populate_km_elevations.py
Reads A and B site elevations from clipped DEM TIFs.
Falls back to GLO-30 remote VRT for any site that returns nodata from the TIF.

TIFs: D:\KinderMorgan\clipped_dems\{Path_ID}.tif
CSV:  D:\KinderMorgan\Kinder_Morgan_Network_Definition_v2.csv

Run: python populate_km_elevations.py
"""
import os
import pandas as pd
import numpy as np
import rasterio

DEM_DIR  = r"D:\KinderMorgan\clipped_dems"
CSV      = r"D:\KinderMorgan\Kinder_Morgan_Network_Definition_v2.csv"

# GLO-30 remote VRT fallback (same as Oncor pipeline)
GLO30_VRT = "/vsicurl/https://copernicus-dem-30m.s3.amazonaws.com/Copernicus_DSM_COG_10_2024_1/mosaic.vrt"

df = pd.read_csv(CSV, dtype=str)
print(f"Loaded: {len(df)} paths")

for col in ["A_Latitude","A_Longitude","b_Latitude","b_Longitude","a_raat","b_raat"]:
    df[col] = pd.to_numeric(df[col], errors="coerce")

def sample_tif(tif_path, lon, lat):
    """Sample elevation at (lon, lat) from a rasterio-readable source. Returns float or None."""
    try:
        with rasterio.open(tif_path) as src:
            vals = list(src.sample([(lon, lat)]))
            v = float(vals[0][0])
            nodata = src.nodata
            if nodata is not None and v == nodata: return None
            if v < -500 or v > 9000: return None
            return round(v, 1)
    except Exception:
        return None

def get_elev(path_id, lon, lat):
    """Try corridor TIF first, fall back to GLO-30 VRT."""
    if pd.isna(lon) or pd.isna(lat):
        return None
    tif = os.path.join(DEM_DIR, f"{path_id}.tif")
    if os.path.exists(tif):
        v = sample_tif(tif, lon, lat)
        if v is not None:
            return v
    # Fallback: GLO-30
    v = sample_tif(GLO30_VRT, lon, lat)
    return v

a_elevs = []
b_elevs = []
missing_tif = []

for i, row in df.iterrows():
    path_id = row["Path_ID"]
    tif = os.path.join(DEM_DIR, f"{path_id}.tif")
    if not os.path.exists(tif):
        missing_tif.append(path_id)

    a_e = get_elev(path_id, row["A_Longitude"], row["A_Latitude"])
    b_e = get_elev(path_id, row["b_Longitude"], row["b_Latitude"])

    a_elevs.append(a_e)
    b_elevs.append(b_e)

    if (i+1) % 20 == 0:
        print(f"  {i+1}/{len(df)} processed...")

df["a_meanelev"] = a_elevs
df["b_meanelev"] = b_elevs

# AMSL heights
df["A Site Elev AMSL (M)"]            = df["a_meanelev"]
df["B Site Elev AMSL (M)"]            = df["b_meanelev"]
df["A Site Main Ant Height AGL (M)"]  = df["a_raat"]
df["B Site Main Ant Height AGL (M)"]  = df["b_raat"]

def amsl(elev, raat):
    try:
        return round(float(elev) + float(raat), 1)
    except:
        return None

df["A Site Main Ant Height AMSL (M)"] = df.apply(lambda r: amsl(r["a_meanelev"], r["a_raat"]), axis=1)
df["B Site Main Ant Height AMSL (M)"] = df.apply(lambda r: amsl(r["b_meanelev"], r["b_raat"]), axis=1)

df.to_csv(CSV, index=False)
print(f"\nSaved: {CSV}")

filled_a = sum(1 for v in a_elevs if v is not None)
filled_b = sum(1 for v in b_elevs if v is not None)
print(f"A-side elevations: {filled_a}/{len(df)}")
print(f"B-side elevations: {filled_b}/{len(df)}")

if missing_tif:
    print(f"\nMissing TIFs ({len(missing_tif)}): {missing_tif}")

still_nan = df[df["a_meanelev"].isna() | df["b_meanelev"].isna()]
if not still_nan.empty:
    print(f"\nStill missing elevation ({len(still_nan)}):")
    print(still_nan[["path_num","Path_ID","a_meanelev","b_meanelev"]].to_string())
else:
    print("\nAll elevations populated.")

print(f"\nSample:")
print(df[["path_num","Path_ID","a_meanelev","b_meanelev",
          "A Site Main Ant Height AMSL (M)",
          "B Site Main Ant Height AMSL (M)"]].head(10).to_string())