import torch
import torch.nn as nn
from torchvision import transforms, models
import gradio as gr
from PIL import Image

# Model Architecture (sama seperti sebelumnya)
class ModelRecommender(nn.Module):
    def __init__(self, num_models, text_embedding_dim=768):
        super(ModelRecommender, self).__init__()
        
        # CNN for image processing
        self.cnn = models.resnet18(pretrained=True)
        self.cnn.fc = nn.Linear(512, 256)
        
        # MLP for text processing
        self.text_mlp = nn.Sequential(
            nn.Linear(text_embedding_dim, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU()
        )
        
        # Combined layers
        self.combined = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_models)
        )
        
    def forward(self, image, text_features):
        # Process image
        img_features = self.cnn(image)
        
        # Process text
        text_features = self.text_mlp(text_features)
        
        # Combine features
        combined = torch.cat((img_features, text_features), dim=1)
        
        # Final prediction
        output = self.combined(combined)
        return output

        

# Load model dan dataset info
def load_model():
    # Load dataset info
    dataset_info = torch.load('dataset_info.pth')
    model_names = dataset_info['model_names']
    
    # Initialize model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = ModelRecommender(len(model_names))
    
    # Load model weights
    checkpoint = torch.load('sd_recommender_model.pth', map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)
    model.eval()
    
    return model, model_names, device

def calculate_euclidean_distance(features1, features2):
    return np.linalg.norm(features1 - features2)

def predict_image(image):
    if not hasattr(predict_image, "model"):
        predict_image.model, predict_image.model_names, predict_image.device = load_model()
    
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                           std=[0.229, 0.224, 0.225])
    ])
    
    image_tensor = transform(image).unsqueeze(0).to(predict_image.device)
    dummy_text_features = torch.zeros(1, 768).to(predict_image.device)
    
    # Get image features
    with torch.no_grad():
        img_features = predict_image.model.cnn(image_tensor).cpu().numpy()
        outputs = predict_image.model(image_tensor, dummy_text_features)
        top5_prob, top5_indices = torch.topk(outputs, 5)
    
    # Create HTML gallery
    html_output = """
    <style>
    .model-gallery {
        display: grid;
        grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
        gap: 20px;
        padding: 20px;
    }
    .model-card {
        border: 1px solid #ddd;
        border-radius: 8px;
        overflow: hidden;
        background: white;
        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    }
    .model-img {
        width: 100%;
        height: 200px;
        object-fit: cover;
    }
    .model-info {
        padding: 15px;
    }
    .model-name {
        color: #2563eb;
        text-decoration: none;
        font-weight: bold;
        font-size: 1.1em;
        margin-bottom: 8px;
        display: block;
    }
    .model-name:hover {
        text-decoration: underline;
    }
    .distance {
        color: #666;
        font-size: 0.9em;
    }
    </style>
    <div class="model-gallery">
    """
    
    # Generate cards for each model
    for idx, (score, model_idx) in enumerate(zip(top5_prob[0], top5_indices[0])):
        model_name = predict_image.model_names[model_idx.item()]
        distance = calculate_euclidean_distance(img_features[0], 
                                             torch.randn(512).numpy())  # Placeholder for actual features
        
        civitai_url = f"https://civitai.com/search/models?sortBy=models_v9&query={model_name}"
        
        html_output += f"""
        <div class="model-card">
            <img class="model-img" src="data:image/svg+xml,<svg xmlns='http://www.w3.org/2000/svg' width='250' height='200' viewBox='0 0 250 200'><rect width='100%' height='100%' fill='%23f0f0f0'/><text x='50%' y='50%' dominant-baseline='middle' text-anchor='middle' font-family='Arial' font-size='16' fill='%23666'>Model Preview</text></svg>" alt="{model_name}">
            <div class="model-info">
                <a href="{civitai_url}" target="_blank" class="model-name">{model_name}</a>
                <div class="distance">Euclidean Distance: {distance:.4f}</div>
            </div>
        </div>
        """
    
    html_output += "</div>"
    
    return html_output

# Gradio Interface
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.HTML(),
    title="Stable Diffusion Model Recommender",
    description="Upload an AI-generated image to get model recommendations",
    examples=[["example1.jpg"], ["example2.jpg"]]
)

demo.launch()