import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import kornia

from PIL import Image
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2


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


# Define model

class BlurUpSample(nn.Module):
    def __init__(self, c):
        super(BlurUpSample, self).__init__()
        self.blurpool =  kornia.filters.GaussianBlur2d((3, 3), (1.5, 1.5))
        self.upsample = nn.Upsample(scale_factor=(2, 2), mode='bilinear', align_corners=False)

    def forward(self, x):
        x = self.blurpool(x)
        x = self.upsample(x)

        return x

class DownLayer(nn.Module):
    def __init__(self, c_in, c_out):
        super(DownLayer, self).__init__()
        self.maxblurpool = kornia.filters.MaxBlurPool2D(kernel_size=3)
        self.conv1 = nn.Conv2d(c_in, c_out, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(c_out)
        self.leakyrelu = nn.LeakyReLU(inplace=True)
        self.conv2 = nn.Conv2d(c_out, c_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(c_out)

    def forward(self, x):
        x = self.maxblurpool(x)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.leakyrelu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.leakyrelu(x)
        return x


class UpLayer(nn.Module):
    def __init__(self, c_in, c_out):
        super(UpLayer, self).__init__()
        self.upsample = BlurUpSample(c_in)
        self.conv1 = nn.Conv2d(c_in+ c_out, c_out, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(c_out)
        self.leakyrelu = nn.LeakyReLU(inplace=True)
        self.conv2 = nn.Conv2d(c_out, c_out, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(c_out)

    def forward(self, x, skip_x):
        x = self.upsample(x)

        dh = skip_x.size(2) - x.size(2)
        dw = skip_x.size(3) - x.size(3)

        x = F.pad(x, (dw // 2, dw - dw // 2, dh // 2, dh - dh // 2))

        x = torch.cat([x, skip_x], dim=1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.leakyrelu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.leakyrelu(x)
        return x

class Generator(nn.Module):
  def __init__(self):
    super(Generator, self).__init__()
    self.conv1 = nn.Conv2d(5, 64, kernel_size=3, stride=1, padding=1)
    self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
    self.batchnorm1 = nn.BatchNorm2d(64)
    self.leakyrelu = nn.LeakyReLU(inplace=True)
    self.downlayer1 = DownLayer(64, 128)
    self.downlayer2 = DownLayer(128, 256)
    self.downlayer3 = DownLayer(256, 512)
    self.downlayer4 = DownLayer(512, 1024)
    self.uplayer1 = UpLayer(1024, 512)
    self.uplayer2 = UpLayer(512, 256)
    self.uplayer3 = UpLayer(256, 128)
    self.uplayer4 = UpLayer(128, 64)
    self.conv3 = nn.Conv2d(64, 3, kernel_size=1, stride=1, padding=0)

  def forward(self, x):
    #print(f'Input Shape: {x.shape}')
    x1 = self.conv1(x)
    x1 = self.batchnorm1(x1)
    x1 = self.leakyrelu(x1)
    x1 = self.conv2(x1)
    x1 = self.batchnorm1(x1)
    x1 = self.leakyrelu(x1)

    #print(f'Processed Input Shape: {x.shape}')

    x2 = self.downlayer1(x1)
    x3 = self.downlayer2(x2)
    x4 = self.downlayer3(x3)
    x5 = self.downlayer4(x4)

    #print(f'Done Downlayering... Shape: {x5.shape}')

    x = self.uplayer1(x5, x4)
    x = self.uplayer2(x, x3)
    x = self.uplayer3(x, x2)
    x = self.uplayer4(x, x1)
    x = self.conv3(x)

    #print(f'Output Shape: {x.shape}')
    return x


transform_resize = A.Compose([
    A.Resize(512, 512),
    ToTensorV2(),
])




# Load model
generator_model = Generator()
generator_model.load_state_dict(torch.load('large-aging-model.h5',map_location=torch.device(device)))
generator_model.to(device)
#generator_model.eval()
print("")



def age_filter(image, input_age, output_age):

    resized_image = image.resize((512,512))
    
    input_image = transform_resize(image=np.array(image))['image']/255


    age_map1 = torch.full((1, 512, 512), input_age / 100)
    age_map2 = torch.full((1, 512, 512),  output_age / 100)

    input_tensor = torch.cat((input_image, age_map1,age_map2), dim=0)
    
    with torch.no_grad():
        model_output = generator_model(input_tensor.unsqueeze(0).to(device))
        
    np_test = np.array(image)

    new_image = (model_output.squeeze(0).cpu().permute(1,2,0).numpy()*255+np.array(resized_image)).astype('uint8')

    sample_image = np.array(Image.fromarray(new_image).resize((np_test.shape[1],np_test.shape[0]))).astype('uint8')
    return sample_image

import gradio as gr
from torchvision.transforms.functional import crop



def crop_and_process_image(input_img,input_age,output_ag):
    # Crop the image using the provided crop tool coordinates
    processed_image = Image.fromarray(input_img)  # Modify this line to preprocess the cropped image

    # Run the processed image through your model
    output = age_filter(processed_image, input_age, output_ag)
    # Return the output
    return output

input_image = gr.Image(label="Input Image", interactive=True)

output_image = gr.Image(label="Output Image", type="pil")

input_age = gr.Slider(label="Current Age",minimum=18,maximum=80)
output_age = gr.Slider(label="Desired Age",minimum=18,maximum=80)


def process_image(input_img,input_age,output_age):

    output = crop_and_process_image(input_img,input_age,output_age)

    output = Image.fromarray(output)
    return output
    
description="Enter age of input image and desired age. Crop out background. Better results on high resolution (512x512 face). To avoid background/hair artifacts, use with a face parser." 



# Create the Gradio interface
gr.Interface(fn=process_image, inputs=[input_image,input_age,output_age], outputs=output_image,description=description, title="Age Transformation",cache_examples = True, examples=[["example_image_input2.jpeg",30,65],["example_image_output.png",20,40]]).launch()