import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import efficientnet_b0, mobilenet_v2, resnet18
import gradio as gr
import numpy as np
from PIL import Image
import joblib
import os

# Constants
NUM_CLASSES = 4
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CLASS_NAMES = ['Cracked_road', 'Flooded_muddy', 'Good_condition', 'Potholes']

# Model paths with fallbacks for different deployment environments
MODEL_PATHS = {
    "efficientnet": [
        "best_efficientnet_model.pth",
        "models/best_efficientnet_model.pth",
        "assign2_Latest/models/best_efficientnet_model.pth"
    ],
    "mobilenet": [
        "best_mobilenetv2_model.pth",
        "models/best_mobilenetv2_model.pth",
        "assign2_Latest/models/best_mobilenetv2_model.pth"
    ],
    "resnet": [
        "best_resnet18_model.pth",
        "models/best_resnet18_model.pth",
        "assign2_Latest/models/best_resnet18_model.pth"
    ],
    "rf": [
        "rf_model.pkl",
        "models/rf_model.pkl",
        "assign2_Latest/models/rf_model.pkl"
    ]
}

# Build EfficientNet model
def build_efficientnet(num_classes=NUM_CLASSES):
    """Build EfficientNet-B0 based model"""
    model = efficientnet_b0(weights=None)
    
    # Replace classifier
    in_features = model.classifier[1].in_features
    model.classifier = nn.Sequential(
        nn.Linear(in_features, 256),
        nn.BatchNorm1d(256),
        nn.ReLU(),
        nn.Dropout(p=0.2),
        nn.Linear(256, num_classes)
    )
    
    return model

# Build MobileNetV2 model
def build_mobilenet(num_classes=NUM_CLASSES):
    """Build MobileNetV2 model"""
    model = mobilenet_v2(weights=None)
    
    # Replace classifier
    last_channel = model.last_channel
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.2),
        nn.Linear(last_channel, 512),
        nn.BatchNorm1d(512),
        nn.LeakyReLU(0.2),
        nn.Dropout(p=0.2),
        nn.Linear(512, 256),
        nn.BatchNorm1d(256),
        nn.LeakyReLU(0.2),
        nn.Dropout(p=0.1),
        nn.Linear(256, num_classes)
    )
    
    return model

# Build ResNet18 model
def build_resnet(num_classes=NUM_CLASSES):
    """Build ResNet18 model"""
    model = resnet18(weights=None)
    
    # Replace classifier
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(in_features, 256),
        nn.ReLU(),
        nn.Dropout(0.3),
        nn.Linear(256, num_classes)
    )
    
    return model

# Load model with fallback paths
def load_model(model_type):
    """Load trained model with fallback paths"""
    if model_type == "efficientnet":
        model = build_efficientnet()
    elif model_type == "mobilenet":
        model = build_mobilenet()
    elif model_type == "resnet":
        model = build_resnet()
    elif model_type == "rf":
        # For Random Forest, handled differently below
        return load_rf_model()
    else:
        raise ValueError(f"Unknown model type: {model_type}")
    
    # Try loading from various possible paths
    paths = MODEL_PATHS[model_type]
    model_loaded = False
    
    for path in paths:
        if os.path.exists(path):
            try:
                model.load_state_dict(torch.load(path, map_location=DEVICE))
                model_loaded = True
                print(f"{model_type} model loaded from {path}")
                break
            except Exception as e:
                print(f"Failed to load {model_type} from {path}: {e}")
                continue
    
    if not model_loaded:
        raise FileNotFoundError(f"Could not find or load {model_type} model file")
    
    model.to(DEVICE)
    model.eval()
    return model

# Load Random Forest model
def load_rf_model():
    """Load Random Forest model"""
    paths = MODEL_PATHS["rf"]
    
    for path in paths:
        if os.path.exists(path):
            try:
                model = joblib.load(path)
                print(f"Random Forest model loaded from {path}")
                return model
            except Exception as e:
                print(f"Failed to load RF model from {path}: {e}")
                continue
    
    raise FileNotFoundError("Could not find or load Random Forest model file")

# Extract features for Random Forest
def extract_features_single(model, image_tensor):
    """Extract features from a single image for Random Forest"""
    model.eval()
    
    # Create feature extractor (using EfficientNet)
    feature_extractor = nn.Sequential(
        model.features,
        model.avgpool,
        nn.Flatten()
    ).to(DEVICE)
    
    with torch.no_grad():
        features = feature_extractor(image_tensor).cpu().numpy()
    
    return features

# Preprocess image for neural networks
def preprocess_image(image):
    """Preprocess image for model prediction"""
    preprocess = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    return preprocess(image).unsqueeze(0)

# Prediction function
def predict(img, model_choice, models):
    """Make prediction using selected model"""
    if img is None:
        return "No image provided", "Please upload an image"
    
    try:
        # Convert to RGB if needed
        if img.mode != "RGB":
            img = img.convert("RGB")
            
        # Preprocess image
        img_tensor = preprocess_image(img)
        img_tensor = img_tensor.to(DEVICE)
        
        if model_choice == "EfficientNet":
            model = models["efficientnet"]
            # Make prediction
            with torch.no_grad():
                outputs = model(img_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
                predicted_class = torch.argmax(probabilities).item()
            
            # Format confidence scores
            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i].item()*100:.2f}%" 
                               for i in range(len(CLASS_NAMES))}
            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
            
            return CLASS_NAMES[predicted_class], confidence_text
            
        elif model_choice == "MobileNetV2":
            model = models["mobilenet"]
            # Make prediction
            with torch.no_grad():
                outputs = model(img_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
                predicted_class = torch.argmax(probabilities).item()
            
            # Format confidence scores
            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i].item()*100:.2f}%" 
                               for i in range(len(CLASS_NAMES))}
            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
            
            return CLASS_NAMES[predicted_class], confidence_text
            
        elif model_choice == "ResNet18":
            model = models["resnet"]
            # Make prediction
            with torch.no_grad():
                outputs = model(img_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
                predicted_class = torch.argmax(probabilities).item()
            
            # Format confidence scores
            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i].item()*100:.2f}%" 
                               for i in range(len(CLASS_NAMES))}
            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
            
            return CLASS_NAMES[predicted_class], confidence_text
            
        elif model_choice == "Random Forest":
            # Extract features using EfficientNet
            features = extract_features_single(models["efficientnet"], img_tensor)
            
            # Make prediction with Random Forest
            rf_model = models["rf"]
            predicted_class = rf_model.predict(features)[0]
            probabilities = rf_model.predict_proba(features)[0]
            
            # Format confidence scores
            confidence_scores = {CLASS_NAMES[i]: f"{probabilities[i]*100:.2f}%" 
                               for i in range(len(CLASS_NAMES))}
            confidence_text = "\n".join([f"{cls}: {score}" for cls, score in confidence_scores.items()])
            
            return CLASS_NAMES[predicted_class], confidence_text
        else:
            return "Invalid model selection", "Please select a valid model"
            
    except Exception as e:
        return f"Error: {str(e)}", "An error occurred during prediction."

# Initialize models
print("Loading models...")
models = {}
try:
    models["efficientnet"] = load_model("efficientnet")
    models["mobilenet"] = load_model("mobilenet")
    models["resnet"] = load_model("resnet")
    models["rf"] = load_model("rf")
    print("All models loaded successfully!")
except Exception as e:
    print(f"Error loading models: {e}")
    # In case of error, models will be empty or partial

# Create Gradio interface
with gr.Blocks(title="Road Surface Classification") as demo:
    gr.Markdown("# Road Surface Classification")
    gr.Markdown("Upload an image of a road surface to classify its condition using various models.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Road Image")
            model_selector = gr.Radio(
                choices=["EfficientNet", "MobileNetV2", "ResNet18", "Random Forest"], 
                value="EfficientNet", 
                label="Select Model"
            )
            classify_btn = gr.Button("Classify", variant="primary")
            
        with gr.Column():
            label_output = gr.Textbox(label="Predicted Class", interactive=False)
            confidence_output = gr.Textbox(
                label="Confidence Scores", 
                lines=5, 
                interactive=False
            )

    # Event handler
    classify_btn.click(
        fn=lambda img, model_choice: predict(img, model_choice, models),
        inputs=[image_input, model_selector],
        outputs=[label_output, confidence_output]
    )
    
    # Add information about classes and models
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Classes")
            gr.Markdown("- **Cracked_road**: Roads with visible cracks")
            gr.Markdown("- **Flooded_muddy**: Roads affected by flooding or mud")
            gr.Markdown("- **Good_condition**: Roads in good condition")
            gr.Markdown("- **Potholes**: Roads with potholes")
        
        with gr.Column():
            gr.Markdown("### Models")
            gr.Markdown("- **EfficientNet**: Efficient deep learning model")
            gr.Markdown("- **MobileNetV2**: Lightweight mobile-friendly model")
            gr.Markdown("- **ResNet18**: Residual network architecture")
            gr.Markdown("- **Random Forest**: Traditional ML using CNN features")

# Launch the app
if __name__ == "__main__":
    demo.launch()