import time
import torch
import torch.nn as nn
from metrics.abstract_metrics import CrossEntropyMetric
from torchmetrics import Metric, MeanSquaredError
import wandb
import torch.nn.functional as F

# from 2:He to 119:*
valencies_check = [0, 1, 2, 3, 4, 3, 2, 1, 0, 1, 2, 6, 6, 7, 6, 1, 0, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 4, 7, 6, 1, 0, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 4, 7, 6, 5, 6, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 4, 7, 6, 5, 0, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
valencies_check = torch.tensor(valencies_check)

bond_check = torch.tensor([0, 1, 2, 3, 1])
       
weight_check = [4.003, 6.941, 9.012, 10.812, 12.011, 14.007, 15.999, 18.998, 20.18, 22.99, 24.305, 26.982, 28.086, 30.974, 32.067, 35.453, 39.948, 39.098, 40.078, 44.956, 47.867, 50.942, 51.996, 54.938, 55.845, 58.933, 58.693, 63.546, 65.39, 69.723, 72.61, 74.922, 78.96, 79.904, 83.8, 85.468, 87.62, 88.906, 91.224, 92.906, 95.94, 98.0, 101.07, 102.906, 106.42, 107.868, 112.412, 114.818, 118.711, 121.76, 127.6, 126.904, 131.29, 132.905, 137.328, 138.906, 140.116, 140.908, 144.24, 145.0, 150.36, 151.964, 157.25, 158.925, 162.5, 164.93, 167.26, 168.934, 173.04, 174.967, 178.49, 180.948, 183.84, 186.207, 190.23, 192.217, 195.078, 196.967, 200.59, 204.383, 207.2, 208.98, 209.0, 210.0, 222.0, 223.0, 226.0, 227.0, 232.038, 231.036, 238.029, 237.0, 244.0, 243.0, 247.0, 247.0, 251.0, 252.0, 257.0, 258.0, 259.0, 262.0, 267.0, 268.0, 269.0, 270.0, 269.0, 278.0, 281.0, 281.0, 285.0, 284.0, 289.0, 288.0, 293.0, 292.0, 294.0, 294.0]
weight_check = torch.tensor(weight_check)

class AtomWeightMetric(Metric):
    def __init__(self):
        super().__init__()
        self.add_state('total_loss', default=torch.tensor(0.), dist_reduce_fx="sum")
        self.add_state('total_samples', default=torch.tensor(0.), dist_reduce_fx="sum")
        global weight_check
        self.weight_check = weight_check

    def update(self, X, Y):
        atom_pred_num = X.argmax(dim=-1)
        atom_real_num = Y.argmax(dim=-1)
        self.weight_check = self.weight_check.type_as(X)

        pred_weight = self.weight_check[atom_pred_num]
        real_weight = self.weight_check[atom_real_num]

        lss = 0
        lss += torch.abs(pred_weight.sum(dim=-1) - real_weight.sum(dim=-1)).sum()
        self.total_loss += lss
        self.total_samples += X.size(0)

    def compute(self):
        return self.total_loss / self.total_samples
            

class TrainLossDiscrete(nn.Module):
    """ Train with Cross entropy"""
    def __init__(self, lambda_train, weight_node=None, weight_edge=None):
        super().__init__()
        self.node_loss = CrossEntropyMetric()
        self.edge_loss = CrossEntropyMetric()
        self.weight_loss = AtomWeightMetric() # total molecular weight of the pred mol == that of the true mol

        self.y_loss = MeanSquaredError()
        self.lambda_train = lambda_train
        
    def forward(self, masked_pred_X, masked_pred_E, pred_y, true_X, true_E, true_y, node_mask, log: bool):
        """ Compute train metrics
        masked_pred_X : tensor -- (bs, n, dx)
        masked_pred_E : tensor -- (bs, n, n, de)
        pred_y : tensor -- (bs, )
        true_X : tensor -- (bs, n, dx)
        true_E : tensor -- (bs, n, n, de)
        true_y : tensor -- (bs, )
        log : boolean. """

        loss_weight = self.weight_loss(masked_pred_X, true_X)
        
        true_X = torch.reshape(true_X, (-1, true_X.size(-1)))  # (bs * n, dx)
        true_E = torch.reshape(true_E, (-1, true_E.size(-1)))  # (bs * n * n, de)
        masked_pred_X = torch.reshape(masked_pred_X, (-1, masked_pred_X.size(-1)))  # (bs * n, dx)
        masked_pred_E = torch.reshape(masked_pred_E, (-1, masked_pred_E.size(-1)))   # (bs * n * n, de)

        # Remove masked rows
        mask_X = (true_X != 0.).any(dim=-1)
        mask_E = (true_E != 0.).any(dim=-1)

        flat_true_X = true_X[mask_X, :]
        flat_pred_X = masked_pred_X[mask_X, :]

        flat_true_E = true_E[mask_E, :]
        flat_pred_E = masked_pred_E[mask_E, :]
        
        pred_X_0 = flat_pred_X.argmax(dim=-1) == 0
        true_X_0 = flat_true_X.argmax(dim=-1) == 0
        
        pred_E_0 = flat_pred_E.argmax(dim=-1) == 0
        true_E_0 = flat_true_E.argmax(dim=-1) == 0
        
        # accuracy of X
        acc_X = (pred_X_0 == true_X_0).sum() / len(pred_X_0)
        
        # accuracy of E
        acc_E = (pred_E_0 == true_E_0).sum() / len(pred_E_0)

        loss_X = self.node_loss(flat_pred_X, flat_true_X) if true_X.numel() > 0 else 0.0
        loss_E = self.edge_loss(flat_pred_E, flat_true_E) if true_E.numel() > 0 else 0.0

        return self.lambda_train[0] * loss_X + self.lambda_train[1] * loss_E + loss_weight

    def forward_w_mask(self, pred_X, pred_E, pred_y, true_X, true_E, true_y, node_mask, cls_mask_node, cls_mask_edge, 
                       dgamma_times_alpha, log: bool, observation: bool=False, noise_schedule: str=None):
        # Reshape tensors
        bs = pred_X.size(0)
        n_max_node = pred_X.size(1)
        true_X = torch.reshape(true_X, (-1, true_X.size(-1)))  # (bs * n, dx)
        true_E = torch.reshape(true_E, (-1, true_E.size(-1)))  # (bs * n * n, de)
        pred_X = torch.reshape(pred_X, (-1, pred_X.size(-1)))  # (bs * n, dx)
        pred_E = torch.reshape(pred_E, (-1, pred_E.size(-1)))   # (bs * n * n, de)
        
        # Reshape masks
        edge_mask = (node_mask.unsqueeze(1) * node_mask.unsqueeze(2))
        diag = torch.eye(n_max_node, dtype=torch.bool, device=node_mask.device)
        edge_mask[diag.unsqueeze(0).expand(bs, -1, -1)] = False
        edge_mask = edge_mask.reshape(-1)
        node_mask = node_mask.view(-1)  # (bs * n)
        cls_mask_node = cls_mask_node.reshape(-1)  # (bs * n)
        cls_mask_edge = cls_mask_edge.reshape(-1)  # (bs * n * n)
        
        final_node_mask = node_mask & cls_mask_node
        final_edge_mask = edge_mask & cls_mask_edge
        # Apply maskss
        flat_true_X = true_X[final_node_mask, :-1]
        flat_pred_X = pred_X[final_node_mask, :-1]

        flat_true_E = true_E[final_edge_mask, :-1]
        flat_pred_E = pred_E[final_edge_mask, :-1]

        if not observation:
            loss_weight = self.weight_loss(flat_pred_X, flat_true_X)

            if "learnable" in noise_schedule:
                weights_node, weights_edge = dgamma_times_alpha
                weights_node = weights_node.reshape(-1)
                weights_edge = weights_edge.reshape(-1)

                weights_node = weights_node[final_node_mask]
                weights_edge = weights_edge[final_edge_mask]
            else:
                if "spindle" in noise_schedule:
                    weights_node, weights_edge = dgamma_times_alpha
                    weights_node = weights_node.reshape(-1)
                    weights_edge = weights_edge.reshape(-1)
                else:
                    weights = dgamma_times_alpha
                    weights_node = weights.expand(-1, n_max_node).reshape(-1)
                    weights_edge = weights.unsqueeze(1).expand(-1, n_max_node, n_max_node).reshape(-1)
                    
                weights_node = weights_node[final_node_mask]
                weights_edge = weights_edge[final_edge_mask]

            if "spindle" in noise_schedule:
                loss_X = self.node_loss(flat_pred_X, flat_true_X) if true_X.numel() > 0 else 0.0
                loss_E = self.edge_loss(flat_pred_E, flat_true_E) if true_E.numel() > 0 else 0.0
            else:
                # Calculate losses
                loss_X = self.node_loss(flat_pred_X, flat_true_X, sample_weight=weights_node) if true_X.numel() > 0 else 0.0
                loss_E = self.edge_loss(flat_pred_E, flat_true_E, sample_weight=weights_edge) if true_E.numel() > 0 else 0.0

        # Calculate accuracies
        true_X = flat_true_X.argmax(dim=-1)
        true_E = flat_true_E.argmax(dim=-1)
        
        if observation:
            if cls_mask_node.sum() == 0:
                loss_X = 0.0
            else:
                loss_X = F.cross_entropy(flat_pred_X, true_X, reduction='none')
                loss_X = loss_X.sum() / len(loss_X)
            if cls_mask_edge.sum() == 0:
                loss_E = 0.0
            else:
                loss_E = F.cross_entropy(flat_pred_E, true_E, reduction='none')
                loss_E = loss_E.sum() / len(loss_E)
            return self.lambda_train[0] * loss_X + self.lambda_train[1] * loss_E 
        else:
            return self.lambda_train[0] * loss_X + self.lambda_train[1] * loss_E + loss_weight

    def reset(self):
        for metric in [self.node_loss, self.edge_loss, self.y_loss]:
            metric.reset()

    def log_epoch_metrics(self, current_epoch, start_epoch_time, log=True):
        epoch_node_loss = self.node_loss.compute() if self.node_loss.total_samples > 0 else -1
        epoch_edge_loss = self.edge_loss.compute() if self.edge_loss.total_samples > 0 else -1
        epoch_weight_loss = self.weight_loss.compute() if self.weight_loss.total_samples > 0 else -1

        if log:
            print(f"Epoch {current_epoch} finished: X_CE: {epoch_node_loss :.4f} -- E_CE: {epoch_edge_loss :.4f} "
                f"Weight: {epoch_weight_loss :.4f} "
                f"-- Time taken {time.time() - start_epoch_time:.1f}s ")

