api/main.py

"""
AquiScore FastAPI Server — Groundwater Security API v2.0

Wraps pipeline/aquifer_model.py as a REST API with v2.0 scientific enrichments:
  - CCME Water Quality Index (replaces linear quality scoring)
  - DRASTIC groundwater vulnerability assessment
  Plus v1.x: batch scoring, response caching, GRACE decomposition, kriging

Endpoints:
  POST /v1/score              — Run aquifer security score (v2.0 enriched)
  POST /v1/score/batch        — Batch score multiple sites
  POST /v1/certificate        — Generate aquifer certificate
  POST /v1/grace/decompose    — GRACE TWS signal decomposition
  POST /v1/wells/interpolate  — Kriging well field interpolation
  POST /v1/wqi                — Standalone CCME Water Quality Index
  POST /v1/drastic            — Standalone DRASTIC vulnerability assessment
  GET  /v1/aquifer-types      — List aquifer types
  GET  /v1/extraction-regimes — List extraction regimes
  GET  /v1/quality-thresholds — WHO/EPA water quality thresholds
  GET  /v1/grace-hotspots     — GRACE-FO depletion hotspots
  GET  /v1/grace-regions      — GRACE regional references
  GET  /v1/presets/{name}     — Run canonical preset
  GET  /v1/cache/stats        — Cache performance statistics
  POST /v1/cache/clear        — Clear response cache
  GET  /v1/health             — Health check

Usage:
  uvicorn api.main:app --host 0.0.0.0 --port 8001 --reload
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, field_validator

from pipeline.aquifer_model import (
    run_aquifer_score,
    AQUIFER_TYPES,
    EXTRACTION_REGIMES,
    WATER_QUALITY_THRESHOLDS,
    RISK_THRESHOLDS,
    PRESET_PRISTINE_ALLUVIAL,
    PRESET_STRESSED_AGRICULTURAL,
    PRESET_CRITICAL_DEPLETION,
)
from pipeline.certificate_generator import generate_certificate_json, generate_certificate_text
from pipeline.grace_fetch import REGIONAL_GRACE_REFS, list_depletion_hotspots
from pipeline.cache import get_default_cache, ScoreCache
from pipeline.grace_decomposition import decompose_grace_signal, classify_depletion_pattern
from pipeline.kriging import interpolate_well_field
from pipeline.ccme_wqi import compute_ccme_wqi, GUIDELINES as CCME_GUIDELINES
from pipeline.drastic_rating import compute_drastic, compute_drastic_from_aquifer_type

# Module-level cache instance
_cache = get_default_cache()

app = FastAPI(
    title="AquiScore API",
    description=(
        "Groundwater Security Scoring Engine v2.0 — "
        "GRACE-FO × NWIS × GLHYMPS. v2.0: CCME Water Quality Index, "
        "DRASTIC vulnerability framework, batch scoring, response caching, "
        "GRACE decomposition, and kriging interpolation."
    ),
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

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


# ── REQUEST / RESPONSE MODELS ───────────────────────────────────────────────

class GRACEInput(BaseModel):
    tws_anomaly_cm: float = 0.0
    trend_cm_yr: float = 0.0
    seasonal_amplitude_cm: float = 5.0
    uncertainty_cm: float = 2.0


class WellInput(BaseModel):
    mean_depth_m: float = 15.0
    trend_m_yr: float = 0.0
    n_wells: int = 10
    record_years: int = 20
    seasonal_range_m: float = 3.0


class SubstrateInput(BaseModel):
    hydraulic_conductivity_m_s: float = 1e-5
    porosity: float = 0.25
    depth_to_water_m: float = 15.0
    specific_yield: float = 0.15
    transmissivity_m2_d: float = 500.0


class QualityInput(BaseModel):
    nitrate_mg_l: float = 5.0
    arsenic_ug_l: float = 3.0
    fluoride_mg_l: float = 0.5
    tds_mg_l: float = 300.0
    ph: float = 7.2
    e_coli_cfu_100ml: float = 0.0


class ScoreRequest(BaseModel):
    grace: GRACEInput
    well: WellInput
    substrate: SubstrateInput
    quality: QualityInput
    aquifer_type: str
    extraction_regime: str
    site_id: Optional[str] = None
    site_name: Optional[str] = None
    coordinates: Optional[dict] = None

    @field_validator("aquifer_type")
    @classmethod
    def validate_aquifer_type(cls, v):
        if v not in AQUIFER_TYPES:
            raise ValueError(f"Invalid aquifer_type. Must be one of: {list(AQUIFER_TYPES.keys())}")
        return v

    @field_validator("extraction_regime")
    @classmethod
    def validate_extraction_regime(cls, v):
        if v not in EXTRACTION_REGIMES:
            raise ValueError(f"Invalid extraction_regime. Must be one of: {list(EXTRACTION_REGIMES.keys())}")
        return v


# ── ENDPOINTS ────────────────────────────────────────────────────────────────

@app.get("/v1/health")
async def health_check():
    return {
        "status": "healthy",
        "platform": "AquiScore",
        "version": "2.0.0",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "v1_optimizations": [
            "batch_scoring", "response_caching", "grace_decomposition",
            "kriging_interpolation", "well_field_analysis",
        ],
        "v2_capabilities": [
            "ccme_water_quality_index", "drastic_vulnerability",
        ],
    }


@app.post("/v1/score")
async def run_score(request: ScoreRequest):
    """
    Run aquifer security score with response caching.

    Fuses satellite (GRACE-FO), well measurement, substrate (GLHYMPS),
    and water quality signals with extraction regime multiplier.
    """
    # Build cache key from all scoring inputs
    cache_key = _cache.make_key(
        grace=request.grace.model_dump(),
        well=request.well.model_dump(),
        substrate=request.substrate.model_dump(),
        quality=request.quality.model_dump(),
        aquifer_type=request.aquifer_type,
        extraction_regime=request.extraction_regime,
    )

    cached = _cache.get(cache_key)
    if cached is not None:
        cached["_cached"] = True
        return cached

    result = run_aquifer_score(
        grace=request.grace.model_dump(),
        well=request.well.model_dump(),
        substrate=request.substrate.model_dump(),
        quality=request.quality.model_dump(),
        aquifer_type=request.aquifer_type,
        extraction_regime=request.extraction_regime,
    )

    response = {
        "score": result.score,
        "confidence_interval": result.confidence_interval,
        "confidence_pct": result.confidence_pct,
        "risk_class": result.risk_class,
        "depletion_rate_cm_yr": result.depletion_rate_cm_yr,
        "years_to_critical": result.years_to_critical,
        "satellite_score": result.satellite_score,
        "well_score": result.well_score,
        "substrate_score": result.substrate_score,
        "quality_score": result.quality_score,
        "extraction_multiplier": result.extraction_multiplier,
        "quality_flags": result.quality_flags,
        "feature_importances": result.feature_importances,
        "citations": result.citations,
        "ccme_wqi": result.ccme_wqi,
        "drastic": result.drastic,
        "_cached": False,
    }

    _cache.set(cache_key, response)
    return response


@app.post("/v1/certificate")
async def generate_cert(request: ScoreRequest):
    result = run_aquifer_score(
        grace=request.grace.model_dump(),
        well=request.well.model_dump(),
        substrate=request.substrate.model_dump(),
        quality=request.quality.model_dump(),
        aquifer_type=request.aquifer_type,
        extraction_regime=request.extraction_regime,
    )

    cert = generate_certificate_json(
        result,
        site_id=request.site_id or "API-REQUEST",
        site_name=request.site_name or "Unknown Site",
        coordinates=request.coordinates,
    )

    return {"certificate": cert, "text_display": generate_certificate_text(cert)}


@app.get("/v1/aquifer-types")
async def get_aquifer_types():
    return AQUIFER_TYPES


@app.get("/v1/extraction-regimes")
async def get_extraction_regimes():
    return EXTRACTION_REGIMES


@app.get("/v1/quality-thresholds")
async def get_quality_thresholds():
    return WATER_QUALITY_THRESHOLDS


@app.get("/v1/grace-hotspots")
async def get_grace_hotspots():
    return list_depletion_hotspots(threshold_cm_yr=-1.0)


@app.get("/v1/grace-regions")
async def get_grace_regions():
    return REGIONAL_GRACE_REFS


@app.get("/v1/presets/{preset_name}/score")
async def run_preset(preset_name: str):
    presets = {
        "pristine": PRESET_PRISTINE_ALLUVIAL,
        "stressed": PRESET_STRESSED_AGRICULTURAL,
        "critical": PRESET_CRITICAL_DEPLETION,
    }
    if preset_name not in presets:
        raise HTTPException(404, f"Unknown preset. Must be one of: {list(presets.keys())}")

    p = presets[preset_name]
    result = run_aquifer_score(
        p["grace"], p["well"], p["substrate"], p["quality"],
        p["aquifer_type"], p["extraction_regime"],
    )
    return {
        "preset": preset_name,
        "score": result.score,
        "risk_class": result.risk_class,
        "confidence_interval": result.confidence_interval,
        "depletion_rate_cm_yr": result.depletion_rate_cm_yr,
        "years_to_critical": result.years_to_critical,
        "quality_flags": result.quality_flags,
    }


# ── v1.1 OPTIMIZATION ENDPOINTS ────────────────────────────────────────────


class BatchScoreRequest(BaseModel):
    """Batch aquifer scoring -- score multiple sites in one request."""
    sites: list[ScoreRequest] = Field(
        ...,
        min_length=1,
        max_length=500,
        description="Array of ScoreRequest objects (max 500 per batch)",
    )


@app.post("/v1/score/batch")
async def run_batch_score(request: BatchScoreRequest):
    """
    Score multiple aquifer sites in a single request.

    Regional groundwater assessments may cover dozens of monitoring
    locations. Batch scoring processes all sites and returns an array
    of results with summary statistics.

    Max 500 sites per request.
    """
    results = []
    errors = []

    for i, site in enumerate(request.sites):
        try:
            result = run_aquifer_score(
                grace=site.grace.model_dump(),
                well=site.well.model_dump(),
                substrate=site.substrate.model_dump(),
                quality=site.quality.model_dump(),
                aquifer_type=site.aquifer_type,
                extraction_regime=site.extraction_regime,
            )
            results.append({
                "index": i,
                "site_id": site.site_id or f"batch-{i}",
                "score": result.score,
                "confidence_interval": result.confidence_interval,
                "confidence_pct": result.confidence_pct,
                "risk_class": result.risk_class,
                "depletion_rate_cm_yr": result.depletion_rate_cm_yr,
                "years_to_critical": result.years_to_critical,
                "satellite_score": result.satellite_score,
                "well_score": result.well_score,
                "substrate_score": result.substrate_score,
                "quality_score": result.quality_score,
                "quality_flags": result.quality_flags,
            })
        except Exception as e:
            errors.append({"index": i, "error": str(e)})

    scores = [r["score"] for r in results]
    summary = {}
    if scores:
        summary = {
            "total_sites": len(request.sites),
            "scored": len(results),
            "errors": len(errors),
            "mean_score": round(sum(scores) / len(scores), 1),
            "min_score": min(scores),
            "max_score": max(scores),
            "secure_count": sum(1 for s in scores if s >= 70),
            "stressed_count": sum(1 for s in scores if 40 <= s < 70),
            "critical_count": sum(1 for s in scores if 20 <= s < 40),
            "failing_count": sum(1 for s in scores if s < 20),
            "quality_flagged_count": sum(1 for r in results if r["quality_flags"]),
        }

    return {
        "results": results,
        "errors": errors,
        "summary": summary,
    }


# ── GRACE Decomposition Endpoint ────────────────────────────────────────────

class GRACEDecomposeRequest(BaseModel):
    """GRACE TWS monthly time series for decomposition."""
    monthly_tws_cm: list[float] = Field(
        ...,
        min_length=24,
        max_length=600,
        description="Monthly TWS anomaly values in cm EWT (min 24 months)",
    )
    trend_window: int = Field(13, ge=3, le=25, description="Moving average window for trend extraction")
    site_id: Optional[str] = None


@app.post("/v1/grace/decompose")
async def grace_decompose(request: GRACEDecomposeRequest):
    """
    Decompose a GRACE TWS time series into trend, seasonal, and residual.

    Uses STL-like decomposition (Cleveland et al. 1990) to extract:
    - Trend: long-term storage change direction
    - Seasonal: annual recharge/discharge cycle
    - Residual: anomalous signals (drought, pumping events)

    Classifies depletion pattern: stable, linear_decline,
    accelerating_decline, seasonal_stress, or recovery.

    Minimum 24 months of data required.
    """
    try:
        decomp = decompose_grace_signal(
            request.monthly_tws_cm,
            trend_window=request.trend_window,
        )
        pattern = classify_depletion_pattern(decomp)
    except ValueError as e:
        raise HTTPException(400, str(e))

    return {
        "site_id": request.site_id,
        "n_months": decomp.n_months,
        "trend_slope_cm_yr": decomp.trend_slope_cm_yr,
        "trend_r_squared": decomp.trend_r_squared,
        "is_accelerating": decomp.is_accelerating,
        "acceleration_rate": decomp.acceleration_rate,
        "seasonal_amplitude_cm": decomp.seasonal_amplitude_cm,
        "depletion_pattern": pattern,
        "trend": decomp.trend,
        "seasonal": decomp.seasonal,
        "residual": decomp.residual,
        "citation": "Cleveland, R.B. et al. (1990) STL: Seasonal-Trend Decomposition. J. Official Statistics 6:3-73.",
    }


# ── Kriging Well Field Interpolation Endpoint ───────────────────────────────

class WellPoint(BaseModel):
    lat: float = Field(..., ge=-90, le=90)
    lon: float = Field(..., ge=-180, le=180)
    depth_m: float = Field(..., ge=0, description="Depth to water table in meters")


class KrigingRequest(BaseModel):
    """Kriging interpolation request for a target location."""
    wells: list[WellPoint] = Field(
        ...,
        min_length=2,
        max_length=500,
        description="Known well data points (min 2 required for kriging)",
    )
    target_lat: float = Field(..., ge=-90, le=90)
    target_lon: float = Field(..., ge=-180, le=180)
    max_distance_km: float = Field(100.0, ge=1, le=500)
    site_id: Optional[str] = None


@app.post("/v1/wells/interpolate")
async def wells_interpolate(request: KrigingRequest):
    """
    Predict groundwater depth at a target location from surrounding wells.

    Uses Ordinary Kriging (Cressie 1993) with a spherical variogram model
    to provide the Best Linear Unbiased Predictor (BLUP) for spatial data.

    Returns prediction with uncertainty (kriging variance) and 95% CI.
    """
    wells = [{"lat": w.lat, "lon": w.lon, "depth_m": w.depth_m} for w in request.wells]

    try:
        result = interpolate_well_field(
            wells,
            target_lat=request.target_lat,
            target_lon=request.target_lon,
            max_distance_km=request.max_distance_km,
        )
    except ValueError as e:
        raise HTTPException(400, str(e))

    return {
        "site_id": request.site_id,
        "target_lat": request.target_lat,
        "target_lon": request.target_lon,
        "predicted_depth_m": result["predicted_depth_m"],
        "prediction_variance": result["prediction_variance"],
        "prediction_std_m": result["prediction_std_m"],
        "confidence_interval_m": result["confidence_interval_m"],
        "n_wells_used": result["n_wells_used"],
        "variogram_params": result["variogram_params"],
        "citation": "Cressie, N.A.C. (1993) Statistics for Spatial Data. Wiley.",
    }


# ── Cache Endpoints ─────────────────────────────────────────────────────────

@app.get("/v1/cache/stats")
async def cache_stats():
    """Return cache performance statistics."""
    return _cache.stats()


@app.post("/v1/cache/clear")
async def cache_clear():
    """Clear the response cache. Returns count of evicted entries."""
    count = _cache.clear()
    return {"cleared": count, "status": "ok"}


# ── v2.0 SCIENTIFIC ENDPOINTS ─────────────────────────────────────────────


class WQIRequest(BaseModel):
    """Standalone CCME Water Quality Index request."""
    nitrate_mg_l: float = Field(5.0, ge=0)
    arsenic_ug_l: float = Field(3.0, ge=0)
    fluoride_mg_l: float = Field(0.5, ge=0)
    tds_mg_l: float = Field(300.0, ge=0)
    ph: float = Field(7.2, ge=0, le=14)
    e_coli_cfu_100ml: float = Field(0.0, ge=0)


@app.post("/v1/wqi")
async def run_ccme_wqi(request: WQIRequest):
    """
    Compute CCME Water Quality Index from chemical/microbial measurements.

    The CCME WQI is the internationally recognized standard for drinking
    water quality assessment, combining three factors:
      F1 (Scope):     % of parameters exceeding guidelines
      F2 (Frequency): % of individual tests failing
      F3 (Amplitude): degree of exceedance (asymptotic to 100)

      WQI = 100 - (sqrt(F1² + F2² + F3²) / 1.732)

    Categories:
      EXCELLENT (95-100), GOOD (80-94), FAIR (65-79),
      MARGINAL (45-64), POOR (0-44)

    Citation: CCME (2001) Canadian Water Quality Guidelines for the
              Protection of Aquatic Life. CCME Water Quality Index 1.0.
    """
    measurements = {
        "nitrate_mg_l": request.nitrate_mg_l,
        "arsenic_ug_l": request.arsenic_ug_l,
        "fluoride_mg_l": request.fluoride_mg_l,
        "tds_mg_l": request.tds_mg_l,
        "ph": request.ph,
        "e_coli_cfu_100ml": request.e_coli_cfu_100ml,
    }

    result = compute_ccme_wqi(measurements)

    return {
        "wqi": result.wqi,
        "category": result.category,
        "f1_scope": result.f1_scope,
        "f2_frequency": result.f2_frequency,
        "f3_amplitude": result.f3_amplitude,
        "n_parameters": result.n_variables,
        "n_failed": result.n_failed_variables,
        "exceedances": result.exceedances,
        "interpretation": result.interpretation,
        "guidelines_used": CCME_GUIDELINES,
        "citation": "CCME (2001) Canadian Water Quality Guidelines. CCME WQI 1.0.",
    }


class DRASTICRequest(BaseModel):
    """Standalone DRASTIC vulnerability assessment request."""
    aquifer_type: str
    depth_to_water_m: float = Field(15.0, ge=0, description="Depth to water table in meters")
    hydraulic_conductivity_m_s: float = Field(1e-5, ge=0, description="Hydraulic conductivity in m/s")

    @field_validator("aquifer_type")
    @classmethod
    def validate_aquifer_type(cls, v):
        if v not in AQUIFER_TYPES:
            raise ValueError(f"Invalid aquifer_type. Must be one of: {list(AQUIFER_TYPES.keys())}")
        return v


@app.post("/v1/drastic")
async def run_drastic_assessment(request: DRASTICRequest):
    """
    Compute DRASTIC groundwater vulnerability index.

    DRASTIC is the EPA-published framework for evaluating intrinsic
    groundwater vulnerability using seven hydrogeological parameters:
      D = Depth to water (weight 5)
      R = Net Recharge (weight 4)
      A = Aquifer media (weight 3)
      S = Soil media (weight 2)
      T = Topography/slope (weight 1)
      I = Impact of vadose zone (weight 5)
      C = Conductivity of aquifer (weight 3)

    Index range: 23-230.
    Converted to security score: 100 × (1 - (DI - 23) / 207)

    Citation: Aller, L. et al. (1987) DRASTIC: A Standardized System
              to Evaluate Ground Water Pollution Potential. US EPA/600/2-87/035.
    """
    result = compute_drastic_from_aquifer_type(
        aquifer_type=request.aquifer_type,
        depth_m=request.depth_to_water_m,
        conductivity_m_s=request.hydraulic_conductivity_m_s,
    )

    return {
        "drastic_index": result.drastic_index,
        "vulnerability_class": result.vulnerability_class,
        "security_score": result.security_score,
        "parameter_ratings": result.parameter_ratings,
        "dominant_factor": result.dominant_factor,
        "interpretation": result.interpretation,
        "citation": "Aller, L. et al. (1987) DRASTIC. US EPA/600/2-87/035.",
    }