main.py

"""
FastAPI Backend for Virus Prediction System
Optimized for Hugging Face Spaces Free Tier with MongoDB Atlas
Version: 1.0.1
"""

from fastapi import FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any
from datetime import datetime
import logging

# Model and prediction imports
from model_handler import (
    get_virus_predictor,
    refresh_virus_mappings,
    VIRUS_MAPPING,
    OTHER_VIRUS_MAPPING,
    ALL_SYMPTOMS
)
from location_mappings import LocationMappingService

# Database imports
from data_handler import save_prediction_to_db, save_validation_to_db, get_db_health, get_prediction_stats

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="Virus Prediction API",
    description="AI-powered viral infection prediction system",
    version="1.0.0",
    docs_url="/",  # Swagger UI at root
    redoc_url="/redoc"
)

# CORS middleware for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update with specific origins in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global predictor instance (loaded on startup)
predictor = None
location_mapping_service = LocationMappingService()


def _normalize_location_name(value: Optional[str]) -> Optional[str]:
    """Return a trimmed location label or None when empty."""
    if value is None:
        return None

    cleaned = value.strip()
    return cleaned or None


def _resolve_location_names(patient_dict: Dict[str, Any]) -> tuple[Optional[str], Optional[str]]:
    """Resolve human-readable state and district names for persistence."""
    explicit_state = _normalize_location_name(patient_dict.get("state_name"))
    explicit_district = _normalize_location_name(patient_dict.get("district_name"))

    resolved_state = explicit_state or location_mapping_service.get_state_name(
        patient_dict.get("labstate"),
        predictor=predictor,
    )
    resolved_district = explicit_district or location_mapping_service.get_district_name(
        patient_dict.get("districtencoded"),
        state_name=resolved_state,
        state_code=patient_dict.get("labstate"),
        predictor=predictor,
    )

    return resolved_state, resolved_district

# ============================================================================
# Pydantic Models for Request/Response
# ============================================================================

class PatientData(BaseModel):
    """Patient information and symptoms"""
    # Demographics
    age: float = Field(..., ge=0, le=120, description="Patient age in years (decimals for months)")
    SEX: int = Field(..., ge=0, le=1, description="0=Female, 1=Male")
    PATIENTTYPE: int = Field(..., ge=0, le=1, description="0=Outpatient, 1=Inpatient")
    durationofillness: int = Field(..., ge=0, le=365, description="Duration of illness in days")
    
    # Location
    labstate: int = Field(..., description="Encoded state value")
    districtencoded: int = Field(..., description="Encoded district value")
    state_name: Optional[str] = Field(None, description="Human-readable state name from frontend")
    district_name: Optional[str] = Field(None, description="Human-readable district name from frontend")
    
    # Temporal
    month: int = Field(..., ge=1, le=12, description="Month of illness (1-12)")
    year: int = Field(..., ge=2012, le=2030, description="Year of illness")
    
    # Syndrome
    syndrome: int = Field(..., ge=1, le=19, description="Primary syndrome classification")
    syndrome_name: Optional[str] = Field(None, description="Syndrome name")
    other_syndrome_specification: Optional[str] = Field("", description="Specification for 'Other' syndrome")
    
    # Symptoms (all binary 0/1)
    HEADACHE: int = Field(0, ge=0, le=1)
    IRRITABILITY: int = Field(0, ge=0, le=1)
    ALTEREDSENSORIUM: int = Field(0, ge=0, le=1)
    SOMNOLENCE: int = Field(0, ge=0, le=1)
    NECKRIGIDITY: int = Field(0, ge=0, le=1)
    SEIZURES: int = Field(0, ge=0, le=1)
    DIARRHEA: int = Field(0, ge=0, le=1)
    DYSENTERY: int = Field(0, ge=0, le=1)
    NAUSEA: int = Field(0, ge=0, le=1)
    VOMITING: int = Field(0, ge=0, le=1)
    ABDOMINALPAIN: int = Field(0, ge=0, le=1)
    MALAISE: int = Field(0, ge=0, le=1)
    MYALGIA: int = Field(0, ge=0, le=1)
    ARTHRALGIA: int = Field(0, ge=0, le=1)
    CHILLS: int = Field(0, ge=0, le=1)
    RIGORS: int = Field(0, ge=0, le=1)
    FEVER: int = Field(0, ge=0, le=1)
    BREATHLESSNESS: int = Field(0, ge=0, le=1)
    COUGH: int = Field(0, ge=0, le=1)
    RHINORRHEA: int = Field(0, ge=0, le=1)
    SORETHROAT: int = Field(0, ge=0, le=1)
    BULLAE: int = Field(0, ge=0, le=1)
    PAPULARRASH: int = Field(0, ge=0, le=1)
    PUSTULARRASH: int = Field(0, ge=0, le=1)
    MUSCULARRASH: int = Field(0, ge=0, le=1)
    MACULOPAPULARRASH: int = Field(0, ge=0, le=1)
    ESCHAR: int = Field(0, ge=0, le=1)
    DARKURINE: int = Field(0, ge=0, le=1)
    HEPATOMEGALY: int = Field(0, ge=0, le=1)
    JAUNDICE: int = Field(0, ge=0, le=1)
    REDEYE: int = Field(0, ge=0, le=1)
    DISCHARGEEYES: int = Field(0, ge=0, le=1)
    CRUSHINGEYES: int = Field(0, ge=0, le=1)
    SWELLINGEYES: int = Field(0, ge=0, le=1)
    RETROORBITALPAIN: int = Field(0, ge=0, le=1)

    class Config:
        json_schema_extra = {
            "example": {
                "age": 30.0,
                "SEX": 1,
                "PATIENTTYPE": 1,
                "durationofillness": 3,
                "labstate": 32,
                "districtencoded": 120,
                "month": 8,
                "year": 2024,
                "syndrome": 5,
                "FEVER": 1,
                "HEADACHE": 1,
                "MYALGIA": 1,
                "ARTHRALGIA": 1
            }
        }


class PredictionResponse(BaseModel):
    """Prediction results"""
    success: bool
    predicted_virus: str
    predicted_virus_id: int
    confidence: float
    top_5_predictions: List[Dict[str, Any]]
    sub_classification: Optional[Dict[str, Any]] = None
    models_info: Dict[str, str]
    timestamp: str
    prediction_id: Optional[str] = None


class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    timestamp: str
    models_loaded: bool
    database_connected: bool


class LocationMappingsResponse(BaseModel):
    """Frontend-safe location encoder configuration."""
    states: List[str]
    districts_by_state: Dict[str, List[str]]
    state_mapping: Dict[str, int]
    district_mapping: Dict[str, int]
    district_mapping_by_state: Dict[str, Dict[str, int]]
    source: str
    timestamp: str
    warnings: List[str] = Field(default_factory=list)


class ValidationRequest(BaseModel):
    """Validation feedback request"""
    prediction_id: str = Field(..., description="MongoDB document ID from prediction response")
    actual_virus_category: str = Field(..., description="'Main' or 'Other' virus category")
    actual_virus_id: int = Field(..., description="Virus ID within the category")
    feedback_notes: Optional[str] = Field("", description="Optional medical professional feedback")
    is_correct: bool = Field(..., description="Whether the prediction was correct")
    
    class Config:
        json_schema_extra = {
            "example": {
                "prediction_id": "507f1f77bcf86cd799439011",
                "actual_virus_category": "Main",
                "actual_virus_id": 1,
                "feedback_notes": "Confirmed Dengue Virus via lab test",
                "is_correct": True
            }
        }


# ============================================================================
# Startup Event
# ============================================================================

@app.on_event("startup")
async def startup_event():
    """Load models and initialize predictor on startup"""
    global predictor
    try:
        logger.info("Loading virus prediction models...")
        refresh_virus_mappings()
        predictor = get_virus_predictor()
        
        if predictor.model1 is None or predictor.model2 is None:
            logger.error("Failed to load models!")
            raise RuntimeError("Model loading failed")
        
        logger.info("Models loaded successfully!")
        logger.info(f"Model 1: {predictor.model1.__class__.__name__}")
        logger.info(f"Model 2: {predictor.model2.__class__.__name__}")

        # Load and cache location mappings for frontend use.
        try:
            location_data = location_mapping_service.load(predictor=predictor, force_reload=True)
            logger.info(
                "Location mappings loaded from '%s' (states=%d, districts=%d)",
                location_data.source,
                len(location_data.state_mapping),
                len(location_data.district_mapping)
            )
            if location_data.warnings:
                logger.warning("Location mapping warnings: %s", "; ".join(location_data.warnings))
        except Exception as mapping_error:
            logger.error("Failed to load location mappings: %s", mapping_error, exc_info=True)
            logger.warning("Application will continue without authoritative location mappings")
        
        # Test database connection
        logger.info("Testing database connection...")
        db_health = get_db_health()
        if db_health.get('status') == 'healthy':
            logger.info("✓ Database connection successful!")
        else:
            logger.warning(f"⚠ Database connection failed: {db_health.get('message', 'Unknown error')}")
            logger.warning("Application will continue but predictions won't be saved to database")
        
    except Exception as e:
        logger.error(f"Startup error: {e}")
        raise


# ============================================================================
# API Endpoints
# ============================================================================

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    db_health = get_db_health()
    
    return HealthResponse(
        status="healthy" if predictor is not None else "unhealthy",
        timestamp=datetime.now().isoformat(),
        models_loaded=predictor is not None and predictor.model1 is not None,
        database_connected=db_health.get("status") == "connected"
    )


@app.get("/mappings")
async def get_mappings():
    """Get virus and symptom mappings"""
    response = {
        "virus_mapping": VIRUS_MAPPING,
        "other_virus_mapping": OTHER_VIRUS_MAPPING,
        "symptoms": ALL_SYMPTOMS,
        "total_major_classes": len(VIRUS_MAPPING),
        "total_other_classes": len(OTHER_VIRUS_MAPPING)
    }

    # Backward-compatible extra location keys for frontend convenience.
    try:
        location_data = location_mapping_service.get(predictor=predictor)
        if location_data.state_mapping:
            response["state_mapping"] = location_data.state_mapping
        if location_data.district_mapping:
            response["district_mapping"] = location_data.district_mapping
        if location_data.district_mapping_by_state:
            response["district_mapping_by_state"] = location_data.district_mapping_by_state
        if location_data.states:
            response["states"] = location_data.states
        if location_data.districts_by_state:
            response["districts_by_state"] = location_data.districts_by_state
    except Exception as mapping_error:
        logger.warning("Could not enrich /mappings with location data: %s", mapping_error)

    return response


@app.get("/location-mappings", response_model=LocationMappingsResponse)
@app.get("/locations", response_model=LocationMappingsResponse)
async def get_location_mappings():
    """Return model-compatible location encoders and state/district options."""
    try:
        data = location_mapping_service.get(predictor=predictor)
        return LocationMappingsResponse(**data.to_response_dict())
    except Exception as e:
        logger.error("Location mapping endpoint error: %s", e, exc_info=True)
        # Keep endpoint resilient for frontend bootstrapping.
        return LocationMappingsResponse(
            states=[],
            districts_by_state={},
            state_mapping={},
            district_mapping={},
            district_mapping_by_state={},
            source="unavailable",
            timestamp=datetime.now().isoformat(),
            warnings=["Location mapping service is unavailable"]
        )


@app.post("/predict", response_model=PredictionResponse)
async def predict_virus(patient_data: PatientData):
    """
    Predict virus from patient data
    
    - Accepts patient demographics and symptoms
    - Returns top 5 predictions with confidence scores
    - Includes sub-classification for "Other Viruses"
    - Saves prediction to MongoDB if available
    """
    if predictor is None:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Models not loaded"
        )
    
    try:
        # Convert Pydantic model to dict
        patient_dict = patient_data.dict()
        
        # Validate at least one symptom is present
        symptoms_present = any(
            patient_dict.get(symptom, 0) == 1 
            for symptom in ALL_SYMPTOMS
        )
        
        if not symptoms_present:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="At least one symptom must be selected"
            )
        
        # Make prediction
        prediction_results = predictor.predict(patient_dict)
        
        y_pred = prediction_results['y_pred']
        y_pred_proba = prediction_results['y_pred_proba']
        top_5_indices = prediction_results['top_5_indices']
        second_model_results = prediction_results['second_model_results']
        
        # Prepare response
        prediction_result = {
            'predicted_virus': VIRUS_MAPPING[y_pred],
            'predicted_virus_id': int(y_pred),
            'confidence': float(y_pred_proba[y_pred] * 100),
            'top_5_predictions': [
                {
                    'virus': VIRUS_MAPPING[idx],
                    'virus_id': int(idx),
                    'confidence': float(y_pred_proba[idx] * 100)
                } for idx in top_5_indices
            ]
        }
        
        # Add sub-classification if available
        sub_classification = None
        if second_model_results:
            sub_classification = {
                'predicted_sub_virus': OTHER_VIRUS_MAPPING[second_model_results['prediction']],
                'predicted_sub_virus_id': int(second_model_results['prediction']),
                'sub_confidence': float(second_model_results['probabilities'][second_model_results['prediction']] * 100),
                'top_5_sub_predictions': [
                    {
                        'virus': OTHER_VIRUS_MAPPING[idx],
                        'virus_id': int(idx),
                        'confidence': float(second_model_results['probabilities'][idx] * 100)
                    } for idx in second_model_results['top_5']
                ]
            }
            prediction_result['sub_classification'] = sub_classification
        
        # Save to database (non-blocking)
        saved_id = None
        try:
            state_name, district_name = _resolve_location_names(patient_dict)
            saved_id = save_prediction_to_db(
                patient_data=patient_dict,
                prediction_result=prediction_result,
                models_info={'model1': 'CustomMajor', 'model2': 'CustomOther'},
                state_name=state_name,
                district_name=district_name
            )
        except Exception as db_error:
            logger.warning(f"Database save failed: {db_error}")
        
        # Return response
        return PredictionResponse(
            success=True,
            predicted_virus=prediction_result['predicted_virus'],
            predicted_virus_id=prediction_result['predicted_virus_id'],
            confidence=prediction_result['confidence'],
            top_5_predictions=prediction_result['top_5_predictions'],
            sub_classification=sub_classification,
            models_info={'model1': 'CustomMajor', 'model2': 'CustomOther'},
            timestamp=datetime.now().isoformat(),
            prediction_id=saved_id
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Prediction error: {e}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Prediction failed: {str(e)}"
        )


@app.post("/validate")
async def validate_prediction(validation: ValidationRequest):
    """
    Submit validation feedback for a prediction
    
    - Links actual diagnosis to predicted results
    - Helps track model accuracy
    - Stored in MongoDB for analysis
    """
    try:
        # Map virus ID to virus name
        actual_virus_name = ""
        actual_virus_key = ""
        
        if validation.actual_virus_category.lower() in ['main', 'major']:
            if validation.actual_virus_id in VIRUS_MAPPING:
                actual_virus_name = VIRUS_MAPPING[validation.actual_virus_id]
                actual_virus_key = f"main_{validation.actual_virus_id}"
        elif validation.actual_virus_category.lower() == 'other':
            if validation.actual_virus_id in OTHER_VIRUS_MAPPING:
                actual_virus_name = OTHER_VIRUS_MAPPING[validation.actual_virus_id]
                actual_virus_key = f"other_{validation.actual_virus_id}"
        
        # Build validation data dictionary
        validation_data = {
            'prediction_id': validation.prediction_id,
            'actual_virus_name': actual_virus_name,
            'actual_virus_key': actual_virus_key,
            'notes': validation.feedback_notes or '',
            'is_correct': validation.is_correct
        }
        
        success = save_validation_to_db(validation_data)
        
        if success:
            return {
                "success": True,
                "message": "Validation feedback saved successfully"
            }
        else:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail="Failed to save validation"
            )
            
    except Exception as e:
        logger.error(f"Validation save error: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Validation failed: {str(e)}"
        )


@app.get("/stats")
async def get_statistics():
    """
    Get prediction statistics
    
    - Total predictions made
    - Database health
    - Model usage stats
    """
    try:
        stats = get_prediction_stats()
        return {
            "success": True,
            "statistics": stats,
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        logger.error(f"Stats retrieval error: {e}")
        return {
            "success": False,
            "error": str(e),
            "statistics": {}
        }


@app.get("/info")
async def get_info():
    """Get API information and available endpoints"""
    return {
        "api_name": "Virus Prediction API",
        "version": "1.0.0",
        "description": "AI-powered viral infection prediction system",
        "endpoints": {
            "/": "Interactive API documentation (Swagger UI)",
            "/health": "Health check endpoint",
            "/predict": "Make virus prediction (POST)",
            "/validate": "Submit validation feedback (POST)",
            "/mappings": "Get virus and symptom mappings",
            "/location-mappings": "Get state and district encoder mappings",
            "/locations": "Alias for /location-mappings",
            "/stats": "Get prediction statistics",
            "/info": "API information (this endpoint)"
        },
        "models": {
            "model1": "CustomMajor - 26 virus categories",
            "model2": "CustomOther - 13 sub-categories"
        },
        "deployment": "Hugging Face Spaces (Free Tier)",
        "database": "MongoDB Atlas"
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)