#!/usr/bin/env python3
"""
FastAPI Backend for Cat vs Dog Classification
Provides endpoints for image upload and prediction
"""

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import numpy as np
import joblib
import json
import os
import cv2
from PIL import Image
import io
from huggingface_hub import hf_hub_download

app = FastAPI(title="Cat vs Dog Classification API", version="1.0.0")

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

# Mount static files
os.makedirs("static", exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

# Global variables for models and artifacts
model = None
scaler = None
label_encoder = None
metadata = None

def load_models():
    """Load trained models and artifacts from Hugging Face Hub"""
    global model, scaler, label_encoder, metadata
    
    try:
        # Model repository configuration
        repo_id = os.getenv("HF_MODEL_REPO", "enigmaceo/svm-classification-cat-and-dog")
        
        # Download and load best model (compressed)
        model_path = hf_hub_download(repo_id, "svm_best_model.pkl.gz")
        import gzip
        import pickle
        with gzip.open(model_path, 'rb') as f:
            model = pickle.load(f)
        
        # Download and load scaler
        scaler_path = hf_hub_download(repo_id, "scaler.pkl")
        scaler = joblib.load(scaler_path)
        
        # Download and load label encoder
        encoder_path = hf_hub_download(repo_id, "label_encoder.pkl")
        label_encoder = joblib.load(encoder_path)
        
        # Download and load metadata
        metadata_path = hf_hub_download(repo_id, "metadata.json")
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        print(f"Model and artifacts loaded successfully from {repo_id}")
        return True
    except Exception as e:
        print(f"Error loading models from Hugging Face: {e}")
        return False

def extract_hog_features(image, pixels_per_cell=(8, 8)):
    """Extract HOG features from image"""
    from skimage.feature import hog
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    features, hog_img = hog(
        gray,
        orientations=9,
        pixels_per_cell=pixels_per_cell,
        cells_per_block=(2, 2),
        block_norm='L2-Hys',
        visualize=True,
        transform_sqrt=True
    )
    return features.astype(np.float32)

def extract_color_histogram(image, bins=32):
    """Extract color histogram features from HSV image"""
    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    hist_h = np.histogram(hsv[:,:,0], bins=bins, range=(0, 180))[0]
    hist_s = np.histogram(hsv[:,:,1], bins=bins, range=(0, 256))[0]
    hist_v = np.histogram(hsv[:,:,2], bins=bins, range=(0, 256))[0]
    return np.concatenate([hist_h, hist_s, hist_v]).astype(np.float32)

def extract_lbp_features(image, radius=3, n_points=24):
    """Extract Local Binary Pattern features for texture"""
    from skimage.feature import local_binary_pattern
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    lbp = local_binary_pattern(gray, n_points, radius, method='uniform')
    hist, _ = np.histogram(lbp.ravel(), bins=n_points + 2)
    hist = hist.astype(np.float32)
    hist /= (hist.sum() + 1e-7)  # Normalize
    return hist

def extract_features_from_image(image_data: bytes) -> np.ndarray:
    """
    Extract HOG, color histogram, and LBP features from uploaded image
    Same feature extraction as used in training
    """
    try:
        # Convert bytes to PIL Image
        image = Image.open(io.BytesIO(image_data))
        
        # Convert to numpy array and RGB
        img_array = np.array(image)
        if len(img_array.shape) == 2:  # Grayscale
            img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
        elif img_array.shape[2] == 4:  # RGBA
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
        
        # Resize to match training size
        img_resized = cv2.resize(img_array, (128, 128))
        
        # Extract HOG features
        hog_feat = extract_hog_features(img_resized)
        
        # Extract color histogram
        col_feat = extract_color_histogram(img_resized)
        
        # Extract LBP features
        lbp_feat = extract_lbp_features(img_resized)
        
        # Combine features
        combined_features = np.concatenate([hog_feat, col_feat, lbp_feat])
        
        return combined_features.reshape(1, -1)
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")

@app.on_event("startup")
async def startup_event():
    """Load models on startup"""
    success = load_models()
    if not success:
        print("Warning: Could not load models. Please run training script first.")

@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve the dashboard"""
    try:
        with open("dashboard/index.html", "r") as f:
            return HTMLResponse(content=f.read())
    except FileNotFoundError:
        return HTMLResponse(content="<h1>Cat vs Dog Classification</h1><p>Dashboard not found. Please check dashboard folder.</p>")

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": model is not None,
        "best_kernel": metadata.get('best_kernel') if metadata else None
    }

@app.post("/api/predict")
async def predict_image(file: UploadFile = File(...)):
    """Predict cat or dog from uploaded image"""
    if not model:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    if not scaler:
        raise HTTPException(status_code=503, detail="Scaler not loaded")
    
    if not label_encoder:
        raise HTTPException(status_code=503, detail="Label encoder not loaded")
    
    try:
        # Read image data
        image_data = await file.read()
        
        # Extract features
        features = extract_features_from_image(image_data)
        features_scaled = scaler.transform(features)
        
        # Extract individual features for display
        img_array = np.array(Image.open(io.BytesIO(image_data)))
        if len(img_array.shape) == 2:  # Grayscale
            img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
        elif img_array.shape[2] == 4:  # RGBA
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
        img_resized = cv2.resize(img_array, (128, 128))
        
        hog_feat = extract_hog_features(img_resized)
        col_feat = extract_color_histogram(img_resized)
        lbp_feat = extract_lbp_features(img_resized)
        
        # Make prediction
        prediction = model.predict(features_scaled)[0]
        class_id = int(prediction)
        class_name = label_encoder.inverse_transform([class_id])[0]
        
        # Get confidence if available
        confidence = None
        if hasattr(model, 'decision_function'):
            decision_values = model.decision_function(features_scaled)[0]
            exp_values = np.exp(decision_values - np.max(decision_values))
            probabilities = exp_values / np.sum(exp_values)
            confidence = float(np.max(probabilities))
        
        return {
            "prediction": {
                "class_id": class_id,
                "class_name": class_name,
                "confidence": confidence
            },
            "features": {
                "hog_size": len(hog_feat),
                "color_size": len(col_feat),
                "lbp_size": len(lbp_feat)
            }
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")

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