src/api.py

"""
Extreme Heat Risk Engine — FastAPI Application

Serves synthetic demo data for the dashboard.
When the real pipeline has been run, serves pipeline results instead.
"""

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

import asyncio
import logging
import os
import random
import threading
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
from pathlib import Path

from fastapi import FastAPI, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse

from config import ZONES, ZONE_MAP, CITIES, HEAT_THRESHOLDS, PAYOUT_PER_EVENT_USD
from src.indexing.heat_index import calculate_wbgt, calculate_heat_index, count_consecutive_days, count_trigger_days
from src.downscaling import get_uhi_corrector
from src.pricing.burn_analysis import BurnAnalysisPricer
from src.pricing.budget_optimizer import BudgetOptimizer
from src.database.crud import init_db, upsert_zone

logger = logging.getLogger(__name__)

# Database connection — set in lifespan
_db_conn = None


def _prewarm_graphcast() -> None:
    """Load GraphCast model into memory at startup so the first pipeline
    trigger doesn't pay the ~30-120s download/init cost. Runs in a
    background thread — failures are logged, not fatal.
    """
    try:
        from src.prediction.graphcast_inference import load_model
        import time as _time
        t0 = _time.time()
        load_model()
        logger.info("[PREWARM] GraphCast loaded at startup (%.1fs)", _time.time() - t0)
    except Exception as exc:
        logger.warning("[PREWARM] GraphCast prewarm threw: %s", exc)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global _db_conn
    # Sync DB init — same pattern as Weather AI 2
    try:
        _db_conn = init_db()
        if _db_conn:
            for z in ZONES:
                try:
                    upsert_zone(_db_conn, {
                        "zone_id": z.zone_id, "name": z.name, "city": z.city,
                        "country": z.country, "latitude": z.latitude,
                        "longitude": z.longitude, "elevation_m": z.elevation_m,
                        "area_km2": z.area_km2, "population_est": z.population_est,
                        "settlement_type": z.settlement_type,
                        "worker_population_est": z.worker_population_est,
                        "outdoor_exposure_pct": z.outdoor_exposure_pct,
                        "heat_vulnerability": z.heat_vulnerability,
                        "hot_months": z.hot_months, "notes": z.notes,
                    })
                except Exception as exc:
                    logger.warning("Failed to seed zone %s: %s", z.zone_id, exc)
            logger.info("Database ready (postgres, %d zones seeded)", len(ZONES))
        else:
            logger.info("Database ready (in-memory)")
    except Exception as e:
        logger.warning("DB init failed (non-fatal): %s", e)
        _db_conn = None

    prewarm_thread = threading.Thread(
        target=_prewarm_graphcast, daemon=True, name="graphcast-prewarm",
    )
    prewarm_thread.start()

    scheduler = _start_scheduler()
    yield
    if scheduler:
        scheduler.shutdown(wait=False)
    if _db_conn:
        _db_conn.close()


app = FastAPI(title="Extreme Heat Risk Engine", version="1.0.0", lifespan=lifespan)

# CORS origins configurable via ALLOWED_ORIGINS (comma-separated).
# Defaults to "*" so local dev and HF Spaces preview stay permissive.
_allowed_origins_env = os.environ.get("ALLOWED_ORIGINS", "*").strip()
if _allowed_origins_env == "*" or not _allowed_origins_env:
    _allowed_origins = ["*"]
else:
    _allowed_origins = [o.strip() for o in _allowed_origins_env.split(",") if o.strip()]

app.add_middleware(
    CORSMiddleware,
    allow_origins=_allowed_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

SEED = 42


def _generate_demo_data():
    """Deterministic synthetic data for the dashboard demo, using real ML models."""
    rng = random.Random(SEED)
    now = datetime(2026, 3, 29, 10, 0, 0)

    # Initialize ML models (UHI_MODEL env var selects synthetic or lst)
    uhi_corrector = get_uhi_corrector()

    # City base temperatures (ERA5-Land grid-level — before UHI correction)
    city_climate = {
        "Dar es Salaam": {"base_temp": 31, "temp_var": 2.5, "base_hum": 78, "hum_var": 8},
        "Kampala": {"base_temp": 28, "temp_var": 2.5, "base_hum": 68, "hum_var": 10},
        "Nairobi": {"base_temp": 25, "temp_var": 2.5, "base_hum": 55, "hum_var": 12},
        "Kigali": {"base_temp": 25, "temp_var": 2, "base_hum": 60, "hum_var": 10},
    }

    # Generate 90 days of daily data per zone
    zones = []
    indices = []
    all_triggers = []
    tid = 1

    for z in ZONES:
        clim = city_climate[z.city]

        daily_grid_temps = []
        daily_temps = []
        daily_humidity = []
        daily_dates = []
        daily_wbgt = []
        daily_hi = []
        daily_uhi_deltas = []

        for d in range(90):
            date = now - timedelta(days=89 - d)
            month = date.month
            seasonal = 1.5 if month in z.hot_months else -0.5
            # Grid-level temperature (ERA5-Land equivalent — before UHI)
            grid_temp = clim["base_temp"] + seasonal + rng.gauss(0, clim["temp_var"] * 0.4)
            grid_temp = round(max(18, min(42, grid_temp)), 1)
            hum = clim["base_hum"] + rng.gauss(0, clim["hum_var"] * 0.3)
            hum = round(max(30, min(95, hum)), 1)
            # ML UHI correction
            corrected, uhi_delta, _ = uhi_corrector.correct_temperature(z, grid_temp, hour=14, month=month)
            temp = round(corrected, 1)
            daily_grid_temps.append(grid_temp)
            wbgt = calculate_wbgt(temp, hum)
            hi = calculate_heat_index(temp, hum)

            daily_temps.append(temp)
            daily_humidity.append(hum)
            daily_dates.append(date.strftime("%Y-%m-%d"))
            daily_wbgt.append(wbgt)
            daily_hi.append(hi)
            daily_uhi_deltas.append(round(uhi_delta, 1))

        max_temp = max(daily_temps)
        max_wbgt = max(daily_wbgt)
        recent_temps = daily_temps[-7:]
        recent_wbgt = daily_wbgt[-7:]
        current_temp = daily_temps[-1]
        current_wbgt = daily_wbgt[-1]
        current_hi = daily_hi[-1]
        watch_temp = HEAT_THRESHOLDS["watch"]["temp_c"]
        consec = count_consecutive_days(recent_temps, watch_temp)
        total_above = count_trigger_days(daily_temps, watch_temp)

        # Risk level from config thresholds
        recent_max = max(recent_temps)
        risk_level = "normal"
        for level in ("critical", "warning", "watch"):
            ht = HEAT_THRESHOLDS[level]
            if recent_max >= ht["temp_c"] and consec >= ht["consecutive_days"]:
                risk_level = level
                break

        # Composite score
        temp_score = min(100, max(0, (max_temp - 28) * 10))
        wbgt_score = min(100, max(0, (max_wbgt - 25) * 12))
        vuln_score = {"high": 85, "moderate": 50, "low": 20}[z.heat_vulnerability]
        exposure_score = z.outdoor_exposure_pct * 100
        composite = round(temp_score * 0.3 + wbgt_score * 0.25 + consec * 10 * 0.2 + vuln_score * 0.15 + exposure_score * 0.1, 1)
        composite = min(100, max(0, composite))

        enrolled = int(z.worker_population_est * rng.uniform(0.15, 0.45))

        # Composite-driven trigger probability for demo shape. Not a model output.
        pred_prob = round(min(1.0, composite / 100), 2)
        pred_conf = 0.5
        pred_tier = "composite_heuristic"

        zone_data = {
            "zone_id": z.zone_id,
            "name": z.name,
            "city": z.city,
            "country": z.country,
            "latitude": z.latitude,
            "longitude": z.longitude,
            "elevation_m": z.elevation_m,
            "settlement_type": z.settlement_type,
            "worker_population_est": z.worker_population_est,
            "outdoor_exposure_pct": z.outdoor_exposure_pct,
            "heat_vulnerability": z.heat_vulnerability,
            "risk_level": risk_level,
            "current_temp_c": current_temp,
            "current_wbgt_c": current_wbgt,
            "current_heat_index_c": current_hi,
            "max_temp_c": round(max_temp, 1),
            "max_wbgt_c": round(max_wbgt, 1),
            "consecutive_hot_days": consec,
            "total_days_above_33": total_above,
            "heat_risk_score": composite,
            "grid_temp_c": daily_grid_temps[-1],
            "uhi_delta_c": daily_uhi_deltas[-1],
            "corrected_temp_c": temp,
            "trigger_probability_7d": round(pred_prob, 2),
            "prediction_confidence": round(pred_conf, 2),
            "model_tier": pred_tier,
            "enrolled_workers": enrolled,
            "data_quality": round(rng.uniform(0.80, 0.98), 2),
            "last_updated": now.isoformat(),
        }
        zones.append(zone_data)

        # Index data with daily history
        indices.append({
            "zone_id": z.zone_id,
            "zone_name": z.name,
            "city": z.city,
            "risk_level": risk_level,
            "temp_current": current_temp,
            "wbgt_current": current_wbgt,
            "heat_index_current": current_hi,
            "consecutive_hot_days": consec,
            "heat_risk_score": composite,
            "grid_temp_c": daily_grid_temps[-1],
            "uhi_delta_c": daily_uhi_deltas[-1],
            "trigger_probability_7d": round(pred_prob, 2),
            "prediction_confidence": round(pred_conf, 2),
            "model_tier": pred_tier,
            "daily_history": [
                {"date": daily_dates[i], "temp_c": daily_temps[i], "grid_temp_c": daily_grid_temps[i], "uhi_delta_c": daily_uhi_deltas[i], "humidity_pct": daily_humidity[i], "wbgt_c": daily_wbgt[i], "heat_index_c": daily_hi[i]}
                for i in range(90)
            ],
        })

        # Triggers
        if risk_level != "normal":
            payout = PAYOUT_PER_EVENT_USD.get(risk_level, 5)
            all_triggers.append({
                "trigger_id": f"TRG-{tid:04d}",
                "zone_id": z.zone_id,
                "zone_name": z.name,
                "city": z.city,
                "trigger_level": risk_level,
                "trigger_date": (now - timedelta(hours=rng.randint(2, 48))).isoformat(),
                "heat_risk_score": composite,
                "max_temp_c": round(max_temp, 1),
                "max_wbgt_c": round(max_wbgt, 1),
                "consecutive_days": consec,
                "total_days_above": total_above,
                "settlement_type": z.settlement_type,
                "payout_per_worker_usd": payout,
                "enrolled_workers": enrolled,
                "total_payout_usd": payout * enrolled,
                "status": "active",
            })
            tid += 1

    # Basis risk
    basis_risk = []
    for z_data in zones:
        zone_obj = ZONE_MAP[z_data["zone_id"]]
        if zone_obj.heat_vulnerability == "high" and zone_obj.settlement_type == "informal":
            score = rng.uniform(0.25, 0.40)
        elif zone_obj.heat_vulnerability == "high":
            score = rng.uniform(0.18, 0.32)
        elif zone_obj.heat_vulnerability == "moderate":
            score = rng.uniform(0.10, 0.22)
        else:
            score = rng.uniform(0.05, 0.15)
        basis_risk.append({
            "zone_id": z_data["zone_id"],
            "zone_name": z_data["name"],
            "city": z_data["city"],
            "overall_score": round(score, 3),
            "false_positive_rate": round(score * rng.uniform(0.4, 0.7), 3),
            "false_negative_rate": round(score * rng.uniform(0.3, 0.6), 3),
            "correlation": round(1 - score * rng.uniform(0.8, 1.1), 3),
            "settlement_type": z_data["settlement_type"],
            "heat_vulnerability": z_data["heat_vulnerability"],
            "recommendation": (
                "Urban heat island effect significant — consider localized temperature sensors"
                if zone_obj.settlement_type == "informal"
                else "Station temperature may underestimate worker-experienced heat by 2-3°C"
                if score > 0.2
                else "Current calibration adequate for this zone"
            ),
        })

    # Notifications
    notifications = []
    nid = 1
    for trigger in all_triggers:
        if trigger["trigger_level"] in ("critical", "warning"):
            notifications.append({
                "id": f"NOT-{nid:04d}",
                "zone_id": trigger["zone_id"],
                "zone_name": trigger["zone_name"],
                "city": trigger["city"],
                "trigger_level": trigger["trigger_level"],
                "channel": rng.choice(["sms", "whatsapp"]),
                "language": rng.choice(["en", "sw"]),
                "recipient_count": trigger["enrolled_workers"],
                "message_preview": (
                    f"HEAT ALERT [{trigger['trigger_level'].upper()}]: "
                    f"{trigger['zone_name']}, {trigger['city']}. "
                    f"Temperature {trigger['max_temp_c']}°C (WBGT {trigger['max_wbgt_c']}°C). "
                    f"Payout: ${trigger['payout_per_worker_usd']}."
                ),
                "status": "sent",
                "delivered_at": trigger["trigger_date"],
                "cost_estimate": round(trigger["enrolled_workers"] * 0.0075, 2),
            })
            nid += 1
            notifications.append({
                "id": f"NOT-{nid:04d}",
                "zone_id": trigger["zone_id"],
                "zone_name": trigger["zone_name"],
                "city": trigger["city"],
                "trigger_level": trigger["trigger_level"],
                "channel": "sms",
                "language": "sw",
                "recipient_count": trigger["enrolled_workers"],
                "message_preview": (
                    f"TAHADHARI YA JOTO [{trigger['trigger_level'].upper()}]: "
                    f"{trigger['zone_name']}, {trigger['city']}. "
                    f"Joto {trigger['max_temp_c']}°C. "
                    f"Malipo: ${trigger['payout_per_worker_usd']}."
                ),
                "status": "sent",
                "delivered_at": trigger["trigger_date"],
                "cost_estimate": round(trigger["enrolled_workers"] * 0.0075, 2),
            })
            nid += 1

    # Pipeline runs
    pipeline_runs = []
    for i in range(15):
        run_date = now - timedelta(days=i * 2)
        duration = rng.uniform(30, 120)
        cost = rng.uniform(0.06, 0.18)
        status = "ok" if rng.random() > 0.15 else "partial"
        pipeline_runs.append({
            "run_id": f"run-{1000 + i}",
            "started_at": run_date.isoformat(),
            "ended_at": (run_date + timedelta(seconds=duration)).isoformat(),
            "status": status,
            "duration_s": round(duration, 1),
            "zones_processed": 20,
            "triggers_found": rng.randint(0, 8),
            "notifications_sent": rng.randint(0, 16),
            "total_cost_usd": round(cost, 4),
            "steps": [
                {"step": s, "status": "ok", "duration_s": round(duration / 6, 1)}
                for s in ["ingest", "heal", "index", "calibrate", "explain", "notify"]
            ],
        })

    stats = {
        "total_runs": len(pipeline_runs),
        "successful_runs": sum(1 for r in pipeline_runs if r["status"] == "ok"),
        "success_rate": round(sum(1 for r in pipeline_runs if r["status"] == "ok") / len(pipeline_runs), 2),
        "zones_monitored": len(ZONES),
        "cities": len(CITIES),
        "active_triggers": len(all_triggers),
        "total_enrolled": sum(z["enrolled_workers"] for z in zones),
        "total_cost_usd": round(sum(r["total_cost_usd"] for r in pipeline_runs), 2),
        "avg_cost_per_run_usd": round(sum(r["total_cost_usd"] for r in pipeline_runs) / len(pipeline_runs), 4),
        "last_run": pipeline_runs[0]["started_at"],
        "data_sources": ["NASA POWER"],
    }

    return {
        "zones": zones,
        "indices": indices,
        "triggers": all_triggers,
        "basis_risk": basis_risk,
        "notifications": notifications,
        "pipeline_runs": pipeline_runs,
        "stats": stats,
    }


_demo = None


def _get_demo():
    """Lazy initialization of demo data — only generated on first API request."""
    global _demo
    if _demo is None:
        _demo = _generate_demo_data()
    return _demo


# Singletons for calibrate endpoint (avoid re-instantiation per request)
_actuarial_pricer = BurnAnalysisPricer()
_budget_optimizer = BudgetOptimizer()


# ── API Endpoints ──────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "service": "extreme-heat-risk-engine", "version": "1.0.0"}


@app.get("/api/zones")
def get_zones():
    return {"zones": _get_demo()["zones"], "total": len(_get_demo()["zones"]), "cities": CITIES}


@app.get("/api/indices")
def get_indices():
    return {"indices": _get_demo()["indices"], "total": len(_get_demo()["indices"])}


@app.get("/api/triggers")
def get_triggers():
    triggers = _get_demo()["triggers"]
    return {
        "triggers": triggers,
        "total": len(triggers),
        "active": sum(1 for t in triggers if t["status"] == "active"),
        "by_level": {
            level: sum(1 for t in triggers if t["trigger_level"] == level)
            for level in ["critical", "warning", "watch"]
        },
    }


@app.get("/api/basis-risk")
def get_basis_risk():
    br = _get_demo()["basis_risk"]
    return {
        "assessments": br,
        "total": len(br),
        "avg_score": round(sum(b["overall_score"] for b in br) / max(1, len(br)), 3),
    }


@app.get("/api/notifications")
def get_notifications():
    notifs = _get_demo()["notifications"]
    return {
        "notifications": notifs,
        "total": len(notifs),
        "by_language": {
            lang: sum(1 for n in notifs if n["language"] == lang)
            for lang in ["en", "sw"]
        },
    }


@app.get("/api/enrolled-workers")
def get_enrolled():
    by_zone = [
        {"zone_id": z["zone_id"], "zone_name": z["name"], "city": z["city"], "enrolled": z["enrolled_workers"]}
        for z in _get_demo()["zones"]
    ]
    return {"by_zone": by_zone, "total_enrolled": sum(z["enrolled_workers"] for z in _get_demo()["zones"])}


@app.get("/api/pipeline/runs")
def get_pipeline_runs():
    return {"runs": _get_demo()["pipeline_runs"], "total": len(_get_demo()["pipeline_runs"])}


@app.get("/api/pipeline/stats")
def get_pipeline_stats():
    return _get_demo()["stats"]


@app.get("/api/coverage-recommendation")
def get_coverage_recommendation(payout_usd: float = 10.0):
    """Neural model-driven coverage recommendation.

    The model analyzes current heat conditions across all zones and
    recommends: how much coverage is needed, where, and at what cost.
    No budget input — the model TELLS you what the budget should be.
    """
    demo = _get_demo()
    zones_data = demo["zones"]
    indices = demo["indices"]
    basis = demo["basis_risk"]

    basis_by_id = {b["zone_id"]: b for b in basis}
    indices_by_id = {idx["zone_id"]: idx for idx in indices}

    zone_recommendations = []
    total_recommended_budget = 0.0
    total_workers_at_risk = 0
    total_workers_enrolled = 0

    for z in zones_data:
        zone_id = z["zone_id"]
        zone = ZONE_MAP.get(zone_id)
        if not zone:
            continue

        idx = indices_by_id.get(zone_id, {})
        history = idx.get("daily_history", [])
        br = basis_by_id.get(zone_id, {})

        # Get trigger probability from forecast
        trigger_prob = z.get("trigger_probability_7d", 0)
        current_temp = z.get("corrected_temp_c", z.get("current_temp_c", 30))
        current_wbgt = z.get("current_wbgt_c", 28)
        consecutive = z.get("consecutive_hot_days", 0)
        risk_level = z.get("risk_level", "normal")
        enrolled = z.get("enrolled_workers", 0)

        # Neural pricing (uses climate history if available)
        ar = _actuarial_pricer.price_zone(
            zone=zone,
            predicted_frequency=z.get("events_per_year", 10),
            basis_risk_score=br.get("overall_score", 0.2),
            payout_per_event=payout_usd,
            enrolled=max(enrolled, 1),
            climate_history=history if history else None,
        )

        # Workers at risk this week (based on trigger probability)
        workers_at_risk = int(enrolled * trigger_prob * zone.outdoor_exposure_pct)

        # Recommended weekly payout for this zone
        weekly_payout = workers_at_risk * payout_usd
        annual_cost = ar.cost_per_worker_year * enrolled

        # Urgency level
        if trigger_prob > 0.7 or risk_level == "critical":
            urgency = "critical"
        elif trigger_prob > 0.4 or risk_level in ("warning", "high"):
            urgency = "high"
        elif trigger_prob > 0.15:
            urgency = "moderate"
        else:
            urgency = "low"

        total_recommended_budget += annual_cost
        total_workers_at_risk += workers_at_risk
        total_workers_enrolled += enrolled

        # Cost decomposition
        cb = ar.cost_breakdown
        payout_fraction = ar.expected_annual_payouts / max(annual_cost, 1)
        admin_fraction = ar.admin_loading / max(annual_cost, 1)
        basis_risk_fraction = ar.basis_risk_loading / max(annual_cost, 1)

        zone_recommendations.append({
            "zone_id": zone_id,
            "zone_name": zone.name,
            "city": zone.city,
            "settlement_type": zone.settlement_type,
            "heat_vulnerability": zone.heat_vulnerability,
            "urgency": urgency,
            # Current conditions
            "current_temp_c": round(current_temp, 1),
            "current_wbgt_c": round(current_wbgt, 1),
            "consecutive_hot_days": consecutive,
            "trigger_probability_7d": round(trigger_prob, 3),
            "risk_level": risk_level,
            # Worker impact
            "enrolled_workers": enrolled,
            "outdoor_exposure_pct": zone.outdoor_exposure_pct,
            "workers_at_risk_this_week": workers_at_risk,
            # Cost
            "annual_cost_per_worker": round(ar.cost_per_worker_year, 2),
            "annual_cost_total": round(annual_cost, 0),
            "weekly_recommended_payout": round(weekly_payout, 0),
            "payout_usd_per_event": payout_usd,
            # Decomposition
            "cost_to_workers_pct": round(payout_fraction * 100, 1),
            "cost_admin_pct": round(admin_fraction * 100, 1),
            "cost_basis_risk_pct": round(basis_risk_fraction * 100, 1),
            # Neural model outputs (if available)
            "neural_model": cb.get("neural_correction_pct") is not None,
            "neural_correction_pct": cb.get("neural_correction_pct"),
            "learned_frequency": cb.get("learned_frequency"),
            "learned_basis_risk": cb.get("learned_basis_risk"),
            "productivity_loss_rate": cb.get("productivity_loss_rate"),
            "gpd_shape_xi": cb.get("gpd_shape_xi"),
        })

    # Sort by urgency then annual cost
    urgency_order = {"critical": 0, "high": 1, "moderate": 2, "low": 3}
    zone_recommendations.sort(key=lambda z: (urgency_order.get(z["urgency"], 9), -z["annual_cost_total"]))

    # Weekly budget recommendation
    weekly_budget = sum(z["weekly_recommended_payout"] for z in zone_recommendations)

    return {
        "recommendation": {
            "annual_budget_needed": round(total_recommended_budget, 0),
            "weekly_budget_needed": round(weekly_budget, 0),
            "total_workers_enrolled": total_workers_enrolled,
            "workers_at_risk_this_week": total_workers_at_risk,
            "zones_at_risk": sum(1 for z in zone_recommendations if z["urgency"] in ("critical", "high")),
            "payout_per_event": payout_usd,
            "model_type": "burn_analysis",
        },
        "zones": zone_recommendations,
        "cost_summary": {
            "total_to_workers_pct": round(
                sum(z["annual_cost_total"] * z["cost_to_workers_pct"] / 100 for z in zone_recommendations)
                / max(total_recommended_budget, 1) * 100, 1
            ),
            "total_admin_pct": round(
                sum(z["annual_cost_total"] * z["cost_admin_pct"] / 100 for z in zone_recommendations)
                / max(total_recommended_budget, 1) * 100, 1
            ),
            "total_basis_risk_pct": round(
                sum(z["annual_cost_total"] * z["cost_basis_risk_pct"] / 100 for z in zone_recommendations)
                / max(total_recommended_budget, 1) * 100, 1
            ),
        },
    }


@app.get("/api/calibrate")
def calibrate(
    temp_threshold: float = 35.0,
    consecutive_days: int = 2,
    wbgt_threshold: float = 30.0,
    payout_usd: float = 10.0,
    budget_usd: float = 500000.0,
    worker_contribution_usd: float = 0.0,
):
    """Interactive calibration endpoint.

    Run heat risk scoring with custom thresholds against all zones.
    Returns per-zone trigger analysis and program cost estimates.
    """
    rng = random.Random(SEED)
    results = []
    total_trigger_days = 0
    total_annual_cost = 0.0
    zones_triggered = 0

    zones_by_id = {z["zone_id"]: z for z in _get_demo()["zones"]}
    basis_by_id = {b["zone_id"]: b for b in _get_demo()["basis_risk"]}

    for idx_data in _get_demo()["indices"]:
        zone_id = idx_data["zone_id"]
        zone = ZONE_MAP.get(zone_id)
        if not zone:
            continue

        # Extract daily temps and humidity from history
        history = idx_data.get("daily_history", [])
        temps = [d["temp_c"] for d in history]
        humidity = [d["humidity_pct"] for d in history]
        wbgts = [d["wbgt_c"] for d in history]

        # Apply custom thresholds
        days_above_temp = count_trigger_days(temps, temp_threshold)
        days_above_wbgt = count_trigger_days(wbgts, wbgt_threshold)
        consec_temp = count_consecutive_days(temps, temp_threshold)
        consec_wbgt = count_consecutive_days(wbgts, wbgt_threshold)

        # Count trigger events (consecutive runs above threshold)
        trigger_events = 0
        run_length = 0
        for t in temps:
            if t > temp_threshold:
                run_length += 1
            else:
                if run_length >= consecutive_days:
                    trigger_events += 1
                run_length = 0
        if run_length >= consecutive_days:
            trigger_events += 1

        # Annualize (90 days of data → multiply by 4)
        events_per_year = round(trigger_events * (365 / max(len(temps), 1)), 1)

        zone_demo = zones_by_id.get(zone_id, {})
        enrolled = zone_demo.get("enrolled_workers", 0)

        annual_payout = round(events_per_year * payout_usd * enrolled, 2)
        annual_per_worker = round(events_per_year * payout_usd, 2)

        br = basis_by_id.get(zone_id, {})
        basis_score = br.get("overall_score", 0.15)

        triggered = trigger_events > 0
        if triggered:
            zones_triggered += 1
        total_trigger_days += days_above_temp
        total_annual_cost += annual_payout

        results.append({
            "zone_id": zone_id,
            "zone_name": zone.name,
            "city": zone.city,
            "settlement_type": zone.settlement_type,
            "heat_vulnerability": zone.heat_vulnerability,
            "enrolled_workers": enrolled,
            "days_above_temp": days_above_temp,
            "days_above_wbgt": days_above_wbgt,
            "consecutive_days_temp": consec_temp,
            "consecutive_days_wbgt": consec_wbgt,
            "trigger_events": trigger_events,
            "events_per_year": events_per_year,
            "annual_payout_per_worker": annual_per_worker,
            "annual_payout_total": annual_payout,
            "basis_risk_score": basis_score,
            "triggered": triggered,
        })

    total_enrolled = sum(r["enrolled_workers"] for r in results)

    # Actuarial pricing per zone
    indices_by_id = {idx["zone_id"]: idx for idx in _get_demo()["indices"]}
    actuarial_results = []
    for r in results:
        zone = ZONE_MAP.get(r["zone_id"])
        if not zone:
            continue
        idx_data = indices_by_id.get(r["zone_id"])
        history = idx_data.get("daily_history") if idx_data else None
        ar = _actuarial_pricer.price_zone(
            zone=zone,
            predicted_frequency=r["events_per_year"],
            basis_risk_score=r["basis_risk_score"],
            payout_per_event=payout_usd,
            enrolled=r["enrolled_workers"],
            climate_history=history,
        )
        r["actuarial_cost_per_worker"] = round(ar.cost_per_worker_year, 2)
        r["cost_breakdown"] = ar.cost_breakdown
        actuarial_results.append(ar)

    # Budget allocation
    allocation = _budget_optimizer.optimize(
        budget_usd=budget_usd,
        actuarial_results=actuarial_results,
        payout_per_event=payout_usd,
        worker_contribution=worker_contribution_usd,
    )

    # Merge allocation into zone results
    alloc_map = {a.zone_id: a for a in allocation.allocations}
    for r in results:
        a = alloc_map.get(r["zone_id"])
        if a:
            r["allocated_budget"] = round(a.allocated_budget, 2)
            r["workers_covered"] = a.workers_covered
            r["coverage_pct"] = round(a.coverage_pct, 1)
            r["priority_rank"] = a.priority_rank
        else:
            r["allocated_budget"] = 0
            r["workers_covered"] = 0
            r["coverage_pct"] = 0
            r["priority_rank"] = 99

    return {
        "zones": sorted(results, key=lambda r: r.get("priority_rank", 99)),
        "summary": {
            "total_zones": len(results),
            "zones_triggered": zones_triggered,
            "total_trigger_days": total_trigger_days,
            "avg_events_per_year": round(sum(r["events_per_year"] for r in results) / max(1, len(results)), 1),
            "total_annual_cost": round(total_annual_cost, 2),
            "avg_cost_per_worker": round(total_annual_cost / max(1, total_enrolled), 2),
            "total_enrolled": total_enrolled,
            "avg_basis_risk": round(sum(r["basis_risk_score"] for r in results) / max(1, len(results)), 3),
        },
        "allocation": {
            "budget_usd": budget_usd,
            "worker_contribution_usd": worker_contribution_usd,
            "workers_covered": allocation.total_workers_covered,
            "overall_coverage_pct": round(allocation.overall_coverage_pct, 1),
            "zones_fully_funded": allocation.zones_fully_funded,
            "zones_partially_funded": allocation.zones_partially_funded,
            "zones_unfunded": allocation.zones_unfunded,
            "stretch_analysis": allocation.stretch_analysis,
        },
        "thresholds": {
            "temp_threshold": temp_threshold,
            "consecutive_days": consecutive_days,
            "wbgt_threshold": wbgt_threshold,
            "payout_usd": payout_usd,
        },
    }


# ── Pipeline trigger ──────────────────────────────────────────────────────

_pipeline_status = {
    "running": False,
    "current_step": None,
    "current_step_index": 0,
    "total_steps": 6,
    "last_result": None,
    "last_run": None,
}


async def _run_pipeline_async():
    """Run the full pipeline in background, writing results to Neon."""
    global _db_conn
    from src.pipeline import STEP_LABELS

    _pipeline_status["running"] = True
    _pipeline_status["current_step"] = None
    _pipeline_status["current_step_index"] = 0

    # Refresh the Neon connection before the run. _db_conn is created once
    # in lifespan, but between pipeline runs (the space can sit idle for
    # hours/days) Neon's autosuspend kills the underlying socket. Without
    # this refresh the first DB write would throw, _db_write would set
    # self.db=None inside the pipeline, and every downstream write would
    # silently no-op -- the pipeline reports status=ok but Neon has no row.
    if _db_conn is not None:
        try:
            _db_conn._refresh_conn()
        except Exception as exc:
            logger.warning("[PIPELINE] DB refresh failed, reconnecting: %s", exc)
            try:
                _db_conn.close()
            except Exception:
                pass
            try:
                _db_conn = init_db()
            except Exception as reconnect_exc:
                logger.warning("[PIPELINE] DB reconnect failed: %s", reconnect_exc)
                _db_conn = None

    def _progress_cb(step_name, step_index):
        _pipeline_status["current_step"] = step_name
        _pipeline_status["current_step_index"] = step_index
        if step_name:
            label = STEP_LABELS.get(step_name, step_name)
            print(f"[PIPELINE] Step {step_index}/6: {label}", flush=True)

    try:
        from src.pipeline import run_pipeline_sync
        result = await asyncio.get_event_loop().run_in_executor(
            None,
            lambda: run_pipeline_sync(
                days_back=14,
                use_claude_healer=bool(os.environ.get("ANTHROPIC_API_KEY")),
                use_claude_explainer=bool(os.environ.get("ANTHROPIC_API_KEY")),
                delivery_channel="console",
                db=_db_conn,
                progress_callback=_progress_cb,
            ),
        )
        _pipeline_status["last_result"] = {
            "run_id": result.run_id,
            "status": result.status,
            "zones_processed": result.zones_processed,
            "triggers_found": result.triggers_found,
            "duration_s": round(result.duration_s, 1),
        }
        _pipeline_status["last_run"] = datetime.utcnow().isoformat()
        print(f"[PIPELINE] Complete: {result.status} — {result.zones_processed} zones, {result.triggers_found} triggers, {result.duration_s:.1f}s", flush=True)
    except Exception as e:
        print(f"[PIPELINE] FAILED: {e}", flush=True)
        _pipeline_status["last_result"] = {"status": "failed", "error": str(e)}
    finally:
        _pipeline_status["running"] = False
        _pipeline_status["current_step"] = None
        _pipeline_status["current_step_index"] = 0


@app.post("/api/pipeline/trigger")
async def trigger_pipeline(background_tasks: BackgroundTasks):
    """Trigger a pipeline run. Returns immediately; pipeline runs in background."""
    if _pipeline_status["running"]:
        return {"status": "already_running", "message": "A pipeline run is already in progress"}
    background_tasks.add_task(_run_pipeline_async)
    return {"status": "started", "message": "Pipeline run started in background"}


@app.get("/api/pipeline/status")
def pipeline_status():
    """Check if a pipeline run is in progress and get last result."""
    return _pipeline_status


# ── Scheduled pipeline runs ──────────────────────────────────────────────

def _start_scheduler():
    """Start weekly pipeline scheduler (runs in background thread)."""
    try:
        from apscheduler.schedulers.background import BackgroundScheduler
        scheduler = BackgroundScheduler()
        scheduler.add_job(
            lambda: asyncio.run(_run_pipeline_async()),
            "cron",
            day_of_week="tue",
            hour=0, minute=30,
            id="weekly_pipeline",
        )
        scheduler.start()
        logger.info("Weekly pipeline scheduler started (Tuesdays 00:30 UTC)")
        return scheduler
    except ImportError:
        logger.info("apscheduler not installed — no scheduled runs")
        return None
    except Exception as e:
        logger.warning("Scheduler failed to start: %s", e)
        return None


# ── Status page (lightweight, no React build needed) ─────────────────────

_STEP_NAMES = {
    "ingest": "Collecting climate data",
    "heal": "Fixing data issues",
    "downscale": "Adjusting for urban heat",
    "predict": "Forecasting heat danger",
    "explain": "Generating alerts",
    "review": "AI review & recommendation",
}


_PIPELINE_STEPS = ["ingest", "heal", "downscale", "predict", "explain", "review"]


def _pipeline_tracker_html() -> str:
    """Generate HTML for the vertical pipeline tracker with dots and run button."""
    ps = _pipeline_status
    running = ps["running"]
    current = ps.get("current_step")
    last = ps.get("last_result")

    # Build step rows
    completed_steps = []
    if last and last.get("status") in ("ok", "partial") and not running:
        completed_steps = _PIPELINE_STEPS  # all done

    rows = ""
    for step in _PIPELINE_STEPS:
        label = _STEP_NAMES.get(step, step)
        if running:
            if current and _PIPELINE_STEPS.index(step) < _PIPELINE_STEPS.index(current):
                cls = "done"
            elif step == current:
                cls = "active"
            else:
                cls = "pending"
        elif step in completed_steps:
            cls = "done"
        else:
            cls = "pending"
        rows += f'<div class="step-row"><div class="step-dot {cls}"></div><span class="step-name">{label}</span></div>\n'

    # Last run info
    last_html = ""
    if last and not running:
        status = last.get("status", "unknown")
        dur = last.get("duration_s", 0)
        zones = last.get("zones_processed", 0)
        triggers = last.get("triggers_found", 0)
        cls = "ok" if status in ("ok", "partial") else "failed"
        last_html = f'<div class="last-run"><span class="{cls}">{status.upper()}</span> — {zones} zones, {triggers} triggers, {dur:.0f}s</div>'

    btn_disabled = "disabled" if running else ""
    btn_text = "Running..." if running else "Run Pipeline"

    return f"""
    <div class="pipeline-tracker">
      <h3>Pipeline</h3>
      {rows}
      <button class="trigger-btn" {btn_disabled} onclick="fetch('/api/pipeline/trigger',{{method:'POST'}}).then(()=>location.reload())">{btn_text}</button>
      {last_html}
    </div>"""


@app.get("/", response_class=HTMLResponse)
async def status_page():
    """Pipeline tracker for the HF Space."""
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta http-equiv="refresh" content="5">
<title>Heat Risk Engine</title>
<style>
  * {{ margin: 0; padding: 0; box-sizing: border-box; }}
  body {{ font-family: system-ui, -apple-system, sans-serif; background: #faf8f5; color: #1a1a1a; padding: 32px; max-width: 480px; margin: 0 auto; }}
  h1 {{ font-size: 1.4rem; font-weight: 700; margin-bottom: 4px; }}
  .subtitle {{ color: #888; font-size: 0.85rem; margin-bottom: 24px; }}
  .link {{ color: #e63946; text-decoration: none; font-weight: 600; }}
  .link:hover {{ text-decoration: underline; }}
  .pipeline-tracker {{ background: #fff; border: 1px solid #e0dcd5; border-radius: 8px; padding: 16px; margin-bottom: 24px; }}
  .pipeline-tracker h3 {{ font-size: 0.8rem; font-weight: 700; text-transform: uppercase; letter-spacing: 0.5px; color: #e63946; margin-bottom: 12px; }}
  .step-row {{ display: flex; align-items: center; gap: 10px; padding: 6px 0; font-size: 0.82rem; }}
  .step-dot {{ width: 10px; height: 10px; border-radius: 50%; flex-shrink: 0; }}
  .step-dot.done {{ background: #2a9d8f; }}
  .step-dot.active {{ background: #e63946; animation: pulse 1.2s infinite; }}
  .step-dot.pending {{ background: #e0dcd5; }}
  .step-dot.failed {{ background: #e63946; }}
  .step-name {{ font-weight: 600; min-width: 110px; }}
  .step-time {{ color: #888; font-size: 0.75rem; }}
  .trigger-btn {{ display: inline-block; margin-top: 12px; padding: 8px 20px; background: #e63946; color: #fff; border: none; border-radius: 6px; font-size: 0.8rem; font-weight: 600; letter-spacing: 0.5px; text-transform: uppercase; cursor: pointer; }}
  .trigger-btn:hover {{ background: #c5303c; }}
  .trigger-btn:disabled {{ opacity: 0.5; cursor: not-allowed; }}
  .last-run {{ font-size: 0.78rem; color: #888; margin-top: 8px; }}
  .last-run .ok {{ color: #2a9d8f; font-weight: 600; }}
  .last-run .failed {{ color: #e63946; font-weight: 600; }}
  @keyframes pulse {{ 0%, 100% {{ opacity: 1; }} 50% {{ opacity: 0.4; }} }}
</style>
</head>
<body>
  <h1>Heat Risk Engine</h1>
  <p class="subtitle"><a class="link" href="https://climate-risk-engine.vercel.app" target="_blank">Open Dashboard</a></p>

  {_pipeline_tracker_html()}
</body>
</html>"""