api.py

"""
NUSARA Mining Inference REST API - FastAPI Edition
Wraps the existing InferencePipeline (inference_pipeline.py) as HTTP endpoints.
Designed to run locally, in Docker, or on a HuggingFace Space (port 7860).

Run locally:
    uvicorn api:app --host 0.0.0.0 --port 7860

Interactive docs:
    http://localhost:7860/docs
"""

import os
import logging
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Optional

import pandas as pd
from fastapi import FastAPI, Header, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from inference_pipeline import InferencePipeline, MODEL_DIR, DATA_DIR

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")
logger = logging.getLogger(__name__)

API_KEY = os.getenv("API_KEY")  # if set, requests must send header X-API-Key

# Holds the single shared pipeline instance (models loaded once at startup)
state: Dict[str, Any] = {"pipeline": None}


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Loading inference pipeline (models)...")
    state["pipeline"] = InferencePipeline(model_dir=MODEL_DIR)
    loaded = list(state["pipeline"].models.keys())
    logger.info(f"Pipeline ready. Models loaded: {loaded or 'NONE'}")
    yield
    state["pipeline"] = None


app = FastAPI(
    title="NUSARA Mining Inference API",
    description=(
        "REST API for mining ML inference: equipment failure, predictive "
        "maintenance, cost anomaly detection, and what-if simulation."
    ),
    version="1.0.0",
    lifespan=lifespan,
)

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


# ── Auth (optional) ──────────────────────────────────────────────────

def require_api_key(x_api_key: Optional[str] = Header(default=None)):
    """If API_KEY env is set, require a matching X-API-Key header."""
    if API_KEY and x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid or missing X-API-Key")
    return True


# ── Request / Response models ────────────────────────────────────────

class PredictRequest(BaseModel):
    """Optional input. If 'records' is omitted, the bundled feature CSV is used."""
    records: Optional[List[Dict[str, Any]]] = Field(
        default=None,
        description="List of row objects. If null, the pipeline reads the bundled CSV.",
    )

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "records": [
                        {
                            "date": "2025-05-01",
                            "site_name": "Tambang Nikel E",
                            "equipment_name": "Bulldozer BD-002",
                            "equipment_type": "Bulldozer",
                            "operating_hours": 286.0,
                            "downtime_hours": 23.0,
                            "fuel_consumption": 1667.0,
                            "maintenance_cost": 849403.0,
                        }
                    ]
                }
            ]
        }
    }


class AnalysisRequest(BaseModel):
    """Request for analysis modules with optional parameters"""
    days_back: int = Field(default=90, description="Number of days to analyze")
    laziness_level: int = Field(default=5, description="Simulated laziness level (0-10) for human factor analysis")
    num_records_per_site: int = Field(default=5, description="Number of GIS road records per site")

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "days_back": 90,
                    "laziness_level": 5,
                    "num_records_per_site": 5
                }
            ]
        }
    }


def _get_pipeline() -> InferencePipeline:
    pipeline = state.get("pipeline")
    if pipeline is None:
        raise HTTPException(status_code=503, detail="Pipeline not initialized")
    return pipeline


def _to_df(req: Optional[PredictRequest]) -> Optional[pd.DataFrame]:
    if req is None or not req.records:
        return None
    try:
        return pd.DataFrame(req.records)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid records payload: {e}")


# ── Info endpoints ───────────────────────────────────────────────────

@app.get("/")
def root():
    pipeline = state.get("pipeline")
    models = list(pipeline.models.keys()) if pipeline else []
    return {
        "name": "NUSARA Mining Inference API",
        "version": "1.0.0",
        "docs": "/docs",
        "models_loaded": models,
        "endpoints": [
            "GET /health",
            "GET /models",
            "POST /predict/equipment-failure",
            "POST /predict/maintenance-priority",
            "POST /predict/cost-anomaly",
            "POST /predict/whatif-simulation",
            "POST /predict/fleet-optimization",
            "POST /predict/weather-analysis",
            "POST /predict/gis-roads-analysis",
            "POST /predict/human-factors-analysis",
            "POST /predict/all",
        ],
    }


@app.get("/health")
def health():
    pipeline = state.get("pipeline")
    models = list(pipeline.models.keys()) if pipeline else []
    return {
        "status": "ok" if pipeline else "initializing",
        "models_loaded": models,
        "model_count": len(models),
        "model_dir": MODEL_DIR,
        "data_dir": DATA_DIR,
    }


@app.get("/models")
def list_models():
    pipeline = _get_pipeline()

    def kind(name: str) -> str:
        if name.startswith("equipment_failure"):
            return "equipment_failure_prediction"
        if name.startswith("maintenance_model"):
            return "predictive_maintenance"
        if name.startswith("cost_anomaly"):
            return "cost_anomaly_detection"
        if name.startswith("whatif"):
            return "whatif_simulation"
        if 'fleet_optimization' in name or name.startswith('fleet_'):
            return "fleet_optimization"
        return "unknown"

    return {
        "model_count": len(pipeline.models),
        "models": [{"name": n, "type": kind(n)} for n in pipeline.models],
    }


# ── Prediction endpoints ─────────────────────────────────────────────

@app.post("/predict/equipment-failure", dependencies=[Depends(require_api_key)])
def predict_equipment_failure(req: Optional[PredictRequest] = None):
    return _get_pipeline().predict_equipment_failure(_to_df(req))


@app.post("/predict/maintenance-priority", dependencies=[Depends(require_api_key)])
def predict_maintenance_priority(req: Optional[PredictRequest] = None):
    return _get_pipeline().predict_maintenance_priority(_to_df(req))


@app.post("/predict/cost-anomaly", dependencies=[Depends(require_api_key)])
def predict_cost_anomaly(req: Optional[PredictRequest] = None):
    return _get_pipeline().detect_cost_anomalies(_to_df(req))


@app.post("/predict/whatif-simulation", dependencies=[Depends(require_api_key)])
def predict_whatif_simulation(req: Optional[PredictRequest] = None):
    return _get_pipeline().run_whatif_simulation(_to_df(req))


@app.post("/predict/fleet-optimization", dependencies=[Depends(require_api_key)])
def predict_fleet_optimization(req: Optional[PredictRequest] = None):
    return _get_pipeline().predict_fleet_optimization(_to_df(req))


@app.post("/predict/weather-analysis", dependencies=[Depends(require_api_key)])
def predict_weather_analysis(req: Optional[AnalysisRequest] = None):
    """Analyze BMKG weather impact on mining operations"""
    days_back = req.days_back if req else 90
    return {"model": "weather_analysis", "results": _get_pipeline().analyze_weather_impact(days_back)}


@app.post("/predict/gis-roads-analysis", dependencies=[Depends(require_api_key)])
def predict_gis_roads_analysis(req: Optional[AnalysisRequest] = None):
    """Analyze GIS road conditions impact on logistics"""
    days_back = req.days_back if req else 90
    num_records = req.num_records_per_site if req else 5
    return {"model": "gis_roads_analysis", "results": _get_pipeline().analyze_gis_roads(days_back, num_records)}


@app.post("/predict/human-factors-analysis", dependencies=[Depends(require_api_key)])
def predict_human_factors_analysis(req: Optional[AnalysisRequest] = None):
    """Analyze simulated human factor (laziness) impact on operations"""
    days_back = req.days_back if req else 90
    laziness_level = req.laziness_level if req else 5
    return {"model": "human_factors_analysis", "results": _get_pipeline().analyze_human_factors(laziness_level, days_back)}


@app.post("/predict/all", dependencies=[Depends(require_api_key)])
def predict_all(req: Optional[AnalysisRequest] = None):
    """Run every model using the bundled feature CSVs and return a consolidated report."""
    days_back = req.days_back if req else 90
    laziness_level = req.laziness_level if req else 5
    return _get_pipeline().run_full_pipeline(input_dir=DATA_DIR, days_back=days_back, laziness_level=laziness_level)


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "7860"))
    uvicorn.run("api:app", host="0.0.0.0", port=port)