# import gradio as gr
import os
import glob
import time
import numpy as np
from PIL import Image
from skimage import io, color

import torch
import torch.nn.functional as F
from torch import nn, optim
from torchvision import transforms
import segmentation_models_pytorch as smp
import gradio as gr

ENCODER = 'resnet18'
ENCODER_WEIGHTS = 'imagenet'
ACTIVATION = 'tanh'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

ResUNet = smp.Unet(
    encoder_name=ENCODER,
    encoder_weights=ENCODER_WEIGHTS,
    in_channels=1,
    classes=2,
    activation=ACTIVATION,
)

class Discriminator(nn.Module):
    def __init__(self, in_channels):
        super(Discriminator, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=4, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.conv5 = nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        return x

def init_weights(net, init_gain=0.02):
    def init(m):
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
            nn.init.normal_(m.weight.data, 0.0, init_gain)
            if m.bias is not None:
                nn.init.constant_(m.bias.data, 0.0)
        elif isinstance(m, (nn.BatchNorm2d)):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)
    net.apply(init)
    return net

class GAN(nn.Module):
    def __init__(self, in_channels, out_channels, generator=None):
        super(GAN, self).__init__()
        self.generator = generator if generator else init_weights(ResUNet(in_channels, out_channels))
        self.generator.to(device)
        self.discriminator = init_weights(Discriminator(in_channels=3).to(device))

        self.GANLoss = nn.BCEWithLogitsLoss()
        self.L1Loss = nn.L1Loss()

        self.gen_optim = optim.Adam(self.generator.parameters(), lr=3e-5, betas=(0.5, 0.99))
        self.disc_optim = optim.Adam(self.discriminator.parameters(), lr=3e-5, betas=(0.5, 0.99))



    def forward(self, L, ab, train):
        self.L = L
        self.ab = ab

        # Generate fake images
        self.fake_ab = self.generator(L)


        fake_img = torch.cat((self.L, self.fake_ab), dim=1)
        disc_fake = self.discriminator(fake_img.detach())
        disc_real = self.discriminator(torch.cat((self.L, self.ab), dim=1))
        self.disc_loss = (self.GANLoss(disc_real, torch.ones_like(disc_real)) + self.GANLoss(disc_fake, torch.zeros_like(disc_fake)))/2

        if train:
            self.disc_optim.zero_grad()
            self.disc_loss.backward()
            self.disc_optim.step()


        # Train generator

        disc_fake = self.discriminator(fake_img)
        gan_loss = self.GANLoss(disc_fake, torch.ones_like(disc_fake))
        L1_loss = self.L1Loss(self.fake_ab, self.ab) * 100
        self.gen_loss = gan_loss + L1_loss

        if train:
            self.gen_optim.zero_grad()
            self.gen_loss.backward()
            self.gen_optim.step()


        return self.gen_loss.item(), self.disc_loss.item()

gan = GAN(in_channels=1, out_channels=2, generator=ResUNet).to(device)
gan.load_state_dict(torch.load('gan.pth',  map_location=torch.device('cpu'),  weights_only=False))

def predict(img):
    L = Image.open(img).convert('L')
    L = transforms.Resize((256, 256), Image.BICUBIC)(L)

    L = transforms.ToTensor()(L)


    gan.eval()
    with torch.no_grad():
        ab = gan.generator(L.unsqueeze(0).to(device))

    ab = ab.squeeze(0).cpu()


    # Denormalize the L and ab channels
    L = L.numpy() * 100          # L is in range [0, 1], scale to [0, 100]
    ab = ab.numpy() * 128        # ab is in range [-1, 1], scale to [-128, 128]

    # Combine L and ab into a single LAB image
    lab_image = np.stack([L[0], ab[0], ab[1]], axis=-1)  # Shape: [256, 256, 3]

    # Convert LAB to RGB using skimage's color.lab2rgb
    rgb_image = color.lab2rgb(lab_image)
    rgb_image = (rgb_image * 255).astype(np.uint8)
    rgb_image = Image.fromarray(rgb_image)
    return rgb_image

image = gr.components.Image(type='filepath')
output = gr.components.Image(type='pil')

demo = gr.Interface(fn=predict, inputs=image, outputs=output, title='Colorize Black and White Images', examples=[['examples/1.jpg'], ['examples/2.jpg']])
demo.launch()