import os
import time
import torch
import torch.nn as nn
from PIL import Image
import pandas as pd
import gradio as gr
from torchvision import models, transforms
from torch.autograd import Function

# 1. Dataset Configuration (Matching your ArchaeoMind settings)
CLASS_NAMES = ['artifact', 'non_artifact']
NUM_CLASSES = len(CLASS_NAMES)

# 2. Define Network Architectures to Load Weights Safely

# --- Model A: ResNet-18 ---
resnet18 = models.resnet18(weights=None)
num_features_res = resnet18.fc.in_features
resnet18.fc = nn.Sequential(
    nn.Linear(num_features_res, 256),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.4),
    nn.Linear(256, NUM_CLASSES)
)

# --- Model B: VGG16 Baseline ---
vgg16_baseline = models.vgg16(weights=None)
num_features_vgg = vgg16_baseline.classifier[0].in_features
vgg16_baseline.classifier = nn.Sequential(
    nn.Linear(num_features_vgg, 256),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.5),
    nn.Linear(256, NUM_CLASSES)
)

# --- Model C: VGG16 Domain-Adversarial (DANN) ---
class GradientReversalLayer(Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        return x.view_as(x)
    @staticmethod
    def backward(ctx, grad_output):
        return grad_output.neg() * ctx.alpha, None

def grad_reverse(x, alpha=1.0):
    return GradientReversalLayer.apply(x, alpha)

class VGG16_DANN(nn.Module):
    def __init__(self, num_classes=2):
        super(VGG16_DANN, self).__init__()
        vgg_base = models.vgg16(weights=None)
        self.feature_extractor = vgg_base.features
        self.pool = nn.AdaptiveAvgPool2d((4, 4))
        in_features = 512 * 4 * 4

        self.class_classifier = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(256, num_classes)
        )
        self.domain_classifier = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(256, 2)
        )

    def forward(self, x, alpha=1.0):
        features = self.feature_extractor(x)
        features = self.pool(features)
        features = torch.flatten(features, 1)
        class_output = self.class_classifier(features)
        reversed_features = grad_reverse(features, alpha)
        domain_output = self.domain_classifier(reversed_features)
        return class_output, domain_output

vgg16_dann = VGG16_DANN(num_classes=NUM_CLASSES)

# 3. Load Model Checkpoints onto CPU
models_dict = {
    "ResNet-18": resnet18,
    "VGG16 Baseline": vgg16_baseline,
    "VGG16 DANN (Domain Adaptation)": vgg16_dann
}

files_dict = {
    "ResNet-18": "resnet18_artifacts.pth",
    "VGG16 Baseline": "vgg16_artifacts_baseline.pth",
    "VGG16 DANN (Domain Adaptation)": "vgg16_dann_artifacts.pth"
}

for name, model in models_dict.items():
    file_path = files_dict[name]
    if os.path.exists(file_path):
        model.load_state_dict(torch.load(file_path, map_location=torch.device('cpu')))
        model.eval()
        print(f"Successfully loaded {name} weights structure.")
    else:
        print(f"WARNING: Checkpoint file '{file_path}' missing!")

# 4. Input Transforms Pipeline (Matches your 128x128 3-Channel Grayscale specs)
img_transforms = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 5. Inference Benchmarking Engine
def benchmark_inference(input_image):
    if input_image is None:
        return None
        
    img_tensor = img_transforms(input_image).unsqueeze(0)
    results_data = []
    
    with torch.no_grad():
        for name, model in models_dict.items():
            # Warm-up pass
            if name == "VGG16 DANN (Domain Adaptation)":
                _, _ = model(img_tensor, alpha=0.0)
            else:
                _ = model(img_tensor)
            
            # Performance timing run
            start_time = time.perf_counter()
            if name == "VGG16 DANN (Domain Adaptation)":
                outputs, _ = model(img_tensor, alpha=0.0)  # Ignore domain output head here
            else:
                outputs = model(img_tensor)
                
            probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
            best_prob, best_class_idx = torch.max(probabilities, dim=0)
            end_time = time.perf_counter()
            
            latency_ms = (end_time - start_time) * 1000
            predicted_label = CLASS_NAMES[best_class_idx.item()]
            confidence_pct = f"{best_prob.item() * 100:.2f}%"
            
            results_data.append({
                "Model Architecture": name,
                "Classification Prediction": predicted_label,
                "Confidence Score": confidence_pct,
                "Inference Latency": f"{latency_ms:.2f} ms"
            })
            
    return pd.DataFrame(results_data)

# 6. Gradio Interface Layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔍 ArchaeoMind: Artifact Detection & Speed Analysis Space")
    gr.Markdown("Upload an excavation find image to classify it as an artifact or non-artifact while profiling model efficiency metrics across architectures.")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_img_slot = gr.Image(type="pil", label="Excavation Find Photo")
            submit_btn = gr.Button("Analyze Image", variant="primary")
            
        with gr.Column(scale=2):
            output_table_slot = gr.Dataframe(
                headers=["Model Architecture", "Classification Prediction", "Confidence Score", "Inference Latency"],
                datatype=["str", "str", "str", "str"],
                label="Benchmark Results Table"
            )
            
    submit_btn.click(
        fn=benchmark_inference, 
        inputs=input_img_slot, 
        outputs=output_table_slot
    )

demo.launch()