import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
from sklearn.decomposition import PCA
import torch.nn as nn
from einops import rearrange
from torchsummary import summary
import torch.optim as optim
import math
from tqdm import tqdm
import cv2
import numpy as np
from torch import Tensor



import lightning as L

# Set global seed to compare same initialized weights
torch.manual_seed(0)

class VQ_VAE(L.LightningModule):
    def __init__(self, input_dim, codebook_size, encoding_dim, hidden_dim, beta):
        super().__init__()

        assert len(input_dim) == 3, "Input dimension must be 3D"
        assert len(encoding_dim) == 3, "Encoding dimension must be 3D"

        # Set the dimensions
        self.input_dim = input_dim
        self.codebook_size = codebook_size
        self.encoding_dim = encoding_dim
        self.hidden_dim = hidden_dim
        self.beta = beta

        # Unwrap input dimension, extract the number of channels, height and width
        self._input_c = input_dim[0]
        self._input_h = input_dim[1]
        self._input_w = input_dim[2]
        self._input_dim_flat = self._input_c * self._input_h * self._input_w

        # Unwrap encoding dimension
        self._embedding_dim = encoding_dim[0]
        self._encoding_h = encoding_dim[1]
        self._encoding_w = encoding_dim[2]

        # Calculate flat encoding dimension
        self._encoding_dim_flat = self._embedding_dim * self._encoding_h * self._encoding_w
    

        # Setup encoder layers
        self.enc_fc1 = nn.Linear(self._input_dim_flat, self.hidden_dim)
        self.enc_fc2 = nn.Linear(self.hidden_dim, self._encoding_dim_flat)
        self.relu = nn.ReLU()

        # Setup decoder layers
        self.dec_fc1 = nn.Linear(self._encoding_dim_flat, self.hidden_dim)
        self.dec_fc2 = nn.Linear(self.hidden_dim, self._input_dim_flat)
        self.sigmoid = nn.Sigmoid()

        # Setup codebook embeddings
        self.codebook = nn.Embedding(self.codebook_size, self._embedding_dim)
        assert self.codebook.weight.requires_grad == True, "Codebook should be learnable"
        # TODO: Initialize codebook uniformly, like in the paper
        nn.init.uniform_(self.codebook.weight, -1/self.codebook_size, 1/self.codebook_size)

        # For now we initialize randomly the codebook vectors. We act like stupid, then if it works, we use uniform like in other implementations.
 
    def encode(self, x):
        """
        x: Input image of shape (b, c, h, w)
        Output: z - encoded image of shape (b, h, w, c)
        """
        assert x.shape[1:] == (self._input_c, self._input_h, self._input_w)

        # Run the layers
        x = rearrange(x, 'b c h w -> b (c h w)') # Flatten the input
        x = self.relu(self.enc_fc1(x))
        x = self.relu(self.enc_fc2(x))

        # Unflatten, last dimension is embedding dim
        x = rearrange(x, 'b (c h w) -> b h w c', c=self._embedding_dim, h=self._encoding_h, w=self._encoding_w)

        return x
    
    def decode(self, z_q):
        """        
        Input: z_q: Quantized encodings of shape (b h w c)

        Returns the reconstructed image x_hat of shape (b c h w)
        """
        # Rearrange so we work on the values, also flatten because of fully-connected layers
        z_q = rearrange(z_q, 'b h w c -> b (c h w)')

        # Run the layers
        x = self.relu(self.dec_fc1(z_q))
        x = self.relu(self.dec_fc2(x))
        x_hat = self.sigmoid(x)

        # Rearrange back to image shape
        x_hat = rearrange(x_hat, 'b (c h w) -> b c h w', c=self._input_c, h=self._input_h, w=self._input_w)

        return x_hat
    
    def quantize(self, z):
        """
        z_e: (batch_size, encoding_height, encoding_width, embedding_dim)

        Get closest (euclidean distance) codebook vector z_q given z.

        Returns quantized encodings z_q of shape (b h w c)
        """
        assert z.shape[1:] == (self._encoding_h, self._encoding_w, self._embedding_dim)

        flat_input = rearrange(z, 'b h w c -> (b h w) c')

        # TODO: Replace this by my own distances calculation
        # Equation: ||z - e||^2 = ||z||^2 + ||e||^2 - 2 * <z,e>
        distances = (torch.sum(flat_input**2, dim=1, keepdim=True)
            + torch.sum(self.codebook.weight**2, dim=1)
            - 2 * torch.matmul(flat_input, self.codebook.weight.t()))

        #distances = self.calculate_distances(z)

        # Distances is shape (batch, height, width, embed_dim)
        # Get the index of the closest codebook vector
        min_indices = torch.argmin(distances, dim=1)

        # We now have matrix of one-hot vectors, with "1" at the index of the closest codebook vector, given by "min_indices"
        one_hot = nn.functional.one_hot(min_indices, num_classes=self.codebook_size)

        # Now to get the value of the closest codebook vector, we multiply the one-hot matrix with the codebook matrix
        z_q = torch.matmul(one_hot.float(), self.codebook.weight)

        # Reshape back
        z_q = rearrange(z_q, '(b h w) c -> b h w c', b=z.shape[0], h=self._encoding_h, w=self._encoding_w)

        return z_q

    def forward(self, x):
        assert x.shape[1:] == (self._input_c, self._input_h, self._input_w)

        # Encode
        z = self.encode(x)

        # Quantize the latent vector
        z_q = self.quantize(z)

        # Add streight through estimator
        z_q = z + (z_q - z).detach()
        
        # Decode
        x_reconstructed = self.decode(z_q)

        # The output image should have the same shape as the input image
        assert x_reconstructed.shape == x.shape

        # Return x hat (and also some other stuff for loss calculation and debugging)
        return x_reconstructed, z, z_q

    def training_step(self, batch, batch_idx):
        images, _ = batch
    
        # Forward pass
        x_hat, z_e, z_q = self(images)

        # Calculate loss
        # Reconstruction loss
        recon_loss = nn.BCELoss(reduction='sum')(x_hat, images)
        # recon_loss = nn.functional.mse_loss(x_hat, images)

        # Quantization loss
        quant_loss = nn.functional.mse_loss(z_e.detach(), z_q) # TODO: This loss is increasing
        # quant_loss = torch.mean(torch.norm((emb - z_e.detach())**2, 2, 1))

        # Commitment loss
        commit_loss = nn.functional.mse_loss(z_q.detach(), z_e)

        # Total loss
        loss = recon_loss + quant_loss + self.beta * commit_loss

        values = {"loss": loss, "recon_loss": recon_loss, "quant_loss": quant_loss, "commit_loss": commit_loss}
        self.log_dict(values)

        return loss
    
    def configure_optimizers(self):
        lr = 1e-3
        optimizer = optim.Adam(self.parameters(), lr=lr)
        return optimizer

# Model
input_dim = (3, 32, 32)
codebook_size = 512

embedding_dim = 64 # Dimension of each codebook vector
encoding_dim = (embedding_dim, 8, 8)

hidden_dim = 1024

beta = 0.25