import os
import torch
from tqdm import tqdm, trange
import torch.nn as nn
import logging
import torch.nn.functional as F
from tabulate import tabulate 
from torch.utils.data import DataLoader
from utils.cluster_frames import VideoClusterer
from utils.metrics import compute_metrics

class SalientFrameSamplerTrainer:
    def __init__(self, model, train_dataloader: DataLoader, val_dataloader: DataLoader, cfg):
        self.cfg = cfg
        self.model = model
        self.train_dataloader = train_dataloader
        self.val_dataloader = val_dataloader
        self.init_logging()  
        self.criterion1 = kl_divergence_loss()
        self.criterion2 = nn.BCEWithLogitsLoss()
        self.criterion3 = CrossEn()
        self.optimizer = self.init_optimizer()

        if self.cfg.use_scheduler:
            self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=self.cfg.scheduler_step_size, gamma=self.cfg.scheduler_gamma)
            
    def train(self):

        for epoch in trange(self.cfg.epochs, desc="Epoch", ncols=100): 
                            
            log_data = {}

            if self.cfg.do_train:
                train_loss = self.run_training_epoch()
                log_data.update({
                    "train_loss(f2f)": train_loss['f2f'],
                    "train_loss(f2t)": train_loss['f2t'],
                    "train_loss(v2t)": train_loss['v2t'],
                })

            if self.cfg.do_eval:
                valid_loss, t2v_metrics, v2t_metrics = self.run_validation_epoch()
                log_data.update({
                    "valid_loss(f2f)": valid_loss['f2f'],
                    "valid_loss(f2t)": valid_loss['f2t'],
                    "valid_loss(v2t)": valid_loss['v2t'],
                })

                for key in t2v_metrics.keys():
                    for r in [1, 5, 10]:
                        log_data[f"[t2v] R@{r} (k={key})"] = t2v_metrics[key][f'R@{r}']
                        log_data[f"[v2t] R@{r} (k={key})"] = v2t_metrics[key][f'R@{r}']
                        
            self.log_metrics(epoch, log_data)

            self.model.anneal_tau()

            if hasattr(self,'scheduler'):
                self.scheduler.step()
            
            if self.cfg.store_backbone:
                torch.save(self.model.pretrained.state_dict(), os.path.join(self.cfg.trained_backbone_path, f"epoch_{epoch+1}.pt"))

            if self.cfg.store_complete_model:
                torch.save(self.model.state_dict(), os.path.join(self.args.trained_model_path, f"epoch_{epoch+1}.pt"))


    def run_training_epoch(self):
        self.model.train()

        epoch_train_loss = {'f2f': 0.0, 'f2t': 0.0, 'v2t': 0.0}

        video_embeddings = []
        text_embeddings = []
        f2f_loss_accum, f2t_loss_accum = 0.0, 0.0

        for i, (frames, teacher_text_features, teacher_frame_features, gt_score) in enumerate(tqdm(self.train_dataloader, desc='Train', leave=False, ncols=100)):
            torch.cuda.empty_cache()
            frames = frames[0].to(self.cfg.device)
            teacher_frame_features = teacher_frame_features[0].to(self.cfg.device)
            teacher_text_features = teacher_text_features[0].to(self.cfg.device)
            gt_score = gt_score[0].to(self.cfg.device)

            frame_scores, frame_weights, student_frame_features = self.model(frames)

            # Frame-to-Frame loss
            f2f_loss = self.criterion1(student_frame_features, teacher_frame_features)

            # Frame-to-Text loss
            f2t_loss = self.criterion2(frame_scores, gt_score)

            # Update accumulators
            f2f_loss_accum += f2f_loss.item() / self.cfg.batch_size
            f2t_loss_accum += f2t_loss.item() / self.cfg.batch_size

            # Calculating the in-pipeline loss for v2t
            video_features = torch.matmul(teacher_frame_features.T, frame_weights) / torch.sum(frame_weights)
            video_embeddings.append(video_features)
            text_embeddings.append(teacher_text_features)

            if (i + 1) % self.cfg.batch_size == 0 or i == len(self.train_dataloader) - 1:
                sim_matrix = self.compute_sim_matrix(video_embeddings, text_embeddings)
                sim_loss1 = self.criterion3(sim_matrix)
                sim_loss2 = self.criterion3(sim_matrix.T)
                v2t_loss = (sim_loss1 + sim_loss2) / 2
                
                # Aggregate the total loss with regularization terms
                total_loss = self.alpha * f2f_loss_accum + self.beta * f2t_loss_accum + self.gamma * v2t_loss
   
                self.optimizer.zero_grad()
                total_loss.backward()
                
                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
                self.optimizer.step()

                # Update epoch losses
                epoch_train_loss['f2f'] += f2f_loss_accum
                epoch_train_loss['f2t'] += f2t_loss_accum
                epoch_train_loss['v2t'] += v2t_loss.item()

                # Reset accumulators
                f2f_loss_accum, f2t_loss_accum = 0.0, 0.0
                video_embeddings = []
                text_embeddings = []

        # Normalize epoch losses by the number of batches
        for key in epoch_train_loss.keys():
            epoch_train_loss[key] /= len(self.train_dataloader)
            
        return epoch_train_loss

    def run_validation_epoch(self):
        self.model.eval()
        
        epoch_valid_loss = {'f2f': 0.0, 'f2t': 0.0, 'v2t': 0.0}

        video_embeddings = []
        text_embeddings = []
        all_video_embeddings = {str(k): {'without_clustering': [], 
                                         'uniform': [], 'spectral': [], 'agglomerative': [], 'kmeans': []} for k in [1,2,4,8,16]}
        all_text_embeddings = []
        all_st_scores = []

        f2f_loss_accum, f2t_loss_accum = 0.0, 0.0

        for i, (frames, teacher_text_features, teacher_frame_features, gt_score) in enumerate(tqdm(self.val_dataloader, desc='Valid', leave=False, ncols=100)):
            torch.cuda.empty_cache()
            
            frames = frames[0].to(self.cfg.device)
            teacher_frame_features = teacher_frame_features[0].to(self.cfg.device)
            teacher_text_features = teacher_text_features[0].to(self.cfg.device)
            gt_score = gt_score[0].to(self.cfg.device)

            with torch.no_grad():
                frame_scores, frame_weights, student_frame_features = self.model(frames)

            f2f_loss = self.criterion1(student_frame_features, teacher_frame_features)
            f2t_loss = self.criterion2(frame_scores, gt_score)

            f2f_loss_accum += f2f_loss.item()
            f2t_loss_accum += f2t_loss.item()

            # Calculating the all embeddings, used for calculating t2v_retrieval metrics
            all_text_embeddings.append(teacher_text_features)

            n_frames = teacher_frame_features.size(0)
            sorted_indices = sorted(range(n_frames), key=lambda i: frame_scores[i], reverse=True).copy()

            for k in [1,2,4,8,16]:
                n_selected_frames = min(n_frames, k)
                indices = sorted_indices[:n_selected_frames]
                all_video_embeddings[str(k)]['without_clustering'].append(teacher_frame_features[indices].mean(dim=0))

                if self.cfg.do_cluster and k>2:
                    for clustering_method in ['uniform', 'spectral', 'agglomerative', 'kmeans']:
                        if len(frame_scores) <= k: #mean of all frames
                            all_video_embeddings[str(k)][clustering_method].append(teacher_frame_features.mean(dim=0))
                        else:
                            clusterer = VideoClusterer(clustering_method, k)
                            clusters = clusterer.get_clusters(student_frame_features)
                            selected_indices = self.select_representative_frame_per_cluster(clusters, frame_scores)
                            all_video_embeddings[str(k)][clustering_method].append(teacher_frame_features[selected_indices].mean(dim=0))
            
            # Calculating the in-pipeline loss
            video_features = torch.matmul(teacher_frame_features.T, frame_weights) / torch.sum(frame_weights)
            video_embeddings.append(video_features)
            text_embeddings.append(teacher_text_features)

            if (i + 1) % self.cfg.batch_size == 0 or i == len(self.val_dataloader) - 1:
                sim_matrix = self.compute_sim_matrix(video_embeddings, text_embeddings)
                sim_loss1 = self.criterion3(sim_matrix)
                sim_loss2 = self.criterion3(sim_matrix.T)
                v2t_loss = (sim_loss1 + sim_loss2) / 2
                
                v2t_loss = (sim_loss1 + sim_loss2) / 2

                # Update epoch losses
                epoch_valid_loss['f2f'] += f2f_loss_accum
                epoch_valid_loss['f2t'] += f2t_loss_accum
                epoch_valid_loss['v2t'] += v2t_loss.item()
            
                # Reset accumulators
                f2f_loss_accum, f2t_loss_accum = 0.0, 0.0
                video_embeddings = []
                text_embeddings = []

        # Normalize epoch losses by the number of batches
        for key in epoch_valid_loss.keys():
            epoch_valid_loss[key] /= len(self.val_dataloader)

        # Compute t2v metrics
        t2v_metrics = {}
        v2t_metrics = {}

        for k, embeddings_dict in all_video_embeddings.items():
            t2v_metrics[str(k)]={}
            for scenario, embeddings_list in embeddings_dict.items():
                if embeddings_list:  # Check if there are any embeddings for this scenario
                    sim_matrix = self.compute_sim_matrix(embeddings_list, all_text_embeddings)
                    _key = f"{k}" if scenario=='without_clustering' else f"{k}-{scenario}"
                    t2v_metrics[_key] = compute_metrics(sim_matrix.detach().cpu().numpy())
                    v2t_metrics[_key] = compute_metrics(sim_matrix.T.detach().cpu().numpy())

        return epoch_valid_loss, t2v_metrics, v2t_metrics

    def compute_sim_matrix(self, video_embeddings, text_embeddings):
        video_embeddings_tensor = torch.stack(video_embeddings, dim=0)
        text_embeddings_tensor = torch.stack(text_embeddings, dim=0)
        video_embeddings_norm = video_embeddings_tensor / video_embeddings_tensor.norm(dim=-1, keepdim=True)
        text_embeddings_norm = text_embeddings_tensor / text_embeddings_tensor.norm(dim=-1, keepdim=True)
        sim_matrix = torch.matmul(video_embeddings_norm, text_embeddings_norm.T)
        return sim_matrix

    def init_optimizer(self):
        if not hasattr(self.cfg, 'backbone_lr_scale'):
            return torch.optim.Adam(params=self.model.parameters(), lr=self.cfg.lr, weight_decay=self.cfg.weight_decay)

        backbone_params = []
        last_params = []
        for name, param in self.model.named_parameters():
            if 'pretrain' in name:
                backbone_params.append(param)
            else:
                last_params.append(param)

        return torch.optim.Adam([
            {'params': backbone_params, 'lr': self.cfg.lr},
            {'params': last_params, 'lr': self.cfg.lr * self.cfg.backbone_lr_scale}
        ], weight_decay=self.cfg.weight_decay)
   
    def init_logging(self):
        
        available_logs = os.listdir(self.cfg.log_dir)
        last_run = 0
        if available_logs:
            last_run = max([int(name.split('-')[0]) for name in available_logs])
        log_file_name = f"{last_run + 1}-local-run.log"

        log_file = os.path.join(self.cfg.log_dir, log_file_name)

        logging.basicConfig(filename=log_file, level=logging.INFO,
                            format='%(asctime)s [%(levelname)s] - %(message)s',
                            datefmt='%Y-%m-%d %H:%M:%S')

        hyperparameters = {key: value for key, value in vars(self.cfg).items()}
        hyperparameters_str = "\n".join([f"{key}: {value}" for key, value in hyperparameters.items() if ('path' not in key and 'log_dir' not in key)])
        logging.info("Hyperparameters:\n" + hyperparameters_str)
   
    def log_metrics(self, epoch, log_data):
        # Loss Metrics Table
        loss_header = ["", "f2f", "f2t", "v2t"]

        # Check if train_loss exists, otherwise assign "N/A"
        train_loss_f2f = log_data.get("train_loss(f2f)", "N/A")
        train_loss_f2t = log_data.get("train_loss(f2t)", "N/A")
        train_loss_v2t = log_data.get("train_loss(v2t)", "N/A")

        # Check if valid_loss exists, otherwise assign "N/A"
        valid_loss_f2f = log_data.get("valid_loss(f2f)", "N/A")
        valid_loss_f2t = log_data.get("valid_loss(f2t)", "N/A")
        valid_loss_v2t = log_data.get("valid_loss(v2t)", "N/A")

        loss_rows = [
            ["train", train_loss_f2f, train_loss_f2t, train_loss_v2t],
            ["val", valid_loss_f2f, valid_loss_f2t, valid_loss_v2t],
        ]
        loss_table = tabulate(loss_rows, headers=loss_header, tablefmt="grid")

        # Recall Metrics Table
        recall_header = ["", "R@1", "R@5", "R@10"]

        clustering_names = ['uniform', 'spectral', 'agglomerative', 'kmeans']

        recall_rows_t2v = []
        recall_rows_v2t = []
        for k in [1, 2, 4, 8, 16]:
            # Without clustering
            base_key = f"k={k}"
            if any(f"[t2v] R@{r} ({base_key})" in log_data for r in [1, 5, 10]):
                row = [f"{k}"]
                for r in [1, 5, 10]:
                    metric_key = f"[t2v] R@{r} ({base_key})"
                    row.append(log_data.get(metric_key, "N/A"))
                recall_rows_t2v.append(row)

            if any(f"[v2t] R@{r} ({base_key})" in log_data for r in [1, 5, 10]):
                row = [f"{k}"]
                for r in [1, 5, 10]:
                    metric_key = f"[v2t] R@{r} ({base_key})"
                    row.append(log_data.get(metric_key, "N/A"))
                recall_rows_v2t.append(row)

            # With clustering
            for clustering_name in clustering_names:
                cluster_key = f"{base_key}-{clustering_name}"
                if any(f"[t2v] R@{r} ({cluster_key})" in log_data for r in [1, 5, 10]):
                    row = [f"{k}-{clustering_name}"]
                    for r in [1, 5, 10]:
                        metric_key = f"[t2v] R@{r} ({cluster_key})"
                        row.append(log_data.get(metric_key, "N/A"))
                    recall_rows_t2v.append(row)
                
                if any(f"[v2t] R@{r} ({cluster_key})" in log_data for r in [1, 5, 10]):
                    row = [f"{k}-{clustering_name}"]
                    for r in [1, 5, 10]:
                        metric_key = f"[v2t] R@{r} ({cluster_key})"
                        row.append(log_data.get(metric_key, "N/A"))
                    recall_rows_v2t.append(row)

        recall_table_t2v = tabulate(recall_rows_t2v, headers=recall_header, tablefmt="grid")
        recall_table_v2t = tabulate(recall_rows_v2t, headers=recall_header, tablefmt="grid")
        
        logging.info(f"Epoch {epoch+1} - Loss and Recall Metrics:")
        logging.info("\n" + loss_table)
        logging.info("\nText-to-Video Recall Metrics:")
        logging.info("\n" + recall_table_t2v)
        logging.info("\nVideo-to-Text Recall Metrics:")
        logging.info("\n" + recall_table_v2t)

    def select_representative_frame_per_cluster(self, clusters, frame_scores):
        representative_frames = []

        for cluster in clusters:
            best_frame_index = max(cluster, key=lambda index: frame_scores[index])
            representative_frames.append(int(best_frame_index))

        return representative_frames

class kl_divergence_loss(nn.Module):
    def __init__(self,):
        super(kl_divergence_loss, self).__init__()

    def forward(self, features1, features2):
        features1 = F.normalize(features1, p=2, dim=-1)
        features2 = F.normalize(features2, p=2, dim=-1)

        cos_sim_features1 = torch.mm(features1, features1.t())
        cos_sim_features2 = torch.mm(features2, features2.t())

        probs_features1 = F.softmax(cos_sim_features1, dim=-1)
        probs_features2 = F.softmax(cos_sim_features2, dim=-1)

        loss = F.kl_div(probs_features1.log(), probs_features2, reduction='batchmean')

        return loss
    
class CrossEn(nn.Module):
    def __init__(self,):
        super(CrossEn, self).__init__()

    def forward(self, sim_matrix):
        logpt = F.log_softmax(sim_matrix, dim=-1)
        logpt = torch.diag(logpt)
        nce_loss = -logpt
        sim_loss = nce_loss.mean()
        return sim_loss