app.py

"""
AQI Intelligence Engine — Hugging Face Gradio + ZeroGPU Main Application Entry Point

Supports both standalone Gradio Web Interface and REST API backend.
Runs seamlessly on Hugging Face ZeroGPU spaces as well as standard CPU fallback.
"""

import os
import sys
import io
import time
import json
import logging
from PIL import Image
import pandas as pd
import gradio as gr
import gradio_client.utils as client_utils

# Monkey-patch Gradio Client bug where boolean OpenAPI schema fields cause TypeError
if hasattr(client_utils, "_json_schema_to_python_type"):
    _orig_json_schema_to_python_type = client_utils._json_schema_to_python_type
    def safe_json_schema_to_python_type(schema, defs=None):
        if isinstance(schema, bool):
            return "Any"
        try:
            return _orig_json_schema_to_python_type(schema, defs)
        except Exception:
            return "Any"
    client_utils._json_schema_to_python_type = safe_json_schema_to_python_type

import asyncio
from datetime import datetime, timezone

try:
    import spaces
except ImportError:
    spaces = None

def gpu_decorator(func=None, duration=None):
    """ZeroGPU decorator wrapper that applies spaces.GPU on HF or identity on CPU."""
    if spaces is not None:
        try:
            if func is None:
                return spaces.GPU(duration=duration) if duration else spaces.GPU
            if callable(func):
                return spaces.GPU(func)
            return spaces.GPU
        except Exception:
            pass
    if func is None:
        return lambda f: f
    if callable(func):
        return func
    return lambda f: f

from config import get_device, HAS_SPACES
from services.forecast.service import forecast_aqi
from services.vision.service import detect_pollution_sources

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("app_hf")

# Default sample AQI history for quick testing
DEFAULT_AQI_HIST = "145, 150, 162, 178, 185, 210, 240, 260, 255, 230, 215, 195, 180, 175, 182, 198, 220, 245, 270, 290, 310, 295, 280, 260"


async def _async_handle_forecast_ui(aqi_text_input: str, lat_val: str, lon_val: str):
    """Async core logic for AQI Forecasting."""
    try:
        if not aqi_text_input or not aqi_text_input.strip():
            return "❌ Error: Please provide at least 24 hourly AQI readings.", None, ""
        
        # Parse CSV string into list of floats
        raw_vals = [x.strip() for x in aqi_text_input.split(",") if x.strip()]
        aqi_series = [float(x) for x in raw_vals]

        lat = float(lat_val) if lat_val and lat_val.strip() else None
        lon = float(lon_val) if lon_val and lon_val.strip() else None

        start_time = time.time()
        result = await forecast_aqi(aqi_series, lat=lat, lon=lon)
        elapsed = round(time.time() - start_time, 2)

        # Build summary text
        summary = (
            f"### 🔮 Forecast Completed in {elapsed}s\n"
            f"- **Active Hardware**: `{result.get('device_used', get_device())}`\n"
            f"- **Model Used**: `{result.get('model_type', 'Google TimesFM 2.5 / XGBoost')}`\n"
            f"- **24-Hour Horizon Avg AQI**: `{result.get('forecast_24h', {}).get('mean', 'N/A')}`\n"
            f"- **48-Hour Horizon Avg AQI**: `{result.get('forecast_48h', {}).get('mean', 'N/A')}`\n"
            f"- **72-Hour Horizon Avg AQI**: `{result.get('forecast_72h', {}).get('mean', 'N/A')}`\n"
        )

        # Build DataFrame for plot table
        hourly_preds = result.get("hourly_predictions", [])
        if hourly_preds:
            df = pd.DataFrame(hourly_preds)
        else:
            df = pd.DataFrame({"step": list(range(1, len(aqi_series)+1)), "aqi": aqi_series})

        formatted_json = json.dumps(result, indent=2)
        return summary, df, formatted_json

    except Exception as e:
        logger.error(f"Forecast UI Error: {e}")
        return f"❌ Execution Failed: {str(e)}", None, ""


@gpu_decorator
def handle_forecast_ui(aqi_text_input: str, lat_val: str, lon_val: str):
    """Gradio handler for AQI Forecasting wrapped with @gpu_decorator."""
    return asyncio.run(_async_handle_forecast_ui(aqi_text_input, lat_val, lon_val))


async def _async_handle_vision_ui(input_image):
    """Async core logic for Satellite Vision Pollution Detection."""
    try:
        if input_image is None:
            return "❌ Error: Please upload a satellite or aerial image.", ""

        start_time = time.time()
        pil_image = Image.fromarray(input_image) if not isinstance(input_image, Image.Image) else input_image

        result = await detect_pollution_sources(pil_image)
        elapsed = round(time.time() - start_time, 2)

        scene_desc = result.get("scene_description", "No detailed description.")
        sources = result.get("pollution_sources_found", [])
        counts = result.get("source_count", {})
        severity = result.get("severity", "unknown").upper()

        sources_str = ", ".join(sources) if sources else "None identified"

        summary = (
            f"### 🛰️ Vision Analysis Complete ({elapsed}s)\n"
            f"- **Device Active**: `{get_device()}`\n"
            f"- **Severity Level**: **`{severity}`**\n"
            f"- **Identified Sources**: `{sources_str}`\n\n"
            f"#### 📝 Scene Description:\n{scene_desc}\n\n"
            f"#### 📊 Category Counts:\n```json\n{json.dumps(counts, indent=2)}\n```"
        )

        formatted_json = json.dumps(result, indent=2)
        return summary, formatted_json

    except Exception as e:
        logger.error(f"Vision UI Error: {e}")
        return f"❌ Analysis Failed: {str(e)}", ""


@gpu_decorator
def handle_vision_ui(input_image):
    """Gradio handler for Satellite Vision wrapped with @gpu_decorator."""
    return asyncio.run(_async_handle_vision_ui(input_image))


# Build Gradio Blocks Interface
with gr.Blocks(title="AQI Intelligence Engine — HF ZeroGPU Edition", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🌬️ AQI Intelligence Engine (ZeroGPU & CPU Compatible)
        ### AI-Powered Urban Air Quality Forecasting & High-Resolution Satellite Vision Intelligence
        ---
        """
    )

    with gr.Tabs():
        # TAB 1: FORECASTING
        with gr.TabItem("📈 AQI Forecasting (TimesFM 2.5 & XGBoost)"):
            gr.Markdown("Input historical hourly AQI readings to generate 24h, 48h, and 72h predictive forecasts.")
            with gr.Row():
                with gr.Column(scale=1):
                    aqi_in = gr.Textbox(
                        label="Historical Hourly AQI Readings (CSV format, min 24 hours)",
                        value=DEFAULT_AQI_HIST,
                        lines=4,
                    )
                    lat_in = gr.Textbox(label="Latitude (Optional, e.g. 28.6139)", value="28.6139")
                    lon_in = gr.Textbox(label="Longitude (Optional, e.g. 77.2090)", value="77.2090")
                    run_forecast_btn = gr.Button("🚀 Run Forecast", variant="primary")

                with gr.Column(scale=1):
                    forecast_summary_out = gr.Markdown(label="Forecast Overview")
                    forecast_table_out = gr.Dataframe(label="Forecast Step Matrix")

            with gr.Accordion("🔍 Raw JSON Response Payload", open=False):
                forecast_json_out = gr.Code(language="json")

            run_forecast_btn.click(
                fn=handle_forecast_ui,
                inputs=[aqi_in, lat_in, lon_in],
                outputs=[forecast_summary_out, forecast_table_out, forecast_json_out],
            )

        # TAB 2: SATELLITE VISION
        with gr.TabItem("🛰️ Satellite Vision AI (Florence-2 + Grounding DINO)"):
            gr.Markdown("Upload satellite/aerial imagery to detect air pollution sources, smoke plumes, kiln clusters, and industrial emissions.")
            with gr.Row():
                with gr.Column(scale=1):
                    image_in = gr.Image(label="Upload Satellite / Aerial Image", type="numpy")
                    analyze_img_btn = gr.Button("🔍 Analyze Image", variant="primary")

                with gr.Column(scale=1):
                    vision_summary_out = gr.Markdown(label="Vision Intelligence Findings")

            with gr.Accordion("🔍 Detailed Detections JSON", open=False):
                vision_json_out = gr.Code(language="json")

            analyze_img_btn.click(
                fn=handle_vision_ui,
                inputs=[image_in],
                outputs=[vision_summary_out, vision_json_out],
            )

        # TAB 3: HARDWARE & SYSTEM STATUS
        with gr.TabItem("⚡ Hardware & API Status"):
            device_active = get_device()
            gr.Markdown(
                f"""
                ### 🖥️ Active Deployment Info
                - **Primary Compute Hardware**: `{device_active}`
                - **HuggingFace `spaces` SDK Available**: `{spaces is not None}`
                - **ZeroGPU Status**: `{"Active & Ready for dynamic GPU acceleration" if spaces is not None else "Running on standard CPU mode (Universal Compatibility)"}`
                
                ### 🔌 REST API Endpoints
                When deployed on Hugging Face Spaces, you can query REST endpoints directly:
                - `POST /forecast`: Direct historical AQI prediction payload
                - `POST /analyze-image`: Satellite base64 image breakdown
                - `GET /health`: Core server state & device telemetry
                """
            )

from starlette.responses import JSONResponse
from starlette.requests import Request

# Initialize the Gradio queue — this creates the underlying FastAPI app object
# MUST be called before registering REST routes on demo.app
demo.queue()


def register_api_routes(gradio_app):
    """Register custom REST API endpoints on the Gradio FastAPI app."""

    @gradio_app.post("/forecast")
    async def api_forecast(request: Request):
        """REST API endpoint for historical AQI prediction."""
        try:
            payload = await request.json()
            from services.forecast.service import forecast_aqi
            aqi_hist = payload.get("aqi_history", [])
            lat = payload.get("lat")
            lon = payload.get("lon")
            res = await forecast_aqi(aqi_hist, lat=lat, lon=lon)
            return JSONResponse(res)
        except Exception as e:
            logger.error(f"API Forecast error: {e}")
            return JSONResponse({"error": str(e)}, status_code=400)

    @gradio_app.post("/analyze-image")
    async def api_analyze_image(request: Request):
        """REST API endpoint for satellite image pollution breakdown."""
        try:
            payload = await request.json()
            import base64
            from services.vision.service import detect_pollution_sources
            base64_str = payload.get("image_base64", "")
            if not base64_str:
                return JSONResponse({"error": "No image_base64 provided"}, status_code=400)
            if "," in base64_str:
                base64_str = base64_str.split(",", 1)[1]
            img_bytes = base64.b64decode(base64_str)
            pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
            res = await detect_pollution_sources(pil_img)
            return JSONResponse(res)
        except Exception as e:
            logger.error(f"API Vision error: {e}")
            return JSONResponse({"error": str(e)}, status_code=400)

    @gradio_app.get("/health")
    async def api_health(request: Request):
        """Health check endpoint."""
        return JSONResponse({
            "status": "ok",
            "device": get_device(),
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "version": "2.0.0",
        })


from fastapi import FastAPI

# Create a standard FastAPI application
app = FastAPI(
    title="AQI Intelligence Engine API",
    description="REST API for AQI Forecasting and Satellite Vision Analysis",
    version="2.0.0",
)

# Register REST API routes on the FastAPI app before mounting Gradio
try:
    from fastapi.middleware.cors import CORSMiddleware
    # Explicitly configure CORS on the FastAPI app to allow direct requests from the browser
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    register_api_routes(app)
    logger.info("CORS middleware and REST API routes registered successfully on FastAPI app.")
except Exception as e:
    logger.warning(f"Could not register REST API routes or CORS: {e}")

# Mount the Gradio Blocks app onto the FastAPI application at the root
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    # Local development only — on HF Spaces, uvicorn runs the app directly
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)