#!/usr/bin/env python
# coding: utf-8

# In[1]:


import cv2
import os
from pathlib import Path
import numpy as np

def video_to_frames(video_path, output_folder, frame_size=(512, 512)):
    """
    Extracts frames from a video and saves them as images.

    Args:
        video_path (str): Path to the input video file.
        output_folder (str): Directory where extracted frames will be saved.
        frame_size (tuple): Desired size of output frames (width, height).
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Open the video file
    cap = cv2.VideoCapture(video_path)
    frame_count = 0
    success = True

    file_name = Path(video_path).stem
    previous_frame = None

    while success:
        success, frame = cap.read()
        if success:
            frame = cv2.resize(frame, frame_size)
            if previous_frame is not None and np.array_equal(frame, previous_frame):
                continue

            frame_filename = os.path.join(output_folder, f"{file_name}_{frame_count:06d}.png")
            cv2.imwrite(frame_filename, frame)
            frame_count += 1
            previous_frame = frame

    cap.release()
    print(f"Extracted {frame_count} frames from {video_path} to {output_folder}")


# In[2]:


from os import listdir
from os.path import isfile, join

videos_path = "/home/muinez/video_ai/videos/"
output_folder = '/home/muinez/video_ai/frames'

# for r, d, f in os.walk(videos_path):
# 	for file in f:
# 		if file.endswith(".mp4") or file.endswith(".gif"):
# 			video_to_frames(videos_path + file, output_folder)


# In[3]:


import torch
import torch.nn as nn


# In[4]:


from PIL import Image
import requests
from transformers import AutoProcessor, SiglipVisionModel
import torch


processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")


# In[ ]:





# In[5]:


import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, in_dim):
        super(SelfAttention, self).__init__()
        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        batch_size, C, width, height = x.size()
        proj_query = self.query_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1)
        proj_key = self.key_conv(x).view(batch_size, -1, width * height)
        energy = torch.bmm(proj_query, proj_key)
        attention = F.softmax(energy, dim=-1)
        proj_value = self.value_conv(x).view(batch_size, -1, width * height)

        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        out = out.view(batch_size, C, width, height)

        out = self.gamma * out + x
        return out

class ResidualLiner(nn.Module):
    def __init__(self, in_dim, out_dim, use_activation = None):
        super(ResidualLiner, self).__init__()
        self.linear1 = nn.Linear(in_dim, out_dim)
        self.linear2 = nn.Linear(in_dim, out_dim)
        
        if(use_activation):
            self.activation = nn.ReLU()
        else:
            self.activation = None
        

    def forward(self, x):
        y = self.linear1(x)
        x = self.linear2(x) + y
        
        if(self.activation is not None):
            return self.activation(x)
        
        return x
    
class ResidualAttention(nn.Module):
    def __init__(self, dim):
        super(ResidualAttention, self).__init__()
        self.linear1 = ResidualLiner(dim, dim)
        self.attenton = nn.MultiheadAttention(dim, 8, batch_first=True)
                

    def forward(self, x):
        x = self.linear1(x)
        attn_output, attn_weights = self.attenton(x, x, x)

        return x + attn_output

class CustomSigmoid(nn.Module):
    def __init__(self):
        super(CustomSigmoid, self).__init__()
        self.scale = 4  # Scaling factor
        self.shift = -2 # Shifting factor

    def forward(self, x):
        # Apply the standard sigmoid function
        sigmoid = torch.sigmoid(x)
        # Scale and shift the sigmoid output
        return self.scale * sigmoid + self.shift
    

class Multiply(nn.Module):
    def __init__(self):
        super(Multiply, self).__init__()
        self.scale = nn.Parameter(torch.ones(1))

    def forward(self, x):
        print(self.scale.item())
        return x * self.scale

class Decoder2(nn.Module):
    def __init__(self, embedding_dim=256):
        super(Decoder2, self).__init__()

        self.decoder = nn.Sequential(
            ResidualLiner(768, 768),
            ResidualAttention(768),
            ResidualLiner(768, 768),
            ResidualLiner(768, 768),
            ResidualAttention(768),
            ResidualLiner(768, 768),
            ResidualLiner(768, 768),
            ResidualAttention(768),
            ResidualLiner(768, 768),
            ResidualLiner(768, 768),
            ResidualAttention(768),
            ResidualLiner(768, 768),
            ResidualLiner(768, 768, False),
            nn.Sigmoid(),
            #Multiply()
        )

        hiddem_dim = 768
        latent_dim = 768

        self.fc2_mean = nn.Linear(hiddem_dim, latent_dim)
        self.fc2_logvar = nn.Linear(hiddem_dim, latent_dim)
        #self.decoder_input = nn.Linear(latent_dim, hiddem_dim)

    def encode(self, x):
        x = self.first_frame_conv(x)
        x = self.second_frame_conv(x)
        
        mu, logvar = torch.split(x, 512, dim=1)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        #x = self.decoder_input(z)
        x = self.decoder(z)
        return F.sigmoid(x)

    def forward(self, x):
        mu = self.fc2_mean(x)
        logvar = self.fc2_logvar(x)
        z = self.reparameterize(mu, logvar)
        z = self.decode(z)

        z = z.view(-1, 14, 14, 3, 16, 16)  # Reshape to [batch_size, 14, 14, 3, 16, 16]
        z = z.permute(0, 3, 1, 4, 2, 5)  # Permute to [batch_size, 3, 14, 16, 14, 16]
        z = z.contiguous().view(-1, 3, 224, 224)  # Reshape to [batch_size, 3, 224, 224]

        return z, mu, logvar
    
def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return(BCE , KLD)


# In[6]:


from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import glob


class ImageDataset(Dataset):
    def __init__(self, frames_folder, preprocessor):
        """
        Args:
            frames_folder (str): Path to the folder containing extracted frames.
        """
        self.preprocessor = preprocessor
        self.directory = frames_folder
        self.frames_paths = self.load_images_in_pairs(frames_folder)
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((224, 224)),
            #transforms.Normalize(mean=[0.485, 0.456, 0.406],  # ImageNet means
            #                     std=[0.229, 0.224, 0.225])   # ImageNet stds
        ])

    def __len__(self):
        return len(self.frames_paths)   # Minus one to create pairs

    def __getitem__(self, idx):
        # Load two consecutive frames
        frame1path = self.frames_paths[idx]
        frame1 = Image.open(self.directory + '/' + frame1path).convert('RGB')

        # Apply transformations
        processed = self.preprocessor(images=frame1, return_tensors="pt")

        # Return individual frames and the concatenated pair
        return {'frame1': processed['pixel_values'][0], 'real_frame': self.transform(frame1)}
    
    def load_images_in_pairs(self, directory):
        image_pairs = []

        # Iterate over all files in the specified directory
        for filename in os.listdir(directory):
            if filename.endswith('.png') or filename.endswith('.jpg'):
                image_pairs.append(filename)

        return image_pairs


# In[7]:


import torch
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np

# Define the inverse normalization transform
def denormalize(tensor):
    """
    Denormalizes a tensor using the mean and std used during preprocessing.
    Args:
        tensor (torch.Tensor): Normalized tensor.
    Returns:
        torch.Tensor: Denormalized tensor.
    """
    # Mean and std used during preprocessing (ImageNet statistics)
    #mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
    #std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
    #tensor = tensor * std + mean
    return tensor

def postprocess(tensor):
    tensor = tensor.detach().cpu()
    tensor = denormalize(tensor)
    tensor = torch.clamp(tensor, 0, 1)
    tensor_np = tensor.numpy()
    tensor_np = np.transpose(tensor_np, (1, 2, 0))  
    return tensor_np

def save_image(tensor, file_path):
    image_np = postprocess(tensor)
    image_np = image_np * 255
    image = Image.fromarray(image_np.astype(np.uint8))
    image.save(file_path)


# In[8]:


frames_folder = 'frames'
#dataset = FrameDataset(frames_folder)
dataset = ImageDataset('/mnt/h/gelbooru_all/randomshit', processor)
dataloader = DataLoader(dataset, batch_size=96, shuffle=True
						, num_workers=6
						)
len(dataset)


# In[9]:


siglip = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224").to('cuda')



# In[10]:


import torch.optim as optim

embedding_dim = 512
decoder = Decoder2(embedding_dim=embedding_dim)

# Loss function and optimizer
criterion = nn.MSELoss()

optimizer = optim.Adam(list(decoder.parameters()) 
					   #+ list(siglip.parameters()														 )
					   , lr=1e-4)


# In[11]:


# import torch

# tensor = torch.rand((2, 4, 3, 4))
# print(tensor)
# #print(tensor.view(4,3,2,2))
# print(tensor.view(2, 4,3,2,2).permute(0, 2, 3, 1, 4).reshape(2, 3, 4, 4))


# In[ ]:


#test single frame
siglip.train()
device = 'cuda'

# Initialize models
decoder = decoder.to(device)


accumulation_steps = 1
# Training loop
num_epochs = 55
steps = 0
current_bce_loss = 0
current_kld_loss = 0
for epoch in range(num_epochs):
    for batch in dataloader:
        frame1 = batch['frame1'].to(device) 
        real_frame = batch['real_frame'].to(device) 
        #print(real_frame.min())
        #print(real_frame.max())
        real_frame = torch.clamp(real_frame, 0, 1)
        hidden_states = siglip(frame1, output_hidden_states=True)['hidden_states'][-1]
        # Forward pass
        reconstructed_frame, mu, logvar = decoder(hidden_states)  # [batch_size, 3, 512, 512]

        # Compute loss
        BCE , KLD = loss_function(reconstructed_frame, real_frame, mu, logvar) 
        
        loss = (BCE + KLD) / accumulation_steps
        loss.backward()

        if (steps + 1) % accumulation_steps == 0: 
            optimizer.step()
            optimizer.zero_grad()

        steps+=1
        current_bce_loss += BCE.item()
        current_kld_loss += KLD.item()
        if steps % 5 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}] steps: {steps}, BCE: {current_bce_loss/5:.4f}, KLD: {current_kld_loss/5:.4f}")
            current_bce_loss = 0
            current_kld_loss = 0
            save_image(real_frame[0], f'samples/{steps}_real.png')
            save_image(reconstructed_frame[0], f'samples/{steps}.png')
            


# In[35]:


torch.save(encoder, "encoder.pth")
torch.save(decoder, "decoder.pth")


# In[9]:


device = 'cuda'

# Initialize models



# Training loop
num_epochs = 5
steps = 0
current_loss = 0
for epoch in range(num_epochs):
    for batch in dataloader:
        frame1 = batch['frame1']  # [batch_size, 3, 512, 512]
        frame2 = batch['frame2']  # [batch_size, 3, 512, 512]
        frame_pair = batch['frame_pair']  # [batch_size, 6, 512, 512]

        frame1 = frame1.to(device)
        frame2 = frame2.to(device)
        frame_pair = frame_pair.to(device)
        encoder = encoder.to(device)
        decoder = decoder.to(device)

        # Forward pass
        embedding = encoder(frame_pair)  # [batch_size, embedding_dim]
        reconstructed_frame = decoder(frame1, embedding)  # [batch_size, 3, 512, 512]

        # Compute loss
        loss = criterion(reconstructed_frame, frame2)

        # Backpropagation and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        steps+=1
        current_loss += loss.item()
        if steps % 5 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}] steps: {steps}, Loss: {current_loss/10:.4f}")
            current_loss = 0
            save_image(reconstructed_frame, f'samples/{steps}.png')
        


# In[31]:


def display_sample(tensor, axs, name):
    tensor = postprocess(tensor)
    fig, axs = plt.subplots(1, 2, figsize=(12, 6))

# Original second frame
    axs.imshow(tensor)
    axs.set_title(name)
    axs.axis('off')


# In[39]:


import torch
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np

# Define the inverse normalization transform
def denormalize(tensor):
    """
    Denormalizes a tensor using the mean and std used during preprocessing.
    Args:
        tensor (torch.Tensor): Normalized tensor.
    Returns:
        torch.Tensor: Denormalized tensor.
    """
    # Mean and std used during preprocessing (ImageNet statistics)
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
    tensor = tensor * std + mean
    return tensor

def display_sample(tensor, ax, name):
    tensor = tensor.detach().cpu()
    tensor = denormalize(tensor)
    tensor = torch.clamp(tensor, 0, 1)
    tensor_np = tensor.numpy()
    tensor_np = np.transpose(tensor_np, (1, 2, 0))

    ax.imshow(tensor_np)
    ax.set_title(name)
    ax.axis('off')

def display_batch_index(index):
	fig, axs = plt.subplots(1, 3, figsize=(12, 6))
    
	display_sample(batch['frame1'][index], axs[0], 'Frame1')
	#display_sample(batch['frame2'][index], axs[1], 'Frame2')
	display_sample(reconstructed_frame[index], axs[2], 'Reconstructed') #

	plt.show()


display_batch_index(4)


# In[15]:


torch.save(encoder, "encoder.pth")
torch.save(decoder, "decoder.pth")



find issue why reconstructed images are pixel art like, like squares 16x16 per each patch, has almost single color, like inside patch there are almost only one color