app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pandas as pd
import numpy as np
import torch
from pytorch_tabnet.tab_model import TabNetRegressor
from contextlib import asynccontextmanager
import os
import requests
from pathlib import Path

# Define the data schema for Window 7 Conductor requests
class EconRequest(BaseModel):
    virus_name: str
    econ_buffer: float  # Value from 0.1 to 1.0 (IMF-WEO)
    population_exposure: float # Value from 100 to 100,000 (WorldPop)

# Global container for the model
aegis_brain = {}

def download_model():
    """Download the model from Hugging Face with authentication"""
    model_path = "aegis_window2_econ_v1.zip"
    model_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/aegis_window2_econ_v1.zip"
    
    # Get HF token from environment variable
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        print("⚠️ No HF_TOKEN environment variable found. Using public access only.")
        return False
    
    headers = {"Authorization": f"Bearer {hf_token}"}
    
    if not os.path.exists(model_path):
        print("📥 Downloading model from Hugging Face...")
        try:
            response = requests.get(model_url, headers=headers, stream=True)
            response.raise_for_status()
            
            with open(model_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            print("✅ Model downloaded successfully!")
            return True
        except Exception as e:
            print(f"❌ Failed to download model: {e}")
            return False
    else:
        print("✅ Model file already exists")
        return True

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Download and load model on startup
    model_loaded = False
    
    if download_model():
        try:
            model_path = "aegis_window2_econ_v1.zip"
            regressor = TabNetRegressor()
            regressor.load_model(model_path)
            aegis_brain["model"] = regressor
            model_loaded = True
            print("✅ TabNet Model Loaded Successfully!")
        except Exception as e:
            print(f"⚠️ Failed to load TabNet model: {e}")
            print("🔄 Falling back to simulation mode...")
    
    # Load risk data for cross-referencing
    try:
        rankings_url = "https://huggingface.co/gsstec/aegis_window2_econ_v1/resolve/main/Rankings.csv"
        hf_token = os.getenv("HF_TOKEN")
        
        if hf_token:
            headers = {"Authorization": f"Bearer {hf_token}"}
            response = requests.get(rankings_url, headers=headers)
            if response.status_code == 200:
                from io import StringIO
                aegis_brain["rankings"] = pd.read_csv(StringIO(response.text))
                print("✅ Rankings data loaded from Hugging Face")
            else:
                raise Exception(f"Rankings CSV request failed: {response.status_code}")
        else:
            raise Exception("No HF token available")
    except Exception as e:
        print(f"⚠️ Could not load rankings from HF: {e}")
        # Create comprehensive mock rankings data
        aegis_brain["rankings"] = pd.DataFrame({
            'Virus Name': [
                'COVID-19', 'H1N1', 'SARS', 'MERS', 'Ebola', 'Zika', 
                'Influenza A', 'RSV', 'Marburg', 'Lassa', 'Nipah', 'Hendra'
            ],
            'Original Score': [
                0.85, 0.65, 0.75, 0.55, 0.95, 0.45, 
                0.60, 0.40, 0.90, 0.70, 0.80, 0.75
            ]
        })
        print("✅ Comprehensive Mock Rankings Data Loaded")
    
    aegis_brain["model_loaded"] = model_loaded
    yield
    aegis_brain.clear()

app = FastAPI(lifespan=lifespan, title="Aegis Econ API")

@app.post("/predict")
async def get_stability_score(data: EconRequest):
    model = aegis_brain.get("model")
    df_rankings = aegis_brain.get("rankings")
    model_loaded = aegis_brain.get("model_loaded", False)
    
    if df_rankings is None:
        raise HTTPException(status_code=500, detail="Rankings data not initialized.")

    # 1. Fetch risk score from Rankings.csv
    virus_row = df_rankings[df_rankings['Virus Name'] == data.virus_name]
    if virus_row.empty:
        # If virus not found, use a default moderate risk score
        base_score = 0.7
        print(f"⚠️ Virus '{data.virus_name}' not found in database. Using default risk score.")
    else:
        base_score = virus_row.iloc[0]['Original Score']
    
    risk_impact = base_score * 0.7
    
    # 2. Predict Continental Stability
    if model_loaded and model is not None:
        # Use actual TabNet model
        input_vector = np.array([[base_score, risk_impact, data.econ_buffer, data.population_exposure]])
        prediction = model.predict(input_vector)
        stability_score = float(prediction[0][0])
        model_status = "trained_model"
    else:
        # Enhanced mathematical simulation
        econ_factor = data.econ_buffer * 0.35
        population_factor = min(data.population_exposure / 100000, 1.0) * 0.35
        virus_factor = (1 - base_score) * 0.30
        
        # More sophisticated calculation
        stability_score = max(0.0, min(1.0, 
            econ_factor + (1 - population_factor) + (1 - virus_factor)
        ))
        
        # Add realistic variance based on input parameters
        variance = 0.03 + (data.population_exposure / 1000000) * 0.02
        stability_score += np.random.normal(0, variance)
        stability_score = max(0.0, min(1.0, stability_score))
        model_status = "enhanced_simulation"
    
    # Determine alert level with more nuanced thresholds
    if stability_score < 0.25:
        alert_level = "CRITICAL"
    elif stability_score < 0.45:
        alert_level = "HIGH_RISK"
    elif stability_score < 0.65:
        alert_level = "MONITOR"
    else:
        alert_level = "STABLE"
    
    return {
        "virus": data.virus_name,
        "stability_score": round(stability_score, 4),
        "alert_level": alert_level,
        "model_status": model_status,
        "base_risk_score": round(base_score, 3),
        "economic_buffer_impact": round(data.econ_buffer * 0.35, 3),
        "population_risk_factor": round(min(data.population_exposure / 100000, 1.0), 3)
    }

@app.get("/health")
async def health():
    return {
        "status": "operational", 
        "hardware": "T4 GPU Active" if torch.cuda.is_available() else "CPU Mode",
        "model_loaded": aegis_brain.get("model_loaded", False),
        "available_viruses": aegis_brain.get("rankings", pd.DataFrame())['Virus Name'].tolist() if aegis_brain.get("rankings") is not None else [],
        "total_virus_database": len(aegis_brain.get("rankings", pd.DataFrame())),
        "hf_token_available": bool(os.getenv("HF_TOKEN"))
    }

@app.get("/")
async def root():
    return {
        "message": "AEGIS Economic Stability Analysis API",
        "version": "1.2.0",
        "model_source": "https://huggingface.co/gsstec/aegis_window2_econ_v1",
        "features": [
            "Economic stability prediction",
            "Multi-virus risk assessment", 
            "Population exposure modeling",
            "Real-time alert classification"
        ],
        "endpoints": {
            "predict": "POST /predict - Get stability prediction",
            "health": "GET /health - System status",
            "root": "GET / - API information"
        }
    }