"""
FastAPI Application for Driving Behavior Analysis
==================================================

Comprehensive REST API with Swagger documentation for testing the driving behavior
classification model. Includes batch predictions, real-time classification, and 
detailed confidence scores.

Run: uvicorn main:app --reload --host 0.0.0.0 --port 8000
Swagger UI: http://localhost:8000/docs
"""

from fastapi import FastAPI, HTTPException, Query
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder
import pickle
import json
from datetime import datetime
import logging
import uvicorn
from collections import deque
import threading
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ============================================================================
# LOAD MODEL & PREPROCESSING OBJECTS
# ============================================================================

try:
    import joblib
    import os
    
    # Use actual relative paths so they work on Hugging Face servers
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(BASE_DIR, 'model.pkl')
    scaler_path = os.path.join(BASE_DIR, 'scaler.pkl')
    le_path = os.path.join(BASE_DIR, 'label_encoder.pkl')
    fc_path = os.path.join(BASE_DIR, 'feature_columns.pkl')
    
    best_model = joblib.load(model_path)
    scaler = joblib.load(scaler_path)
    label_encoder = joblib.load(le_path)
    
    # Feature names (in correct order)
    with open(fc_path, 'rb') as f:
        feature_columns = pickle.load(f)
    
    logger.info("✅ All models and preprocessors loaded successfully!")
    
except FileNotFoundError as e:
    logger.warning(f"⚠️ Could not load model files: {e}")
    logger.warning("⚠️ Using mock models for demonstration")
    best_model = None
    scaler = None
    label_encoder = None
    feature_columns = None


# Global buffer to store recent sensor readings for proper time-series feature engineering
reading_history = deque(maxlen=15)
history_lock = threading.Lock()

# ============================================================================
# PYDANTIC MODELS (Request/Response Schemas)
# ============================================================================

class SensorInput(BaseModel):
    """Raw sensor input from accelerometer and gyroscope"""
    acc_x: float = Field(
        ..., 
        description="Acceleration in X direction (m/s²)",
        example=0.5
    )
    acc_y: float = Field(
        ..., 
        description="Acceleration in Y direction (m/s²)",
        example=0.2
    )
    acc_z: float = Field(
        ..., 
        description="Acceleration in Z direction (m/s²)",
        example=9.8
    )
    gyro_x: float = Field(
        ..., 
        description="Angular velocity around X axis (rad/s)",
        example=0.01
    )
    gyro_y: float = Field(
        ..., 
        description="Angular velocity around Y axis (rad/s)",
        example=0.02
    )
    gyro_z: float = Field(
        ..., 
        description="Angular velocity around Z axis (rad/s)",
        example=0.03
    )

    class Config:
        json_schema_extra = {
            "example": {
                "acc_x": 0.5,
                "acc_y": 0.2,
                "acc_z": 9.8,
                "gyro_x": 0.01,
                "gyro_y": 0.02,
                "gyro_z": 0.03
            }
        }


class PredictionResponse(BaseModel):
    """Response with prediction and confidence scores"""
    prediction: str = Field(..., description="Predicted driving behavior class")
    confidence: Dict[str, float] = Field(..., description="Confidence scores for each class")
    timestamp: str = Field(..., description="Prediction timestamp")
    
    class Config:
        json_schema_extra = {
            "example": {
                "prediction": "NORMAL",
                "confidence": {
                    "AGGRESSIVE": 0.02,
                    "NORMAL": 0.88,
                    "SLOW": 0.10
                },
                "timestamp": "2024-04-17T12:34:56"
            }
        }


class BatchPredictionRequest(BaseModel):
    """Request for batch predictions"""
    samples: List[SensorInput] = Field(..., description="List of sensor readings")
    return_features: bool = Field(
        False, 
        description="Include engineered features in response"
    )


class BatchPredictionResponse(BaseModel):
    """Response with batch predictions"""
    total_samples: int
    successful_predictions: int
    failed_predictions: int
    predictions: List[Dict] = Field(..., description="List of predictions")
    processing_time_ms: float


class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    model_loaded: bool
    model_version: str
    timestamp: str
    features_count: Optional[int] = None


# ============================================================================
# FEATURE ENGINEERING FUNCTION
# ============================================================================

def engineer_features(data_list: list) -> pd.DataFrame:
    """
    Apply feature engineering to raw sensor data using a sequence of historical readings
    to correctly compute rates of changes (Jerk) and rolling statistics.
    """
    try:
        # Create DataFrame from list
        df = pd.DataFrame(data_list)
        
        # Rename columns to match training
        df = df.rename(columns={
            'acc_x': 'AccX', 'acc_y': 'AccY', 'acc_z': 'AccZ',
            'gyro_x': 'GyroX', 'gyro_y': 'GyroY', 'gyro_z': 'GyroZ'
        })
        
        # ===== JERK CALCULATION =====
        # For single row: jerk = 0
        df['JerkX'] = df['AccX'].diff().fillna(0)
        df['JerkY'] = df['AccY'].diff().fillna(0)
        df['JerkZ'] = df['AccZ'].diff().fillna(0)
        
        # ===== MAGNITUDE FEATURES =====
        df['AccMagnitude'] = np.sqrt(df['AccX']**2 + df['AccY']**2 + df['AccZ']**2)
        df['GyroMagnitude'] = np.sqrt(df['GyroX']**2 + df['GyroY']**2 + df['GyroZ']**2)
        df['JerkMagnitude'] = np.sqrt(df['JerkX']**2 + df['JerkY']**2 + df['JerkZ']**2)
        
        # ===== ROLLING STATISTICS =====
        window_size = 5
        df['AccX_rolling_mean'] = df['AccX'].rolling(window=window_size, min_periods=1).mean()
        df['AccY_rolling_mean'] = df['AccY'].rolling(window=window_size, min_periods=1).mean()
        df['AccZ_rolling_mean'] = df['AccZ'].rolling(window=window_size, min_periods=1).mean()
        
        df['AccX_rolling_std'] = df['AccX'].rolling(window=window_size, min_periods=1).std().fillna(0)
        df['AccY_rolling_std'] = df['AccY'].rolling(window=window_size, min_periods=1).std().fillna(0)
        df['AccZ_rolling_std'] = df['AccZ'].rolling(window=window_size, min_periods=1).std().fillna(0)
        
        df['JerkX_rolling_mean'] = df['JerkX'].rolling(window=window_size, min_periods=1).mean()
        df['JerkY_rolling_mean'] = df['JerkY'].rolling(window=window_size, min_periods=1).mean()
        df['JerkZ_rolling_mean'] = df['JerkZ'].rolling(window=window_size, min_periods=1).mean()
        
        df['JerkX_rolling_max'] = df['JerkX'].rolling(window=window_size, min_periods=1).max()
        df['JerkY_rolling_max'] = df['JerkY'].rolling(window=window_size, min_periods=1).max()
        df['JerkZ_rolling_max'] = df['JerkZ'].rolling(window=window_size, min_periods=1).max()
        
        # ===== VARIANCE & ENERGY =====
        df['AccX_var'] = df['AccX'] ** 2
        df['AccY_var'] = df['AccY'] ** 2
        df['AccZ_var'] = df['AccZ'] ** 2
        
        df['JerkX_var'] = df['JerkX'] ** 2
        df['JerkY_var'] = df['JerkY'] ** 2
        df['JerkZ_var'] = df['JerkZ'] ** 2
        
        # ===== ABSOLUTE VALUES =====
        df['AbsAccX'] = abs(df['AccX'])
        df['AbsAccY'] = abs(df['AccY'])
        df['AbsAccZ'] = abs(df['AccZ'])
        df['AbsJerkX'] = abs(df['JerkX'])
        df['AbsJerkY'] = abs(df['JerkY'])
        df['AbsJerkZ'] = abs(df['JerkZ'])
        
        # ===== ENERGY FEATURES =====
        df['Acc_Energy'] = (df['AccX']**2 + df['AccY']**2 + df['AccZ']**2) / 3
        df['Jerk_Energy'] = (df['JerkX']**2 + df['JerkY']**2 + df['JerkZ']**2) / 3
        
        return df
    
    except Exception as e:
        logger.error(f"Error in feature engineering: {str(e)}")
        raise


# ============================================================================
# PREDICTION FUNCTION
# ============================================================================

def predict_driving_behavior(data_list: list) -> Dict:
    """
    Predict driving behavior from sensor data sequence
    
    Args:
        data_list: List of dictionaries with sensor readings
        
    Returns:
        Dictionary with prediction and confidence scores
    """
    try:
        # Check if models are loaded
        if best_model is None or scaler is None or label_encoder is None:
            raise ValueError("Models not loaded. Cannot make predictions.")
        
        # Engineer features
        df = engineer_features(data_list)
        
        # Select features in correct order
        df_processed = df[feature_columns]
        
        # Scale features
        df_scaled = scaler.transform(df_processed)
        
        # Extract only the latest row for prediction
        latest_row_scaled = df_scaled[-1].reshape(1, -1)
        
        # Make prediction
        pred = best_model.predict(latest_row_scaled)
        proba = best_model.predict_proba(latest_row_scaled)
        
        # Decode prediction
        prediction = label_encoder.inverse_transform(pred)[0]
        
        # Get confidence scores
        confidence_dict = {
            label: float(proba[0][i]) 
            for i, label in enumerate(label_encoder.classes_)
        }
        
        return {
            "prediction": prediction,
            "confidence": confidence_dict,
            "raw_probability": proba[0].tolist()
        }
    
    except Exception as e:
        logger.error(f"Prediction error: {str(e)}")
        raise


# ============================================================================
# FASTAPI APPLICATION
# ============================================================================

app = FastAPI(
    title="🚗 Driving Behavior Analysis API",
    description="""
    Real-time driving behavior classification API using machine learning.
    
    Classify driving patterns into three categories:
    - **NORMAL**: Regular, safe driving
    - **SLOW**: Cautious, slower driving
    - **AGGRESSIVE**: Risky, aggressive driving
    
    ## Features
    - Single prediction endpoint
    - Batch prediction support
    - Real-time confidence scores
    - Feature engineering included
    - Health check endpoint
    - Comprehensive API documentation
    
    ## How to Use
    1. Provide raw sensor data (acceleration & gyroscope readings)
    2. API automatically engineers features
    3. Get driving behavior classification with confidence scores
    
    ## Example Request
    ```json
    {
        "acc_x": 0.5,
        "acc_y": 0.2,
        "acc_z": 9.8,
        "gyro_x": 0.01,
        "gyro_y": 0.02,
        "gyro_z": 0.03
    }
    ```
    """,
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json",
    contact={
        "name": "ML Team",
        "email": "ml@example.com"
    }
)


# ============================================================================
# HEALTH CHECK ENDPOINT
# ============================================================================

@app.get(
    "/health",
    response_model=HealthResponse,
    tags=["Health"],
    summary="Health Check",
    description="Check API status and model availability"
)
async def health_check():
    """
    Check the health status of the API and model availability
    """
    return HealthResponse(
        status="healthy",
        model_loaded=best_model is not None,
        model_version="1.0.0",
        timestamp=datetime.now().isoformat(),
        features_count=len(feature_columns) if feature_columns else 0
    )


# ============================================================================
# SINGLE PREDICTION ENDPOINT
# ============================================================================

@app.post(
    "/predict",
    response_model=PredictionResponse,
    tags=["Prediction"],
    summary="Predict Driving Behavior",
    description="Predict driving behavior from a single sensor reading"
)
async def predict(sensor_input: SensorInput):
    """
    Predict driving behavior from raw sensor data.
    
    **Input Parameters:**
    - acc_x: Acceleration in X direction (m/s²)
    - acc_y: Acceleration in Y direction (m/s²)
    - acc_z: Acceleration in Z direction (m/s²)
    - gyro_x: Angular velocity around X axis (rad/s)
    - gyro_y: Angular velocity around Y axis (rad/s)
    - gyro_z: Angular velocity around Z axis (rad/s)
    
    **Response:**
    - prediction: One of [AGGRESSIVE, NORMAL, SLOW]
    - confidence: Confidence scores for each class
    - timestamp: When prediction was made
    
    **Example Request:**
    ```json
    {
        "acc_x": 0.5,
        "acc_y": 0.2,
        "acc_z": 9.8,
        "gyro_x": 0.01,
        "gyro_y": 0.02,
        "gyro_z": 0.03
    }
    ```
    
    **Example Response:**
    ```json
    {
        "prediction": "NORMAL",
        "confidence": {
            "AGGRESSIVE": 0.02,
            "NORMAL": 0.88,
            "SLOW": 0.10
        },
        "timestamp": "2024-04-17T12:34:56"
    }
    ```
    """
    try:
        # Convert input to dictionary
        input_dict = sensor_input.dict()
        
        # Append to global history
        with history_lock:
            reading_history.append(input_dict)
            history_snapshot = list(reading_history)
        
        # Make prediction using history
        result = predict_driving_behavior(history_snapshot)
        
        return PredictionResponse(
            prediction=result["prediction"],
            confidence=result["confidence"],
            timestamp=datetime.now().isoformat()
        )
    
    except Exception as e:
        logger.error(f"Prediction error: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Prediction failed: {str(e)}"
        )


# ============================================================================
# BATCH PREDICTION ENDPOINT
# ============================================================================

@app.post(
    "/predict-batch",
    response_model=BatchPredictionResponse,
    tags=["Batch Prediction"],
    summary="Batch Predictions",
    description="Make predictions on multiple sensor readings at once"
)
async def predict_batch(request: BatchPredictionRequest):
    """
    Predict driving behavior for multiple sensor readings.
    
    Useful for processing streams or datasets efficiently.
    
    **Request Parameters:**
    - samples: List of sensor readings
    - return_features: Whether to include engineered features in response
    
    **Returns:**
    - total_samples: Number of samples processed
    - successful_predictions: Number of successful predictions
    - failed_predictions: Number of failed predictions
    - predictions: List of prediction results
    - processing_time_ms: Total processing time
    """
    import time
    
    start_time = time.time()
    predictions = []
    successful = 0
    failed = 0
    
    try:
        for i, sensor_input in enumerate(request.samples):
            try:
                input_dict = sensor_input.dict()
                
                # Using the same global history mechanism to accumulate batch over time
                with history_lock:
                    reading_history.append(input_dict)
                    history_snapshot = list(reading_history)
                
                result = predict_driving_behavior(history_snapshot)
                
                prediction_result = {
                    "sample_index": i,
                    "prediction": result["prediction"],
                    "confidence": result["confidence"],
                    "timestamp": datetime.now().isoformat()
                }
                
                predictions.append(prediction_result)
                successful += 1
            
            except Exception as e:
                logger.error(f"Error on sample {i}: {str(e)}")
                predictions.append({
                    "sample_index": i,
                    "error": str(e)
                })
                failed += 1
        
        processing_time = (time.time() - start_time) * 1000  # Convert to ms
        
        return BatchPredictionResponse(
            total_samples=len(request.samples),
            successful_predictions=successful,
            failed_predictions=failed,
            predictions=predictions,
            processing_time_ms=processing_time
        )
    
    except Exception as e:
        logger.error(f"Batch prediction error: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Batch prediction failed: {str(e)}"
        )


# ============================================================================
# DETAILED PREDICTION ENDPOINT
# ============================================================================

@app.post(
    "/predict-detailed",
    tags=["Prediction"],
    summary="Detailed Prediction",
    description="Get detailed prediction with engineered features"
)
async def predict_detailed(sensor_input: SensorInput):
    """
    Get detailed prediction including engineered features.
    
    Useful for understanding which features influenced the prediction.
    """
    try:
        input_dict = sensor_input.dict()
        
        with history_lock:
            reading_history.append(input_dict)
            history_snapshot = list(reading_history)
        
        # Engineer features
        df = engineer_features(history_snapshot)
        
        # Make prediction
        result = predict_driving_behavior(history_snapshot)
        
        return {
            "prediction": result["prediction"],
            "confidence": result["confidence"],
            "engineered_features": df.to_dict(orient='records')[-1],
            "timestamp": datetime.now().isoformat()
        }
    
    except Exception as e:
        logger.error(f"Detailed prediction error: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Detailed prediction failed: {str(e)}"
        )


# ============================================================================
# TEST DATA ENDPOINT
# ============================================================================

@app.get(
    "/test-samples",
    tags=["Testing"],
    summary="Get Test Samples",
    description="Get example sensor readings for testing"
)
async def get_test_samples():
    """
    Get example sensor readings for different driving behaviors.
    
    Useful for testing the API without real sensor data.
    """
    return {
        "NORMAL": {
            "description": "Normal driving - balanced sensor readings",
            "sample": {
                "acc_x": 0.05,
                "acc_y": -0.08,
                "acc_z": 9.81,
                "gyro_x": 0.002,
                "gyro_y": -0.001,
                "gyro_z": 0.008
            }
        },
        "SLOW": {
            "description": "Slow driving - smooth, low acceleration",
            "sample": {
                "acc_x": 0.1,
                "acc_y": -0.02,
                "acc_z": 9.8,
                "gyro_x": 0.0,
                "gyro_y": 0.0,
                "gyro_z": 0.001
            }
        },
        "AGGRESSIVE": {
            "description": "Aggressive driving - high accelerations and jerky movements",
            "sample": {
                "acc_x": 0.8,
                "acc_y": -0.5,
                "acc_z": 9.7,
                "gyro_x": 0.05,
                "gyro_y": 0.03,
                "gyro_z": 0.1
            }
        }
    }


# ============================================================================
# INFO ENDPOINT
# ============================================================================

@app.get(
    "/info",
    tags=["Info"],
    summary="API Information",
    description="Get information about the API and model"
)
async def get_info():
    """
    Get detailed information about the API and trained model.
    """
    return {
        "api_name": "Driving Behavior Analysis API",
        "version": "1.0.0",
        "model_status": "loaded" if best_model is not None else "not_loaded",
        "supported_classes": [
            "AGGRESSIVE",
            "NORMAL",
            "SLOW"
        ],
        "features_count": len(feature_columns) if feature_columns else "unknown",
        "endpoints": {
            "predict": "POST /predict - Single prediction",
            "predict_batch": "POST /predict-batch - Batch predictions",
            "predict_detailed": "POST /predict-detailed - Detailed prediction with features",
            "health_check": "GET /health - Health check",
            "test_samples": "GET /test-samples - Get example test data",
            "info": "GET /info - API information"
        },
        "documentation": {
            "swagger": "http://localhost:8000/docs",
            "redoc": "http://localhost:8000/redoc",
            "openapi": "http://localhost:8000/openapi.json"
        }
    }


# ============================================================================
# ROOT ENDPOINT
# ============================================================================

@app.get(
    "/",
    tags=["Root"],
    summary="Welcome",
    description="Welcome message and quick start guide"
)
async def root():
    """
    Welcome to the Driving Behavior Analysis API!
    
    **Quick Start:**
    1. Go to http://localhost:8000/docs for interactive Swagger UI
    2. Try the /predict endpoint with sample data
    3. Check /test-samples for example inputs
    
    **Endpoints:**
    - POST /predict - Single prediction
    - POST /predict-batch - Batch predictions
    - GET /health - Health check
    - GET /test-samples - Test data
    - GET /info - API information
    """
    return {
        "message": "🚗 Welcome to Driving Behavior Analysis API",
        "status": "running",
        "docs": "http://localhost:8000/docs",
        "quick_start": [
            "1. Visit http://localhost:8000/docs",
            "2. Click on POST /predict",
            "3. Click 'Try it out'",
            "4. Enter sensor data or use example",
            "5. Click 'Execute' to get prediction"
        ]
    }


# ============================================================================
# EXCEPTION HANDLERS
# ============================================================================

@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """Handle HTTP exceptions"""
    return {
        "error": exc.detail,
        "status_code": exc.status_code,
        "timestamp": datetime.now().isoformat()
    }


# ============================================================================
# RUN THE APPLICATION
# ============================================================================


# ==========================================================================
# SAFETY SCORE ENDPOINT (NEW)
# ==========================================================================

from enum import Enum

class SafetyColor(str, Enum):
    GREEN = "Green"
    YELLOW = "Yellow"
    RED = "Red"

class SafetyScoreResponse(BaseModel):
    score: int = Field(..., description="Safety score (0-100)")
    color: SafetyColor = Field(..., description="Color-coded safety status")
    events: Dict[str, bool] = Field(..., description="Detected driving events")
    timestamp: str = Field(..., description="Timestamp of evaluation")

@app.post(
    "/safety-score",
    response_model=SafetyScoreResponse,
    tags=["Safety"],
    summary="Real-Time Safety Score",
    description="Get a real-time safety score and color-coded status from a single sensor reading."
)
async def safety_score(sensor_input: SensorInput):
    """
    Calculate a real-time safety score and color-coded status from sensor data.
    Detects harsh braking, rapid acceleration, aggressive cornering.
    """
    input_dict = sensor_input.dict()
    acc_x = input_dict["acc_x"]
    acc_y = input_dict["acc_y"]
    acc_z = input_dict["acc_z"]
    gyro_x = input_dict["gyro_x"]
    gyro_y = input_dict["gyro_y"]
    gyro_z = input_dict["gyro_z"]

    # Simple event detection thresholds (tune as needed)
    harsh_braking = acc_x < -2.5
    rapid_acceleration = acc_x > 2.5
    aggressive_cornering = abs(acc_y) > 2.0
    # (Advanced: add tailgating/lane weaving if you have more data)

    # Score logic (deduct for each event)
    score = 100
    if harsh_braking:
        score -= 30
    if rapid_acceleration:
        score -= 25
    if aggressive_cornering:
        score -= 20
    score = max(0, min(100, score))

    # Color coding
    if score >= 80:
        color = SafetyColor.GREEN
    elif score >= 50:
        color = SafetyColor.YELLOW
    else:
        color = SafetyColor.RED

    return SafetyScoreResponse(
        score=score,
        color=color,
        events={
            "harsh_braking": harsh_braking,
            "rapid_acceleration": rapid_acceleration,
            "aggressive_cornering": aggressive_cornering
        },
        timestamp=datetime.now().isoformat()
    )

if __name__ == "__main__":

    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )