import numpy as np
import torch
import torch.nn as nn
import gradio as gr
from PIL import Image
import torchvision.transforms as transforms
import os  # 📁 For file operations

# 🧠 Neural network layers
norm_layer = nn.InstanceNorm2d

# 🧱 Building block for the generator
class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()

        conv_block = [  nn.ReflectionPad2d(1),
                        nn.Conv2d(in_features, in_features, 3),
                        norm_layer(in_features),
                        nn.ReLU(inplace=True),
                        nn.ReflectionPad2d(1),
                        nn.Conv2d(in_features, in_features, 3),
                        norm_layer(in_features)
                        ]

        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x):
        return x + self.conv_block(x)

# 🎨 Generator model for creating line drawings
class Generator(nn.Module):
    def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
        super(Generator, self).__init__()

        # 🏁 Initial convolution block
        model0 = [   nn.ReflectionPad2d(3),
                    nn.Conv2d(input_nc, 64, 7),
                    norm_layer(64),
                    nn.ReLU(inplace=True) ]
        self.model0 = nn.Sequential(*model0)

        # 🔽 Downsampling
        model1 = []
        in_features = 64
        out_features = in_features*2
        for _ in range(2):
            model1 += [  nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                        norm_layer(out_features),
                        nn.ReLU(inplace=True) ]
            in_features = out_features
            out_features = in_features*2
        self.model1 = nn.Sequential(*model1)

        # 🔁 Residual blocks
        model2 = []
        for _ in range(n_residual_blocks):
            model2 += [ResidualBlock(in_features)]
        self.model2 = nn.Sequential(*model2)

        # 🔼 Upsampling
        model3 = []
        out_features = in_features//2
        for _ in range(2):
            model3 += [  nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                        norm_layer(out_features),
                        nn.ReLU(inplace=True) ]
            in_features = out_features
            out_features = in_features//2
        self.model3 = nn.Sequential(*model3)

        # 🎭 Output layer
        model4 = [  nn.ReflectionPad2d(3),
                        nn.Conv2d(64, output_nc, 7)]
        if sigmoid:
            model4 += [nn.Sigmoid()]

        self.model4 = nn.Sequential(*model4)

    def forward(self, x, cond=None):
        out = self.model0(x)
        out = self.model1(out)
        out = self.model2(out)
        out = self.model3(out)
        out = self.model4(out)

        return out

# 🔧 Load the models
model1 = Generator(3, 1, 3)
model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu'), weights_only=True))
model1.eval()

model2 = Generator(3, 1, 3)
model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu'), weights_only=True))
model2.eval()

# 🖼️ Function to process the image and create line drawing
def predict(input_img, ver):
    # Open the image and get its original size
    original_img = Image.open(input_img)
    original_size = original_img.size

    # Define the transformation pipeline
    transform = transforms.Compose([
        transforms.Resize(256, Image.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # Apply the transformation
    input_tensor = transform(original_img)
    input_tensor = input_tensor.unsqueeze(0)

    # Process the image through the model
    with torch.no_grad():
        if ver == 'Simple Lines':
            output = model2(input_tensor)
        else:
            output = model1(input_tensor)

    # Convert the output tensor to an image
    output_img = transforms.ToPILImage()(output.squeeze().cpu().clamp(0, 1))

    # Resize the output image back to the original size
    output_img = output_img.resize(original_size, Image.BICUBIC)

    return output_img

# 📝 Title for the Gradio interface
title="🖌️ Image to Artistic Drawing"

# 🖼️ Dynamically generate examples from images in the directory
examples = []
image_dir = '.'  # Assuming images are in the current directory
for file in os.listdir(image_dir):
    if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
        examples.append([file, 'Simple Lines'])
        examples.append([file, 'Complex Lines'])

# 🚀 Create and launch the Gradio interface
iface = gr.Interface(
    fn=predict, 
    inputs=[
        gr.Image(type='filepath'),
        gr.Radio(['Complex Lines', 'Simple Lines'], label='version', value='Simple Lines')
    ],
    outputs=gr.Image(type="pil"),
    title=title,
    examples=examples
)

iface.launch()