# -- Code Cell --
from diffusers import DiffusionPipeline, AutoencoderKL
from transformers import pipeline
import torch
import random
import ast
import pandas as pd
import numpy as np

# -- Code Cell --
train_data = pd.read_csv('/kaggle/input/datasets/david23200/missing-frame-v2/train_data_final.csv')
train_data['answer'] = train_data['answer'].apply(lambda x: np.array(ast.literal_eval(x)).astype(np.float32))
device = "cuda" if torch.cuda.is_available() else "cpu"

# -- Code Cell --
vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
vae.to(device, dtype=torch.float32)
vae.eval()
vae.enable_tiling()
vae.enable_slicing()


# -- Code Cell --
def get_img(index):
 img_ = torch.tensor(train_data["answer"].iloc[index])
 img_ = img_.to(device)
 with torch.no_grad():# crucial for keeping memory stable
  img = vae.decode(img_).sample[0] # make sure latents has shape of type (batch_size,4,90,160)
 return img

# -- Code Cell --
from tqdm.notebook import tqdm
imgs = [get_img(i) for i in tqdm(range(len(train_data)))]

# -- Code Cell --
imgs = torch.stack(imgs)

# -- Code Cell --
from torch.utils.data import Dataset, DataLoader
class data_idk(Dataset):
    def __init__(self, imgs):
        self.img = imgs
    def __len__(self):
        return len(self.img) - 1
    def __getitem__(self, index):
        return self.img[index], self.img[index + 1]

# -- Code Cell --
train_dt = data_idk(imgs)

# -- Code Cell --
train_loader = DataLoader(train_dt, batch_size = 4)

# -- Code Cell --
train_dt[0][0].shape

# -- Code Cell --
import torch.nn as nn
class double_conv(nn.Module):
    def __init__(self, in_ch, out_ch, mid_ch = None):
        super().__init__()
        if mid_ch == None:
            mid_ch = out_ch
        self.conv1 = nn.Conv2d(in_ch, mid_ch, kernel_size=(2, 2), stride= 1)
        self.activation = nn.ReLU()
        self.conv2 = nn.Conv2d(mid_ch, out_channels=out_ch, kernel_size=(2, 2), stride= 1)
    def forward(self, x):
        x = self.conv1(x)
        x = self.activation(x)
        x = self.conv2(x)
        x = self.activation(x)
        return x
class down_sample(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.down = nn.MaxPool2d((2, 2), stride= 2)
    def forward(self, x):
        return self.down(x)
class up_sample(nn.Module):
    def __init__(self,in_ch, out_ch, kernel_size, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.up = nn.ConvTranspose2d(in_ch, out_ch, kernel_size=(kernel_size, kernel_size), stride= 2)
    def forward(self, x):
        return self.up(x)
class double_conv_up(nn.Module):
    def __init__(self, in_ch, out_ch, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conv_block = nn.Sequential(
            nn.ConvTranspose2d(in_ch, in_ch, kernel_size=(2, 2)),
            nn.ReLU(),
            nn.ConvTranspose2d(in_ch, out_ch, kernel_size=(2, 2))
        )
    def forward(self, x):
        x= self.conv_block(x)
        return x
#UNET_Baddie

# -- Code Cell --
import torch
import torch.nn as nn
import torch.nn.functional as F

class UNET_Baddie(nn.Module):
    def __init__(self, in_ch, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.double_conv_1 = double_conv(in_ch, 8)
        self.down_1 = down_sample()
        self.double_conv_2 = double_conv(8, 16)
        self.down_2 = down_sample()
        self.double_conv_3 = double_conv(16, 32)
        self.down_3 = down_sample()
        self.double_conv_4 = double_conv(32, 64)
        self.down_4 =down_sample()
        self.bottle_neck =double_conv(64, 64)
        self.up_1 = up_sample(64, 32, 6)
        self.double_up_1 = double_conv_up(96, 16)
        self.up_2 = up_sample(16, 16, 2)
        self.double_up_2 = double_conv_up(48, 8)
        self.up_3 = up_sample(8, 8, 3)
        self.double_up_3 = double_conv_up(24, 4)
        self.up_4 = up_sample(4, 4, 2)
        self.final_goddamn_it_conv = nn.ConvTranspose2d(12, 3, kernel_size=(3, 3))
        self.signmd = nn.Sigmoid()
    def forward(self, x):
        x1 =self.double_conv_1(x)
        x = self.down_1(x1)
        x2 = self.double_conv_2(x)
        x = self.down_2(x2)
        x3 = self.double_conv_3(x)
        x = self.down_3(x3)
        x4 = self.double_conv_4(x)
        x = self.down_4(x4)
        x = self.bottle_neck(x)
        x = self.up_1(x)
        x = torch.concat([x, x4], dim = 1)
        x = self.double_up_1(x)
        x = self.up_2(x)
        x = torch.concat([x, x3], dim = 1)
        x = self.double_up_2(x)
        x =self.up_3(x)
        x = torch.concat([x, x2], dim= 1)
        x = self.double_up_3(x)
        x = self.up_4(x)
        x = torch.concat([x, x1], dim = 1)
        x = self.signmd(self.final_goddamn_it_conv(x))
        return x

# -- Code Cell --
model = UNET_Baddie(3)
model = model.to(device)

# -- Code Cell --
import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl

class RegressionModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.save_hyperparameters()

        self.model = UNET_Baddie(3)

        self.criterion = nn.MSELoss()

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)

        self.log("train_loss", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)

        self.log("val_loss", loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters())

# -- Code Cell --
lit = RegressionModel()
trainer = pl.Trainer(max_epochs= 25)
trainer.fit(lit, train_loader)

# -- Code Cell --
img = get_img(-1)
img = img.to("cuda")

# -- Code Cell --
lit.model = lit.model.to("cuda")

# -- Code Cell --
lit.model(img.unsqueeze(0)).shape

# -- Code Cell --


# -- Code Cell --
img = get_img(-1).unsqueeze(0)
imgs = []
img = img.to("cuda")
for i in range(5):
    imgs.append(img)
    img = lit.model(img)

# -- Code Cell --
from torchvision.transforms import v2
trns =v2.Compose([
    v2.Resize((360,640)),
    v2.ToPILImage()
])

# -- Code Cell --
len(imgs)

# -- Code Cell --
imgs[0].shape

# -- Code Cell --
df_pred = pd.DataFrame({
    "subtaskID" : 1,
    "datapointID" :[i + 41 for i in range(5)],
    "answer" :[np.array(trns(x.squeeze(0))).tolist() for x in imgs]
})

# -- Code Cell --
df_pred.to_csv("subi.csv", index = False)