# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/13tu6v1reMxLATyBwle-BgpQrql9p4nqn
"""
#import os
#os.system("pip install fastai")
#from fastai.vision.all import *
#from fastai.basics import *

"""cyclegan_inference.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/12lelsBZXqNOe7xaXI724rEHAbppRt07y
"""
import gradio as gr
import torch
import torchvision
from torch import nn
from typing import List

#def ifnone(a, b): # a fastai-specific (fastcore) function used below, redefined so it's independent
#    "`b` if `a` is None else `a`"
#    return b if a is None else a

class ConvBlock(torch.nn.Module):
    def __init__(self,input_size,output_size,kernel_size=4,stride=2,padding=1,activation='relu',batch_norm=True):
        super(ConvBlock,self).__init__()
        self.conv = torch.nn.Conv2d(input_size,output_size,kernel_size,stride,padding)
        self.batch_norm = batch_norm
        self.bn = torch.nn.InstanceNorm2d(output_size)
        self.activation = activation
        self.relu = torch.nn.ReLU(True)
        self.lrelu = torch.nn.LeakyReLU(0.2,True)
        self.tanh = torch.nn.Tanh()
        self.sigmoid = torch.nn.Sigmoid()
    def forward(self,x):
        if self.batch_norm:
            out = self.bn(self.conv(x))
        else:
            out = self.conv(x)
        
        if self.activation == 'relu':
            return self.relu(out)
        elif self.activation == 'lrelu':
            return self.lrelu(out)
        elif self.activation == 'tanh':
            return self.tanh(out)
        elif self.activation == 'no_act':
            return out
        elif self.activation =='sigmoid':
            return self.sigmoid(out)

    
class ResnetBlock(torch.nn.Module):
    def __init__(self,num_filter,kernel_size=3,stride=1,padding=0):
        super(ResnetBlock,self).__init__()
        conv1 = torch.nn.Conv2d(num_filter,num_filter,kernel_size,stride,padding)
        conv2 = torch.nn.Conv2d(num_filter,num_filter,kernel_size,stride,padding)
        bn = torch.nn.InstanceNorm2d(num_filter)
        relu = torch.nn.ReLU(True)
        pad = torch.nn.ReflectionPad2d(1)
        
        self.resnet_block = torch.nn.Sequential(
            pad,
            conv1,
            bn,
            relu,
            pad,
            conv2,
            bn
            )
    def forward(self,x):
        out = self.resnet_block(x)
        return out

class DeconvBlock(torch.nn.Module):
    def __init__(self,input_size,output_size,kernel_size=4,stride=2,padding=1,activation='relu',batch_norm=True):
        super(DeconvBlock,self).__init__()
        self.deconv = torch.nn.ConvTranspose2d(input_size,output_size,kernel_size,stride,padding)
        self.batch_norm = batch_norm
        self.bn = torch.nn.InstanceNorm2d(output_size)
        self.activation = activation
        self.relu = torch.nn.ReLU(True)
        self.tanh = torch.nn.Tanh()
    def forward(self,x):
        if self.batch_norm:
            out = self.bn(self.deconv(x))
        else:
            out = self.deconv(x)
        if self.activation == 'relu':
            return self.relu(out)
        elif self.activation == 'lrelu':
            return self.lrelu(out)
        elif self.activation == 'tanh':
            return self.tanh(out)
        elif self.activation == 'no_act':
            return out

class Generator(torch.nn.Module):
    def __init__(self,input_dim,num_filter,output_dim,num_resnet):
        super(Generator,self).__init__()
        
        #Encoder
        self.conv1 = ConvBlock(input_dim,num_filter,kernel_size=4,stride=2,padding=1)
        self.conv2 = ConvBlock(num_filter,num_filter*2)
        #Resnet blocks
        self.resnet_blocks = []
        for i in range(num_resnet):
            self.resnet_blocks.append(ResnetBlock(num_filter*2))
        self.resnet_blocks = torch.nn.Sequential(*self.resnet_blocks)
        #Decoder
        self.deconv1 = DeconvBlock(num_filter*2,num_filter)
        self.deconv2 = DeconvBlock(num_filter,output_dim,activation='tanh')
    
    def forward(self,x):
        #Encoder
        enc1 = self.conv1(x)
        enc2 = self.conv2(enc1)
        #Resnet blocks
        res = self.resnet_blocks(enc2)
        #Decoder
        dec1 = self.deconv1(res)
        dec2 = self.deconv2(dec1)
        return dec2
    

model = Generator(3, 32, 3, 4).cpu() # input_dim, num_filter, output_dim, num_resnet
model.load_state_dict(torch.load('G_A_HW4_SAVE.pt',map_location=torch.device('cpu')))
print(model)
model.eval()

totensor = torchvision.transforms.ToTensor()
normalize_fn = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
topilimage = torchvision.transforms.ToPILImage()

def predict(input_1):
    im1 = normalize_fn(totensor(input_1))
    print(im1.shape)
    preds1 = model(im1.unsqueeze(0))/2 + 0.5
    print(preds1.shape)
    return topilimage(preds1.squeeze(0).detach())

gr_interface = gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(256,256)), outputs="image", title='Emoji_CycleGAN').launch()