#!/usr/bin/env python3
"""
Distributed DANA Flood Analysis Pipeline — Hybrid Architecture
===============================================================
Stage 1 (Local):   STAC ingestion → 30m tiling → S3 staging
Stage 2 (Lambda):  Distributed spectral-index computation per tile
Stage 3 (Local):   Mosaic assembly → high-resolution visualisation

This design overcomes the PyRun Lambda runtime's lack of geospatial libs
(stackstac/rasterio) by pre-fetching COGs locally and pushing compressed
numpy tiles to S3, then leveraging AWS Lambda's native NumPy for the
CPU-bound index math.
"""

import json
import time
from io import BytesIO
from typing import List, Tuple

import lithops
import numpy as np
import planetary_computer as pc
import pystac
import pystac_client
import stackstac
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize

# =============================================================================
# CONFIGURATION
# =============================================================================
AOI_BBOX = [-0.60, 39.25, -0.30, 39.55]
PRE_EVENT_DATE = "2024-09-26"
POST_EVENT_DATE = "2024-11-10"
STAC_API_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
COLLECTION = "sentinel-2-l2a"
RESOLUTION = 30          # strict requirement
S3_BUCKET = "pyrun-bucket-6j8ox79jz7"
TILE_GRID = (2, 2)       # (rows, cols)
ASSETS = ["B02", "B03", "B04", "B08", "B11", "B12"]

# =============================================================================
# STAGE 0 — LOCAL HELPER FUNCTIONS
# =============================================================================

def generate_tiles(bbox, n_rows, n_cols):
    """Divide AOI into an n_rows × n_cols grid of bounding boxes."""
    min_lon, min_lat, max_lon, max_lat = bbox
    lon_step = (max_lon - min_lon) / n_cols
    lat_step = (max_lat - min_lat) / n_rows
    tiles = []
    for row in range(n_rows):
        for col in range(n_cols):
            tile_bbox = [
                min_lon + col * lon_step,
                min_lat + row * lat_step,
                min_lon + (col + 1) * lon_step,
                min_lat + (row + 1) * lat_step,
            ]
            tiles.append({"tile_id": f"tile_r{row}_c{col}", "bbox": tile_bbox})
    return tiles


def search_items(bbox, datetime_range, max_retries=3):
    """Search Planetary Computer for Sentinel-2 L2A items with retries."""
    import time
    for attempt in range(max_retries):
        try:
            catalog = pystac_client.Client.open(STAC_API_URL, modifier=pc.sign_inplace)
            search = catalog.search(
                collections=[COLLECTION],
                bbox=bbox,
                datetime=datetime_range,
                limit=10,
            )
            items = list(search.items())
            items = [it for it in items if it.properties.get("s2:data_coverage", 100) > 95]
            return items
        except Exception as e:
            print(f"  Search attempt {attempt+1} failed: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)
            else:
                raise


def load_tile_array(items, bbox):
    """Load a single tile's mean composite via stackstac (local)."""
    ds = stackstac.stack(
        items,
        assets=ASSETS,
        bounds_latlon=bbox,
        resolution=RESOLUTION,
        epsg=3857,
        rescale=False,
    ) / 10000.0
    if "time" in ds.dims:
        ds = ds.mean(dim="time")
    return ds


def save_tile_to_s3(ds, bucket, key, storage):
    """Serialise an xarray Dataset tile to a compressed NPZ in S3."""
    buf = BytesIO()
    # Extract band arrays in fixed order
    band_data = {str(b): ds.sel(band=b).values for b in ASSETS}
    band_data["bbox"] = np.array(ds.attrs.get("bbox", []), dtype=np.float64)
    band_data["x_coords"] = ds.x.values
    band_data["y_coords"] = ds.y.values
    np.savez_compressed(buf, **band_data)
    buf.seek(0)
    storage.put_object(bucket=bucket, key=key, body=buf.read())


# =============================================================================
# STAGE 2 — LITHOPS WORKER (runs on AWS Lambda)
# =============================================================================

def worker_compute_indices(tile_info, bucket, storage):
    """
    Distributed worker:
      1. Downloads a pre-staged tile NPZ from S3.
      2. Computes NDVI, NDWI, NBR and their deltas via pure NumPy.
      3. Uploads results back to S3.
    """
    import time as _time
    import traceback as _tb
    import numpy as np
    from io import BytesIO

    tile_id = tile_info["tile_id"]
    s3_in_key = tile_info["s3_input_key"]
    start = _time.time()

    try:
        # -----------------------------------------------------------------
        # 1. Fetch staged data from S3
        # -----------------------------------------------------------------
        raw = storage.get_object(bucket=bucket, key=s3_in_key)
        npz = np.load(BytesIO(raw))

        pre = {b: npz[f"pre_{b}"] for b in ["B02", "B03", "B04", "B08", "B11", "B12"]}
        post = {b: npz[f"post_{b}"] for b in ["B02", "B03", "B04", "B08", "B11", "B12"]}
        x_coords = npz["x_coords"]
        y_coords = npz["y_coords"]

        # -----------------------------------------------------------------
        # 2. Spectral index math (pure NumPy)
        # -----------------------------------------------------------------
        def calc_ndvi(b04, b08):
            return (b08 - b04) / (b08 + b04)

        def calc_ndwi(b03, b08):
            return (b03 - b08) / (b03 + b08)

        def calc_nbr(b08, b12):
            return (b08 - b12) / (b08 + b12)

        pre_ndvi = calc_ndvi(pre["B04"], pre["B08"])
        post_ndvi = calc_ndvi(post["B04"], post["B08"])
        pre_ndwi = calc_ndwi(pre["B03"], pre["B08"])
        post_ndwi = calc_ndwi(post["B03"], post["B08"])
        pre_nbr = calc_nbr(pre["B08"], pre["B12"])
        post_nbr = calc_nbr(post["B08"], post["B12"])

        ndvi_delta = post_ndvi - pre_ndvi
        ndwi_delta = post_ndwi - pre_ndwi
        nbr_delta = post_nbr - pre_nbr

        # -----------------------------------------------------------------
        # 3. Persist results to S3
        # -----------------------------------------------------------------
        buf = BytesIO()
        np.savez_compressed(
            buf,
            pre_ndvi=pre_ndvi,
            post_ndvi=post_ndvi,
            pre_ndwi=pre_ndwi,
            post_ndwi=post_ndwi,
            pre_nbr=pre_nbr,
            post_nbr=post_nbr,
            ndvi_delta=ndvi_delta,
            ndwi_delta=ndwi_delta,
            nbr_delta=nbr_delta,
            x_coords=x_coords,
            y_coords=y_coords,
        )
        buf.seek(0)
        s3_out_key = f"dana_valencia_30m/{tile_id}_results.npz"
        storage.put_object(bucket=bucket, key=s3_out_key, body=buf.read())

        return {
            "tile_id": tile_id,
            "s3_out_key": s3_out_key,
            "shape": pre_ndvi.shape,
            "elapsed_sec": _time.time() - start,
        }

    except Exception as exc:
        return {
            "tile_id": tile_id,
            "error": str(exc),
            "traceback": _tb.format_exc(),
        }


# =============================================================================
# STAGE 3 — LOCAL VISUALISATION
# =============================================================================

def assemble_mosaic(successes, bucket, storage):
    """Download result tiles from S3 and stitch into xarray DataArrays."""
    tiles = {}
    for meta in successes:
        raw = storage.get_object(bucket=bucket, key=meta["s3_out_key"])
        npz = np.load(BytesIO(raw))
        tiles[meta["tile_id"]] = npz

    # Sort tiles by row/col to assemble grid
    grid = [[None for _ in range(TILE_GRID[1])] for _ in range(TILE_GRID[0])]
    for tid, npz in tiles.items():
        parts = tid.split("_")
        row = int(parts[1][1:])
        col = int(parts[2][1:])
        grid[row][col] = npz

    def concat_rows(arr_key):
        rows = []
        for r in range(TILE_GRID[0]):
            cols = [grid[r][c][arr_key] for c in range(TILE_GRID[1])]
            rows.append(np.concatenate(cols, axis=1))
        return np.concatenate(rows, axis=0)

    # Build full mosaics for the key arrays
    mosaic = {}
    for key in ["pre_ndvi", "post_ndvi", "pre_ndwi", "post_ndwi",
                "pre_nbr", "post_nbr", "ndvi_delta", "ndwi_delta", "nbr_delta"]:
        mosaic[key] = concat_rows(key)

    # Use coordinates from first tile (approximate for plotting)
    x_coords = grid[0][0]["x_coords"]
    y_coords = grid[0][0]["y_coords"]
    # Extend coordinates for full mosaic
    # This is a simplified approach; for exact georeferencing we'd use rioxarray
    full_x = np.linspace(x_coords.min(), x_coords.max(), mosaic["pre_ndvi"].shape[1])
    full_y = np.linspace(y_coords.max(), y_coords.min(), mosaic["pre_ndvi"].shape[0])

    da = {}
    for key, arr in mosaic.items():
        da[key] = xr.DataArray(arr, dims=["y", "x"], coords={"y": full_y, "x": full_x})

    return da


def generate_figures(da, output_dir="/work/DANA/output"):
    """Generate publication-quality 30m resolution comparison figures."""
    import os
    os.makedirs(output_dir, exist_ok=True)

    fig_paths = []

    # --- Figure 1: NDVI Pre / Post / Delta ---
    fig, axes = plt.subplots(1, 3, figsize=(22, 7))
    vmin, vmax = -0.2, 0.8
    im0 = axes[0].imshow(da["pre_ndvi"], vmin=vmin, vmax=vmax, cmap="RdYlGn")
    axes[0].set_title(f"PRE-DANA NDVI\n{PRE_EVENT_DATE} — 30 m/px")
    axes[0].axis("off")
    plt.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04, label="NDVI")

    im1 = axes[1].imshow(da["post_ndvi"], vmin=vmin, vmax=vmax, cmap="RdYlGn")
    axes[1].set_title(f"POST-DANA NDVI\n{POST_EVENT_DATE} — 30 m/px")
    axes[1].axis("off")
    plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04, label="NDVI")

    im2 = axes[2].imshow(da["ndvi_delta"], vmin=-0.5, vmax=0.5, cmap="RdBu_r")
    axes[2].set_title("NDVI Change (Post − Pre)\nNegative = Vegetation Loss")
    axes[2].axis("off")
    plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04, label="ΔNDVI")

    plt.suptitle("Valencia DANA Impact — Normalized Difference Vegetation Index (30 m)", fontsize=16)
    plt.tight_layout()
    p1 = f"{output_dir}/ndvi_30m_comparison.png"
    fig.savefig(p1, dpi=300, bbox_inches="tight")
    fig_paths.append(p1)
    plt.close(fig)

    # --- Figure 2: NDWI Pre / Post / Delta (Flood detection) ---
    fig, axes = plt.subplots(1, 3, figsize=(22, 7))
    vmin, vmax = -0.5, 0.5
    im0 = axes[0].imshow(da["pre_ndwi"], vmin=vmin, vmax=vmax, cmap="BrBG")
    axes[0].set_title(f"PRE-DANA NDWI\n{PRE_EVENT_DATE} — 30 m/px")
    axes[0].axis("off")
    plt.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04, label="NDWI")

    im1 = axes[1].imshow(da["post_ndwi"], vmin=vmin, vmax=vmax, cmap="BrBG")
    axes[1].set_title(f"POST-DANA NDWI\n{POST_EVENT_DATE} — 30 m/px")
    axes[1].axis("off")
    plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04, label="NDWI")

    im2 = axes[2].imshow(da["ndwi_delta"], vmin=-0.6, vmax=0.6, cmap="RdBu")
    axes[2].set_title("NDWI Change (Post − Pre)\nPositive = Water/Moisture Increase")
    axes[2].axis("off")
    plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04, label="ΔNDWI")

    plt.suptitle("Valencia DANA Impact — Normalized Difference Water Index (30 m)", fontsize=16)
    plt.tight_layout()
    p2 = f"{output_dir}/ndwi_30m_comparison.png"
    fig.savefig(p2, dpi=300, bbox_inches="tight")
    fig_paths.append(p2)
    plt.close(fig)

    # --- Figure 3: NBR Pre / Post / Delta (Burn / disturbance) ---
    fig, axes = plt.subplots(1, 3, figsize=(22, 7))
    vmin, vmax = -0.5, 0.5
    im0 = axes[0].imshow(da["pre_nbr"], vmin=vmin, vmax=vmax, cmap="RdYlGn")
    axes[0].set_title(f"PRE-DANA NBR\n{PRE_EVENT_DATE} — 30 m/px")
    axes[0].axis("off")
    plt.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04, label="NBR")

    im1 = axes[1].imshow(da["post_nbr"], vmin=vmin, vmax=vmax, cmap="RdYlGn")
    axes[1].set_title(f"POST-DANA NBR\n{POST_EVENT_DATE} — 30 m/px")
    axes[1].axis("off")
    plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04, label="NBR")

    im2 = axes[2].imshow(da["nbr_delta"], vmin=-0.5, vmax=0.5, cmap="RdBu_r")
    axes[2].set_title("NBR Change (Post − Pre)\nNegative = Soil Disturbance")
    axes[2].axis("off")
    plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04, label="ΔNBR")

    plt.suptitle("Valencia DANA Impact — Normalized Burn Ratio (30 m)", fontsize=16)
    plt.tight_layout()
    p3 = f"{output_dir}/nbr_30m_comparison.png"
    fig.savefig(p3, dpi=300, bbox_inches="tight")
    fig_paths.append(p3)
    plt.close(fig)

    return fig_paths


# =============================================================================
# MAIN ORCHESTRATOR
# =============================================================================

def run_pipeline():
    print("=" * 70)
    print("DISTRIBUTED DANA FLOOD ANALYSIS PIPELINE")
    print(f"Resolution: {RESOLUTION} m/px  |  Grid: {TILE_GRID[0]}×{TILE_GRID[1]}")
    print(f"AOI: {AOI_BBOX}")
    print("=" * 70)

    # -----------------------------------------------------------------
    # STAGE 1A — STAC search (local, lightweight)
    # -----------------------------------------------------------------
    print("\n[Stage 1A] Searching Planetary Computer STAC catalog...")
    t0 = time.time()
    pre_items = search_items(AOI_BBOX, f"{PRE_EVENT_DATE}/{PRE_EVENT_DATE}T23:59:59Z")
    post_items = search_items(AOI_BBOX, f"{POST_EVENT_DATE}/{POST_EVENT_DATE}T23:59:59Z")
    print(f"  Pre-event items : {len(pre_items)}")
    print(f"  Post-event items: {len(post_items)}")
    if not pre_items or not post_items:
        raise RuntimeError("Insufficient imagery found.")
    print(f"  Elapsed: {time.time() - t0:.2f}s")

    # -----------------------------------------------------------------
    # STAGE 1B — Spatial tiling & local load → S3 staging
    # -----------------------------------------------------------------
    print("\n[Stage 1B] Generating tiles and staging to S3...")
    tiles = generate_tiles(AOI_BBOX, *TILE_GRID)
    print(f"  Total tiles: {len(tiles)}")

    # Use Lithops storage client locally for S3 I/O
    storage = lithops.Storage()
    staging_meta = []
    for tile in tiles:
        tid = tile["tile_id"]
        tbbox = tile["bbox"]

        pre_ds = load_tile_array(pre_items, tbbox)
        post_ds = load_tile_array(post_items, tbbox)

        # Package both pre and post into one NPZ per tile
        buf = BytesIO()
        np.savez_compressed(
            buf,
            **{f"pre_{b}": pre_ds.sel(band=b).values for b in ASSETS},
            **{f"post_{b}": post_ds.sel(band=b).values for b in ASSETS},
            x_coords=pre_ds.x.values,
            y_coords=pre_ds.y.values,
        )
        buf.seek(0)
        s3_key = f"dana_valencia_30m/{tid}_input.npz"
        storage.put_object(bucket=S3_BUCKET, key=s3_key, body=buf.read())
        tile["s3_input_key"] = s3_key
        staging_meta.append({"tile_info": tile})
        print(f"  Staged {tid} → {s3_key}  shape={pre_ds.sel(band='B04').shape}")

    stage1_elapsed = time.time() - t0
    print(f"  Stage 1 total elapsed: {stage1_elapsed:.2f}s")

    # -----------------------------------------------------------------
    # STAGE 2 — Distributed Lambda computation
    # -----------------------------------------------------------------
    print("\n[Stage 2] Launching distributed workers (AWS Lambda)...")
    t1 = time.time()
    with lithops.FunctionExecutor(
        backend="aws_lambda",
        runtime_memory=1024,
        execution_timeout=300,
    ) as fexec:
        futures = fexec.map(worker_compute_indices, staging_meta, extra_args={"bucket": S3_BUCKET})
        results = fexec.get_result(throw_except=False)

    stage2_elapsed = time.time() - t1
    successes = [r for r in results if "error" not in r]
    failures = [r for r in results if "error" in r]
    print(f"  Successful: {len(successes)}  |  Failed: {len(failures)}")
    for s in successes:
        print(f"  {s['tile_id']}: shape={s['shape']}  worker_time={s['elapsed_sec']:.2f}s")
    for f in failures:
        print(f"  {f['tile_id']}: ERROR → {f['error']}")
    print(f"  Stage 2 total elapsed (wall): {stage2_elapsed:.2f}s")

    # -----------------------------------------------------------------
    # STAGE 3 — Mosaic assembly & visualisation
    # -----------------------------------------------------------------
    print("\n[Stage 3] Assembling mosaic and generating figures...")
    t2 = time.time()
    da = assemble_mosaic(successes, S3_BUCKET, storage)
    fig_paths = generate_figures(da)
    stage3_elapsed = time.time() - t2
    print(f"  Figures saved: {fig_paths}")
    print(f"  Stage 3 elapsed: {stage3_elapsed:.2f}s")

    # -----------------------------------------------------------------
    # PERFORMANCE SUMMARY
    # -----------------------------------------------------------------
    total_wall = stage1_elapsed + stage2_elapsed + stage3_elapsed
    total_worker_cpu = sum(s["elapsed_sec"] for s in successes)
    parallelism_efficiency = total_worker_cpu / stage2_elapsed if stage2_elapsed > 0 else 0

    print("\n" + "=" * 70)
    print("PERFORMANCE SUMMARY")
    print("=" * 70)
    print(f"Stage 1 (Ingest + Staging)      : {stage1_elapsed:.2f}s")
    print(f"Stage 2 (Distributed Compute)   : {stage2_elapsed:.2f}s  wall")
    print(f"  Aggregate worker CPU-time     : {total_worker_cpu:.2f}s")
    print(f"  Effective parallelism         : {parallelism_efficiency:.1f}×")
    print(f"Stage 3 (Mosaic + Viz)          : {stage3_elapsed:.2f}s")
    print(f"TOTAL PIPELINE WALL TIME        : {total_wall:.2f}s")
    print("=" * 70)

    return fig_paths, {
        "stage1_sec": stage1_elapsed,
        "stage2_sec": stage2_elapsed,
        "stage3_sec": stage3_elapsed,
        "total_sec": total_wall,
        "worker_cpu_sec": total_worker_cpu,
        "parallelism_efficiency": parallelism_efficiency,
        "tiles": len(tiles),
        "successes": len(successes),
        "failures": len(failures),
    }


if __name__ == "__main__":
    run_pipeline()
