import torch
from torchvision import transforms
from torch.utils.data import random_split
from torch.utils.data import DataLoader, Dataset

import os
import random
import pickle
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from train_mm_moco import MultiModalMoCo
from sklearn.neighbors import NearestNeighbors


def denormalize(tensor, mean, std):
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

def compute_tsne(model, test_dataloader):
    with torch.no_grad():
        test_data_list = list(test_dataloader)
        x_vision_test, x_tactile_test = random.choice(test_data_list)
        random_indices = random.sample(range(x_vision_test.shape[0]), 100)
        x_vision_test = x_vision_test[random_indices].to('cuda')
        x_tactile_test = x_tactile_test[random_indices].to('cuda')
        vision_base_q = model.vision_base_q(x_vision_test)
        tactile_base_q = model.tactile_base_q(x_tactile_test)
    
    x_vision_test = denormalize(x_vision_test, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    x_tactile_test = denormalize(x_tactile_test, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

    vision_base_q = vision_base_q.cpu().numpy()
    tactile_base_q = tactile_base_q.cpu().numpy()

    image_data = np.concatenate((x_vision_test.cpu().numpy(), x_tactile_test.cpu().numpy()), axis=0)

    tsne = TSNE(n_components=2, random_state=0, perplexity=75,n_iter=50000)

    # Create pairs of corresponding representations and labels
    num_samples = min(vision_base_q.shape[0], tactile_base_q.shape[0])
    data = np.concatenate((vision_base_q[:num_samples], tactile_base_q[:num_samples]), axis=0)
    labels = np.arange(1, 2*(num_samples)+1)

    tsne_data = tsne.fit_transform(data)
    nn_all = find_knn(tsne_data, labels)
    plot_images_by_labels(image_data, nn_all[0])
    print(nn_all[:5])

    fig = plt.figure(figsize=(10, 10))
    
    for i, (x, y) in enumerate(tsne_data):
        plt.scatter(x, y, color='blue')
        plt.text(x, y, f"{labels[i]}", fontsize=12, ha='center', va='bottom')
    plt.savefig('temp_figure.png')
    plt.close(fig)

    image = Image.open('temp_figure.png')
    image = np.array(image)  # Convert image to a NumPy array
    image_rgb = image[:, :, :3] # Extract RGB channels and change format to CHW
    plt.imshow(image_rgb)
    plt.title('t-SNE plot')
    plt.axis('off')
    plt.show()

def find_knn(tsne_data, labels):
    neigh = NearestNeighbors(n_neighbors=8)
    neigh.fit(tsne_data)
    knn = neigh.kneighbors(tsne_data, return_distance=False)
    return labels[knn]

def plot_images_by_labels(image_data, labels_to_plot):
    fig, axes = plt.subplots(1, len(labels_to_plot), figsize=(15, 5))
    for i, label in enumerate(labels_to_plot):
        axes[i].imshow(image_data[label].transpose(1, 2, 0))
        axes[i].set_title(label)
        axes[i].axis('off')
    plt.show()

if __name__ == "__main__":
    class CustomMultiModalDataset(Dataset):
        def __init__(self, vision_folder, tactile_folder, transform=None):
            self.vision_folder = vision_folder
            self.tactile_folder = tactile_folder
            self.transform = transform

            self.vision_files = sorted(os.listdir(vision_folder))
            self.tactile_files = sorted(os.listdir(tactile_folder))

        def __len__(self):
            return len(self.vision_files)

        def __getitem__(self, idx):
            vision_path = os.path.join(self.vision_folder, self.vision_files[idx])
            tactile_path = os.path.join(self.tactile_folder, self.tactile_files[idx])

            vision_image = Image.open(vision_path).convert("RGB")
            tactile_image = Image.open(tactile_path).convert("RGB")

            if self.transform:
                vision_image = self.transform(vision_image)
                tactile_image = self.transform(tactile_image)

            return vision_image, tactile_image

    # Initialize augmentation
    simple_transforms = transforms.Compose([
        transforms.CenterCrop(500),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # Load the indices from disk
    with open('indices/train_indices.pkl', 'rb') as f:
        train_indices = pickle.load(f)
        
    with open('indices/test_indices.pkl', 'rb') as f:
        test_indices = pickle.load(f)

    # Initialize dataset and dataloader
    vision_folder = "/home/vedant/Downloads/ssvtp_data/images_rgb"
    tactile_folder = "/home/vedant/Downloads/ssvtp_data/images_tac"
    dataset = CustomMultiModalDataset(vision_folder, tactile_folder, transform=simple_transforms)

    # Create subset datasets and DataLoaders
    test_subset = torch.utils.data.Subset(dataset, test_indices)
    test_dataloader = DataLoader(test_subset, batch_size=150, shuffle=False)

    from torch.utils.tensorboard import SummaryWriter
    writer = SummaryWriter('runs/mmssl1')
    model = MultiModalMoCo(writer).to('cuda')
    model.load_state_dict(torch.load('/home/vedant/TacSSL/models/model.pth'))
    compute_tsne(model, test_dataloader)