main.py

"""
Forest Fire Simulation Pipeline — main entry point.

Usage:
  python main.py --step all             # run full pipeline end-to-end
  python main.py --step download        # download all raw data
  python main.py --step preprocess      # align & build patches
  python main.py --step train           # train U-Net
  python main.py --step predict         # predict tomorrow's fire probability
  python main.py --step simulate        # run CA spread simulation
  python main.py --step visualize       # generate plots and animation

  python main.py --step predict --date 2024-05-15   # predict for a specific date
  python main.py --step simulate --ignition-threshold 3   # only 'high' class ignites

Environment variables (or .env file):
  FIRMS_MAP_KEY        – NASA FIRMS API key
  OPENTOPO_API_KEY     – OpenTopography API key
  (ERA-5 key lives in ~/.cdsapirc)
"""

import argparse
import io
import json
import os
import ssl
import sys
from pathlib import Path

# Force UTF-8 output on Windows so Unicode log characters don't crash
if hasattr(sys.stdout, "buffer"):
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")

import yaml
from dotenv import load_dotenv

load_dotenv()

# Fix SSL certificate verification on Windows (Microsoft Store Python)
try:
    import certifi
    os.environ.setdefault("REQUESTS_CA_BUNDLE", certifi.where())
    os.environ.setdefault("SSL_CERT_FILE", certifi.where())
except ImportError:
    ssl._create_default_https_context = ssl._create_unverified_context

# ---------------------------------------------------------------------------
# Config loading
# ---------------------------------------------------------------------------

def load_config(path: str = "config/config.yaml") -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


# ---------------------------------------------------------------------------
# Pipeline steps
# ---------------------------------------------------------------------------

def step_download(cfg: dict) -> dict:
    from src.data_pipeline.download_dem import download_dem
    from src.data_pipeline.download_lulc import download_lulc
    from src.data_pipeline.download_viirs import download_viirs
    from src.data_pipeline.download_era5 import download_era5

    bbox = tuple(cfg["region"]["bbox"])
    epsg = cfg["region"]["epsg_utm"]
    res = cfg["data"]["target_resolution"]
    raw = cfg["data"]["raw_dir"]

    print("\n[1/4] Downloading DEM …")
    dem_paths = download_dem(bbox, epsg, res, raw)

    print("\n[2/4] Downloading LULC (ESA WorldCover) …")
    lulc_paths = download_lulc(bbox, epsg, res, raw, dem_paths["dem"],
                               cfg["data"]["lulc"].get("fuel_map"))

    print("\n[3/4] Downloading VIIRS fire detections …")
    viirs_paths = download_viirs(
        years=cfg["data"]["viirs"]["years"],
        bbox=bbox,
        epsg_utm=epsg,
        resolution_m=res,
        raw_dir=raw,
        confidence_min=cfg["data"]["viirs"]["confidence_min"],
        firms_map_key=os.environ.get("FIRMS_MAP_KEY"),
    )

    print("\n[4/4] Downloading ERA-5 weather data …")
    era5_records = download_era5(
        years=cfg["data"]["era5"]["years"],
        months=cfg["data"]["era5"]["months"],
        bbox=bbox,
        raw_dir=raw,
    )

    return {
        "dem": dem_paths,
        "lulc": lulc_paths,
        "viirs": viirs_paths,
        "era5_records": era5_records,
    }


def step_preprocess(cfg: dict) -> dict:
    from src.data_pipeline.preprocess import run_preprocessing

    proc_dir = Path(cfg["data"]["processed_dir"])
    weather_base = proc_dir / "weather_tifs"
    if not weather_base.exists():
        print("Weather TIFs not found — run --step download first.")
        sys.exit(1)

    day_dirs = sorted(weather_base.iterdir())
    label_raw = str(proc_dir / "viirs_labels.tif")
    if not Path(label_raw).exists():
        label_raw = str(Path(cfg["data"]["raw_dir"]) / "viirs" / "viirs_labels.tif")

    result = run_preprocessing(cfg, day_dirs, label_raw)
    print(f"\nPreprocessing complete: {result['total_patches']} patches")
    return result


def step_train(cfg: dict) -> str:
    from src.models.train import train

    patch_dir = Path(cfg["data"]["processed_dir"]) / "patches"
    if not (patch_dir / "images").exists():
        print("Patches not found — run --step preprocess first.")
        sys.exit(1)

    print("\nStarting U-Net training …")
    best_ckpt = train(cfg)
    print(f"Best checkpoint: {best_ckpt}")
    return best_ckpt


def step_predict(cfg: dict, date_str: str = "latest") -> dict:
    from src.data_pipeline.preprocess import build_inference_stack
    from src.models.predict import predict_fire_probability

    proc_dir = Path(cfg["data"]["processed_dir"])

    # Find weather day dir
    if date_str == "latest":
        weather_dirs = sorted((proc_dir / "weather_tifs").iterdir())
        if not weather_dirs:
            print("No weather data — run --step download first.")
            sys.exit(1)
        day_dir = weather_dirs[-1]
        date_str = day_dir.name
    else:
        day_dir = proc_dir / "weather_tifs" / date_str
        if not day_dir.exists():
            print(f"Weather data for {date_str} not found.")
            sys.exit(1)

    print(f"\nBuilding inference feature stack for {date_str} …")
    stack, transform, crs = build_inference_stack(day_dir, proc_dir)

    ckpt = str(Path(cfg["model"]["checkpoint_dir"]) / "best.pth")
    if not Path(ckpt).exists():
        print("No checkpoint found — run --step train first.")
        sys.exit(1)

    result = predict_fire_probability(
        checkpoint_path=ckpt,
        feature_stack=stack,
        transform=transform,
        crs=crs,
        cfg=cfg,
        output_dir=cfg["output"]["prediction_dir"],
        date_str=date_str,
    )
    print(f"Probability raster → {result['probability_raster']}")
    print(f"Class raster       → {result['class_raster']}")
    return result


def step_simulate(cfg: dict, prediction_result: dict = None, ignition_threshold: int = 3) -> dict:
    import os
    os.environ.setdefault("GDAL_CACHEMAX", "64")
    import rasterio
    from rasterio.enums import Resampling as RsEnum
    import numpy as np
    from src.simulation.cellular_automata import CAGrid, simulate, make_ignition_mask

    proc_dir  = Path(cfg["data"]["processed_dir"])
    r300_dir  = proc_dir / "rasters_300m"

    def _load(path, dtype=np.float32):
        with rasterio.open(path) as s:
            return s.read(1).astype(dtype)

    def _meta(path):
        with rasterio.open(path) as s:
            return s.transform, s.crs

    print("\nLoading pre-resampled 300 m layers …")
    fuel   = _load(r300_dir / "fuel_300m.tif")
    slope  = _load(r300_dir / "slope_300m.tif")
    aspect = _load(r300_dir / "aspect_300m.tif")
    lulc   = _load(r300_dir / "lulc_300m.tif", dtype=np.int32)
    transform, crs = _meta(r300_dir / "lulc_300m.tif")
    print(f"Grid: {fuel.shape[0]}x{fuel.shape[1]}  ({fuel.nbytes/1e6:.1f} MB per layer)")

    # Probability / class maps — also at 300 m (output of make_prediction.py)
    if prediction_result is not None:
        prob_map  = prediction_result["prob_map"]
        class_map = prediction_result["class_map"]
    else:
        pred_dir   = Path(cfg["output"]["prediction_dir"])
        prob_files = sorted(pred_dir.glob("fire_probability_*.tif"))
        if not prob_files:
            print("No prediction found — run make_prediction.py first.")
            sys.exit(1)
        with rasterio.open(str(prob_files[-1])) as s:
            ph, pw = max(1, s.height // 10), max(1, s.width // 10)
            prob_map = s.read(1, out_shape=(ph, pw),
                              resampling=RsEnum.average).astype(np.float32)
        cls_files = sorted(pred_dir.glob("fire_class_*.tif"))
        with rasterio.open(str(cls_files[-1])) as s:
            class_map = s.read(1, out_shape=(ph, pw),
                               resampling=RsEnum.nearest).astype(np.uint8)

    # Use latest ERA-5 wind/humidity aligned to grid
    weather_aligned = proc_dir / "weather_aligned"
    day_dirs = sorted(weather_aligned.iterdir()) if weather_aligned.exists() else []
    if day_dirs:
        latest = day_dirs[-1]
        def _load_w(var):
            p = latest / f"{var}.tif"
            if p.exists():
                with rasterio.open(p) as s:
                    arr = s.read(1).astype(np.float32)
                    # Resize if needed
                    if arr.shape != fuel.shape:
                        from scipy.ndimage import zoom
                        arr = zoom(arr, (fuel.shape[0]/arr.shape[0], fuel.shape[1]/arr.shape[1]), order=1)
                return arr
            return np.full(fuel.shape, v, dtype=np.float32)
        wind_speed  = _load_w("wind_speed")
        wind_dir    = _load_w("wind_dir")
        humidity    = _load_w("humidity")
    else:
        # Reasonable defaults for Uttarakhand fire season (May)
        wind_speed = np.full(fuel.shape, 5.0,  dtype=np.float32)
        wind_dir   = np.full(fuel.shape, 270.0, dtype=np.float32)  # westerly
        humidity   = np.full(fuel.shape, 25.0,  dtype=np.float32)

    print(f"Creating ignition mask (class ≥ {ignition_threshold}) …")
    ignition_mask = make_ignition_mask(prob_map, class_map, threshold_class=ignition_threshold)
    n_seeds = int(ignition_mask.sum())
    print(f"Ignition seeds: {n_seeds} cells ({n_seeds * 900 / 10000:.1f} ha)")

    if n_seeds == 0:
        print("No ignition points found. Lower --ignition-threshold or run prediction again.")
        sys.exit(1)

    grid = CAGrid.from_layers(
        fuel=fuel, slope=slope, aspect=aspect,
        wind_speed=wind_speed, wind_dir=wind_dir, humidity=humidity,
        lulc=lulc, ignition_mask=ignition_mask, cfg=cfg,
    )

    print("\nRunning CA fire spread simulation …")
    results = simulate(
        grid=grid,
        time_steps_hours=cfg["simulation"]["time_steps_hours"],
        step_minutes=cfg["simulation"]["step_minutes"],
        transform=transform,
        crs=crs,
        output_dir=cfg["output"]["simulation_dir"],
    )

    return {
        "simulation_results": results,
        "prob_map": prob_map,
        "ignition_mask": ignition_mask,
        "transform": transform,
        "crs": crs,
    }


def step_visualize(cfg: dict, prediction_result: dict = None, simulation_data: dict = None):
    import json
    import rasterio
    import numpy as np
    from src.simulation.visualize import (
        plot_probability_map,
        plot_spread_snapshot,
        create_spread_animation,
        plot_spread_statistics,
        plot_training_history,
    )

    pred_dir  = Path(cfg["output"]["prediction_dir"])
    sim_dir   = Path(cfg["output"]["simulation_dir"])
    anim_dir  = Path(cfg["output"]["animation_dir"])
    ckpt_dir  = Path(cfg["model"]["checkpoint_dir"])

    # Probability map — load at 300 m to stay within RAM
    VIZ_SCALE = 10
    if prediction_result:
        prob_map = prediction_result["prob_map"]
    else:
        prob_files = sorted(pred_dir.glob("fire_probability_*.tif"))
        if prob_files:
            from rasterio.enums import Resampling as RsEnum
            with rasterio.open(str(prob_files[-1])) as s:
                h = max(1, s.height // VIZ_SCALE)
                w = max(1, s.width  // VIZ_SCALE)
                prob_map = s.read(1, out_shape=(h, w),
                                  resampling=RsEnum.average).astype(np.float32)
        else:
            prob_map = None

    if prob_map is not None:
        plot_probability_map(
            prob_map,
            str(anim_dir / "fire_probability_map.png"),
            title="Uttarakhand Fire Probability",
            bbox_label="Uttarakhand, India",
        )

    # Spread snapshots + animation
    if simulation_data:
        sim_results = simulation_data["simulation_results"]
        ign_mask    = simulation_data.get("ignition_mask")
    else:
        sim_tifs = sorted(sim_dir.glob("fire_spread_*.tif"))
        if not sim_tifs:
            print("No simulation output found — run --step simulate first.")
            return
        sim_results = {}
        for p in sim_tifs:
            hours = float(p.stem.replace("fire_spread_", "").replace("h", ""))
            with rasterio.open(str(p)) as s:
                state    = s.read(1).astype(np.int8)
                cell_m2  = abs(s.transform.a * s.transform.e)
            burned  = int((state == 2).sum())
            burning = int((state == 1).sum())
            sim_results[hours] = {
                "state": state,
                "burned_cells": burned,
                "burning_cells": burning,
                "affected_ha": (burned + burning) * cell_m2 / 10000,
            }
        ign_mask = None

    for hours, data in sim_results.items():
        plot_spread_snapshot(
            state=data["state"],
            hours=hours,
            output_path=str(anim_dir / f"spread_{hours:.1f}h.png"),
            prob_map=prob_map,
            ignition_mask=ign_mask,
        )

    create_spread_animation(
        sim_results,
        str(anim_dir / "fire_spread_animation.gif"),
        fps=2,
        prob_map=prob_map,
        ignition_mask=ign_mask,
    )

    plot_spread_statistics(sim_results, str(anim_dir / "spread_statistics.png"))

    # Training history
    hist_path = ckpt_dir / "history.json"
    if hist_path.exists():
        with open(hist_path) as f:
            history = json.load(f)
        plot_training_history(history, str(anim_dir / "training_history.png"))

    print(f"\nAll visualizations saved to {anim_dir}/")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Forest Fire Simulation Pipeline")
    parser.add_argument("--step", choices=["download", "preprocess", "train", "predict",
                                           "simulate", "visualize", "all"],
                        default="all")
    parser.add_argument("--config", default="config/config.yaml")
    parser.add_argument("--date", default="latest",
                        help="Date (YYYY-MM-DD) for prediction step")
    parser.add_argument("--ignition-threshold", type=int, default=3,
                        help="Minimum fire class (1-3) to use as CA seed")
    args = parser.parse_args()

    cfg = load_config(args.config)

    # Ensure output dirs exist
    for d in [cfg["data"]["raw_dir"], cfg["data"]["processed_dir"],
              cfg["output"]["prediction_dir"], cfg["output"]["simulation_dir"],
              cfg["output"]["animation_dir"], cfg["model"]["checkpoint_dir"]]:
        Path(d).mkdir(parents=True, exist_ok=True)

    steps = (
        ["download", "preprocess", "train", "predict", "simulate", "visualize"]
        if args.step == "all"
        else [args.step]
    )

    prediction_result = None
    simulation_data = None

    for s in steps:
        print(f"\n{'='*60}")
        print(f"  STEP: {s.upper()}")
        print(f"{'='*60}")

        if s == "download":
            step_download(cfg)
        elif s == "preprocess":
            step_preprocess(cfg)
        elif s == "train":
            step_train(cfg)
        elif s == "predict":
            prediction_result = step_predict(cfg, args.date)
        elif s == "simulate":
            simulation_data = step_simulate(cfg, prediction_result, args.ignition_threshold)
        elif s == "visualize":
            step_visualize(cfg, prediction_result, simulation_data)

    print("\n Pipeline complete.")


if __name__ == "__main__":
    main()