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

# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Image loader
def load_image(image, max_size=400):
    transform = transforms.Compose([
        transforms.Resize(max_size),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x[:3, :, :]),  # remove alpha channel
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    image = transform(image).unsqueeze(0)
    return image.to(device)

# Reconvert image
def tensor_to_pil(tensor):
    image = tensor.cpu().clone().squeeze(0)
    image = image * torch.tensor([0.229, 0.224, 0.225]).view(3,1,1) + torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
    image = image.clamp(0, 1)
    image = transforms.ToPILImage()(image)
    return image

# VGG19 Model
class VGGFeatures(nn.Module):
    def __init__(self):
        super(VGGFeatures, self).__init__()
        vgg = models.vgg19(pretrained=True).features.eval()
        self.slice = nn.Sequential(*list(vgg.children())[:21])
        for param in self.slice.parameters():
            param.requires_grad = False

    def forward(self, x):
        return self.slice(x)

# Gram Matrix
def gram_matrix(tensor):
    b, c, h, w = tensor.size()
    features = tensor.view(b * c, h * w)
    G = torch.mm(features, features.t())
    return G.div(b * c * h * w)

# Style Transfer Function
def style_transfer(content_img, style_img, steps=300, style_weight=1e6, content_weight=1):
    content = load_image(content_img)
    style = load_image(style_img)
    generated = content.clone().requires_grad_(True)

    model = VGGFeatures().to(device)
    optimizer = torch.optim.LBFGS([generated])

    content_features = model(content)
    style_features = model(style)
    style_gram = gram_matrix(style_features)

    def closure():
        optimizer.zero_grad()
        generated_features = model(generated)
        generated_gram = gram_matrix(generated_features)

        content_loss = content_weight * nn.functional.mse_loss(generated_features, content_features)
        style_loss = style_weight * nn.functional.mse_loss(generated_gram, style_gram)

        total_loss = content_loss + style_loss
        total_loss.backward()
        return total_loss

    for i in range(steps):
        optimizer.step(closure)

    return tensor_to_pil(generated)

# Gradio Interface
interface = gr.Interface(
    fn=style_transfer,
    inputs=[
        gr.Image(type="pil", label="Content Image"),
        gr.Image(type="pil", label="Style Image")
    ],
    outputs=gr.Image(type="pil", label="Stylized Output"),
    title="🎨 Neural Style Transfer",
    description="Upload a content image and a style image. The model will apply the style to the content!"
)

if __name__ == "__main__":
    interface.launch()