# Instalar dependencias en Colab
import sys
IN_COLAB = 'google.colab' in sys.modules

if IN_COLAB:
    !pip install -q rasterio pyproj netCDF4 xarray scipy cartopy elevation
    # GDAL normalmente ya está disponible en Colab
    print("Dependencias instaladas.")
else:
    print("Ejecutando localmente – asegúrese de tener el entorno 'tsunami-tumaco' activo.")

import os
import requests
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from pathlib import Path

import rasterio
from rasterio.crs import CRS
from rasterio.transform import from_bounds
from rasterio.warp import calculate_default_transform, reproject, Resampling
from rasterio.enums import Resampling as RIOResampling
from scipy.ndimage import gaussian_filter

# Directorio de trabajo
WORK_DIR = Path('/content') if IN_COLAB else Path('.')
WORK_DIR.mkdir(parents=True, exist_ok=True)

print(f"Directorio de trabajo: {WORK_DIR}")

# Límites del dominio en coordenadas geográficas (WGS84)
LAT_MIN =  0.5   # grados N
LAT_MAX =  3.5   # grados N
LON_MIN = -79.5  # grados W
LON_MAX = -77.0  # grados W

# Proyección de destino: UTM Zona 18N (Colombia Pacífico)
CRS_SRC  = CRS.from_epsg(4326)   # WGS84 geográfico
CRS_DEST = CRS.from_epsg(32618)  # UTM Zona 18N

# Parámetros del epicentro del terremoto de 1979
EPICENTRO_LAT =  1.598  # °N
EPICENTRO_LON = -79.359  # °W

print(f"Dominio: ({LAT_MIN}°N, {LON_MIN}°W) → ({LAT_MAX}°N, {LON_MAX}°W)")
print(f"Epicentro 1979: {EPICENTRO_LAT}°N, {EPICENTRO_LON}°W")

def download_gebco(lat_min, lat_max, lon_min, lon_max, output_path):
    """Descarga datos GEBCO 2023 para la región especificada."""
    # Endpoint de la herramienta de descarga por sub-región de GEBCO
    url = (
        f"https://download.gebco.net/a/gebco_2023"
        f"?lat1={lat_min}&lat2={lat_max}"
        f"&lon1={lon_min}&lon2={lon_max}"
        f"&format=netcdf4"
    )
    print(f"Descargando GEBCO 2023 desde:\n  {url}")
    r = requests.get(url, timeout=120)
    r.raise_for_status()
    if len(r.content) < 5000:
        raise ValueError("Respuesta demasiado pequeña – posible error en la API de GEBCO.")
    with open(output_path, 'wb') as f:
        f.write(r.content)
    print(f"  Guardado en: {output_path} ({len(r.content)/1e6:.1f} MB)")
    return output_path


def download_etopo_ncei(lat_min, lat_max, lon_min, lon_max, output_path):
    """Descarga ETOPO 2022 vía el servicio REST de NOAA NCEI (15 arc-sec)."""
    # Calcula el tamaño de la grilla en píxeles para ~450m de resolución
    # 15 arc-sec = 0.004167°; ancho / res → número de columnas
    ncols = int((lon_max - lon_min) / 0.004167)
    nrows = int((lat_max - lat_min) / 0.004167)

    url = (
        "https://gis.ngdc.noaa.gov/arcgis/rest/services/DEM_mosaics/"
        "DEM_global_mosaic/ImageServer/exportImage"
        f"?bbox={lon_min},{lat_min},{lon_max},{lat_max}"
        "&bboxSR=4326"
        f"&size={ncols},{nrows}"
        "&imageSR=4326"
        "&format=tiff"
        "&pixelType=F32"
        "&f=image"
    )
    print(f"Descargando ETOPO 2022 (NOAA NCEI):\n  Grid: {ncols}×{nrows} píxeles")
    r = requests.get(url, timeout=120)
    r.raise_for_status()
    with open(output_path, 'wb') as f:
        f.write(r.content)
    print(f"  Guardado en: {output_path} ({len(r.content)/1e6:.1f} MB)")
    return output_path


def build_synthetic_dem(lat_min, lat_max, lon_min, lon_max, output_path):
    """
    Genera un DEM sintético simplificado como fallback para pruebas.
    Aproxima la batimetría del Pacífico colombiano con perfiles lineales
    y una plataforma continental.
    ADVERTENCIA: solo para pruebas – no usar para análisis de riesgo.
    """
    print("⚠ Generando DEM SINTÉTICO de prueba (no usar para análisis real).")
    res = 0.004167  # 15 arc-sec en grados
    lons = np.arange(lon_min, lon_max, res)
    lats = np.arange(lat_min, lat_max, res)
    nrows, ncols = len(lats), len(lons)

    # Coordenadas en grados desde el extremo oeste (océano)
    x_rel = (lons - lon_min) / (lon_max - lon_min)  # 0 = océano, 1 = tierra

    # Perfil de profundidad simplificado (oeste → este):
    # Cuenca oceánica (~4000m) → plataforma (~200m) → costa (~0m) → tierra (+10m)
    shelf_start = 0.70   # fracción del dominio donde empieza la plataforma
    coast_pos   = 0.88   # fracción donde está la costa

    z_profile = np.where(
        x_rel < shelf_start,
        -4000 + 3800 * (x_rel / shelf_start),       # cuenca a plataforma
        np.where(
            x_rel < coast_pos,
            -200 + 200 * ((x_rel - shelf_start) / (coast_pos - shelf_start)),
            10.0   # tierra continental
        )
    )

    # Extender a 2D y añadir rugosidad suave
    dem = np.tile(z_profile, (nrows, 1)).astype(np.float32)
    dem += gaussian_filter(np.random.randn(nrows, ncols) * 50, sigma=5)

    # Añadir islotes de Tumaco (lat ~1.7-1.9°N, lon ~78.9-78.7°W)
    tumaco_lat_idx = np.where((lats >= 1.7) & (lats <= 1.9))[0]
    tumaco_lon_idx = np.where((lons >= -78.9) & (lons <= -78.65))[0]
    if len(tumaco_lat_idx) and len(tumaco_lon_idx):
        dem[np.ix_(tumaco_lat_idx, tumaco_lon_idx)] = np.clip(
            dem[np.ix_(tumaco_lat_idx, tumaco_lon_idx)], 1.5, 10.0
        )

    transform = from_bounds(lon_min, lat_min, lon_max, lat_max, ncols, nrows)
    with rasterio.open(
        output_path, 'w',
        driver='GTiff', height=nrows, width=ncols,
        count=1, dtype='float32',
        crs=CRS_SRC, transform=transform
    ) as dst:
        dst.write(dem, 1)

    print(f"  DEM sintético guardado en: {output_path} ({nrows}×{ncols} px)")
    return output_path

# Intentar descarga en cascada
geo_tiff_raw = WORK_DIR / 'tumaco_bathy_raw.nc'
geo_tiff_etopo = WORK_DIR / 'tumaco_bathy_etopo.tif'
geo_tiff_synth = WORK_DIR / 'tumaco_bathy_synthetic.tif'

bathy_source = None

# Intento 1: GEBCO 2023
try:
    download_gebco(LAT_MIN, LAT_MAX, LON_MIN, LON_MAX, geo_tiff_raw)
    bathy_source = 'gebco'
    print("✓ GEBCO 2023 descargado exitosamente.")
except Exception as e:
    print(f"✗ GEBCO falló: {e}")

# Intento 2: ETOPO 2022 (NOAA NCEI)
if bathy_source is None:
    try:
        download_etopo_ncei(LAT_MIN, LAT_MAX, LON_MIN, LON_MAX, geo_tiff_etopo)
        bathy_source = 'etopo'
        print("✓ ETOPO 2022 descargado exitosamente.")
    except Exception as e:
        print(f"✗ ETOPO falló: {e}")

# Intento 3: DEM sintético (fallback)
if bathy_source is None:
    build_synthetic_dem(LAT_MIN, LAT_MAX, LON_MIN, LON_MAX, geo_tiff_synth)
    bathy_source = 'synthetic'

print(f"\nFuente activa: {bathy_source}")

def load_dem_to_array(source, raw_nc, etopo_tif, synth_tif):
    """Carga el DEM desde la fuente disponible y retorna (array, crs, transform)."""
    if source == 'gebco':
        ds = xr.open_dataset(raw_nc)
        # GEBCO usa 'elevation' como variable, lat/lon como coordenadas
        varname = [v for v in ds.data_vars if 'elev' in v.lower() or 'z' in v.lower()][0]
        da = ds[varname]
        # Asegura orientación norte-arriba
        lat_dim = [d for d in da.dims if 'lat' in d.lower()][0]
        lon_dim = [d for d in da.dims if 'lon' in d.lower()][0]
        if da[lat_dim].values[0] < da[lat_dim].values[-1]:
            da = da.isel({lat_dim: slice(None, None, -1)})  # flip norte-sur
        arr = da.values.astype(np.float32)
        lats = da[lat_dim].values
        lons = da[lon_dim].values
        transform = from_bounds(lons.min(), lats.min(), lons.max(), lats.max(),
                                arr.shape[1], arr.shape[0])
        crs = CRS_SRC
    else:
        path = etopo_tif if source == 'etopo' else synth_tif
        with rasterio.open(path) as src:
            arr = src.read(1)
            transform = src.transform
            crs = src.crs

    print(f"DEM cargado: {arr.shape} | Min={arr.min():.0f}m  Max={arr.max():.0f}m")
    return arr, crs, transform


dem_geo, src_crs, src_transform = load_dem_to_array(
    bathy_source, geo_tiff_raw, geo_tiff_etopo, geo_tiff_synth
)

# Guardar DEM en GeoTIFF geográfico (WGS84) como paso intermedio
geo_path = WORK_DIR / 'tumaco_bathy_geo.tif'
nrows, ncols = dem_geo.shape

with rasterio.open(
    geo_path, 'w',
    driver='GTiff', height=nrows, width=ncols,
    count=1, dtype='float32',
    crs=CRS_SRC, transform=src_transform
) as dst:
    dst.write(dem_geo, 1)

print(f"Guardado DEM geográfico: {geo_path}")

utm_path = WORK_DIR / 'tumaco_dem_utm.tif'

with rasterio.open(geo_path) as src:
    transform_utm, width_utm, height_utm = calculate_default_transform(
        src.crs, CRS_DEST,
        src.width, src.height,
        *src.bounds
    )
    kwargs = src.meta.copy()
    kwargs.update({
        'crs': CRS_DEST,
        'transform': transform_utm,
        'width': width_utm,
        'height': height_utm,
        'dtype': 'float32'
    })

    with rasterio.open(utm_path, 'w', **kwargs) as dst:
        reproject(
            source=rasterio.band(src, 1),
            destination=rasterio.band(dst, 1),
            src_transform=src.transform,
            src_crs=src.crs,
            dst_transform=transform_utm,
            dst_crs=CRS_DEST,
            resampling=Resampling.bilinear
        )

# Verificar el resultado
with rasterio.open(utm_path) as src:
    dem_utm = src.read(1)
    gt = src.transform
    pixel_size_m = abs(gt.a)  # tamaño del píxel en metros

print(f"DEM UTM: {dem_utm.shape} píxeles")
print(f"Tamaño de píxel: {pixel_size_m:.0f} m")
print(f"Rango de elevación: {dem_utm.min():.0f} m a {dem_utm.max():.0f} m")
print(f"Archivo guardado: {utm_path}")

from pyproj import Transformer

# Convertir coordenadas del epicentro a UTM
transformer = Transformer.from_crs(4326, 32618, always_xy=True)
epi_x, epi_y = transformer.transform(EPICENTRO_LON, EPICENTRO_LAT)

with rasterio.open(utm_path) as src:
    # Convertir coordenadas UTM a índice de píxel
    row, col = src.index(epi_x, epi_y)
    epi_elev = src.read(1)[row, col]

print(f"Epicentro 1979 en UTM: X={epi_x/1000:.1f} km, Y={epi_y/1000:.1f} km")
print(f"Elevación en el epicentro: {epi_elev:.0f} m")
assert epi_elev < 0, "ERROR: el epicentro debería estar en zona marina (z < 0)"
print("✓ El epicentro está correctamente en zona marina.")

from matplotlib.colors import TwoSlopeNorm

fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Colormap combinado: azul para oceano, verde-marrón para tierra
cmap_ocean = plt.cm.Blues_r
cmap_land = plt.cm.terrain

# --- Panel izquierdo: DEM completo ---
ax = axes[0]
norm = TwoSlopeNorm(vmin=-5000, vcenter=0, vmax=500)
im = ax.imshow(
    dem_utm, cmap='RdBu', norm=norm,
    extent=[
        gt.c / 1000, (gt.c + gt.a * dem_utm.shape[1]) / 1000,
        (gt.f + gt.e * dem_utm.shape[0]) / 1000, gt.f / 1000
    ]
)
plt.colorbar(im, ax=ax, label='Elevación (m)', shrink=0.8)

# Marcar epicentro
ax.plot(epi_x / 1000, epi_y / 1000, 'r*', markersize=15,
        label=f'Epicentro 1979 ({epi_elev:.0f} m)')
ax.plot(0, 0, 'ko', markersize=1)  # placeholder para Tumaco
ax.set_xlabel('Easting UTM 18N (km)')
ax.set_ylabel('Northing UTM 18N (km)')
ax.set_title('DEM Tumaco – Dominio completo')
ax.legend(fontsize=9)

# --- Panel derecho: histograma de elevaciones ---
ax2 = axes[1]
ocean_vals = dem_utm[dem_utm < 0].flatten()
land_vals = dem_utm[dem_utm >= 0].flatten()
ax2.hist(ocean_vals, bins=80, color='steelblue', alpha=0.7, label=f'Océano (n={len(ocean_vals):,})')
ax2.hist(land_vals,  bins=40, color='sienna',    alpha=0.7, label=f'Tierra (n={len(land_vals):,})')
ax2.axvline(0, color='k', linewidth=1.5, linestyle='--', label='Nivel del mar')
ax2.set_xlabel('Elevación (m)')
ax2.set_ylabel('Número de celdas')
ax2.set_title('Distribución de elevaciones')
ax2.legend()

plt.tight_layout()
plt.savefig(WORK_DIR / 'dem_overview.png', dpi=150, bbox_inches='tight')
plt.show()
print(f"Figura guardada en: {WORK_DIR}/dem_overview.png")

utm_flipped_path = WORK_DIR / 'tumaco_dem_utm_sim.tif'

with rasterio.open(utm_path) as src:
    dem_arr = src.read(1)
    meta = src.meta.copy()
    old_transform = src.transform

# Invertir columnas (eje x = W-E → E-W)
dem_flipped = dem_arr[:, ::-1]

# Actualizar el geotransform: el origen X ahora es el extremo ESTE
# y el paso en X es negativo
new_transform = rasterio.transform.from_origin(
    west=old_transform.c + old_transform.a * dem_arr.shape[1],  # extremo este
    north=old_transform.f,
    xsize=-old_transform.a,   # paso negativo (de este a oeste)
    ysize=-old_transform.e    # paso en y (positivo hacia abajo)
)

meta['transform'] = new_transform
with rasterio.open(utm_flipped_path, 'w', **meta) as dst:
    dst.write(dem_flipped, 1)

print(f"DEM orientado para simulación: {utm_flipped_path}")
print(f"  Columna 0 (y=0): extremo ESTE (costa)")
print(f"  Columna {dem_flipped.shape[1]-1} (y=L): extremo OESTE (océano)")
print(f"  Shape: {dem_flipped.shape[0]} filas × {dem_flipped.shape[1]} columnas")

outputs = [
    utm_flipped_path,
    WORK_DIR / 'dem_overview.png'
]

print("=" * 55)
print("SALIDAS DEL CUADERNO 1")
print("=" * 55)
for f in outputs:
    size_mb = Path(f).stat().st_size / 1e6 if Path(f).exists() else 0
    print(f"  {Path(f).name:40s} {size_mb:.1f} MB")

print()
print("El archivo 'tumaco_dem_utm_sim.tif' es la entrada al Cuaderno 2.")

# Guardar parámetros del DEM para el Cuaderno 2
with rasterio.open(utm_flipped_path) as src:
    pixel_m = abs(src.transform.a)
    nrows_sim, ncols_sim = src.height, src.width

print(f"\nParámetros para el Cuaderno 2:")
print(f"  resolution = {pixel_m:.0f}  # metros por píxel")
print(f"  DEM shape  = ({nrows_sim}, {ncols_sim})")