import wandb
import torch
import numpy as np
from monai.visualize import blend_images


class WandBModel:
    """
    Enable WandB features to the model using multiple inheritance
    """
    def __init__(self, *args, **kwargs):
        # the following attributes should be initialized by class `BaseSegmentationModel`
        self.visual_pairs = None
        self.train_loss = None
        self.val_loss = None
        self.metric_meter = None
        self.name = None
        # the following attributes should be initialized by the child class
        self.val_table = None

    def volume2videos(self, time_dim=3, tag=''):
        """
        Convert 3D volumes to video in favor of WandB logging
        Args:
            time_dim: the spatial dimension to be converted as the time dimension, default is the axial axis (dim 3)
            tag: extra information for logging
        """
        videos = []
        for image_pair in self.visual_pairs:
            try:
                pair_name = getattr(self, image_pair['name'])
                image = getattr(self, image_pair['image'])
                mask = getattr(self, image_pair['mask'])
                vis_type = image_pair['type']
            except:
                continue
            for i in range(image.shape[0]):  # deallocate the batch dim
                image2save = image[i, ...]
                mask2save = mask[i, ...]
                item_name = pair_name[i]
                # detach the tensor, format [C, H, W, D]
                image_numpy = image2save.detach()
                mask_numpy = mask2save.detach()
                if mask_numpy.shape[0] > 1:
                    mask_numpy = torch.argmax(mask_numpy, dim=0, keepdim=True)
                # (C, H, W, D), torch.Tensor on device
                pair_blend = blend_images(image_numpy, mask_numpy, alpha=0.5) * 255
                # permute the axes to (time, channel, height, width)
                spatial_dim = list(range(1, len(pair_blend.shape[1:]) + 1))
                spatial_dim.remove(time_dim)
                pair_blend = pair_blend.permute([time_dim, 0] + spatial_dim).cpu().numpy().astype(np.uint8)
                # record in the wandb.Video class
                video = wandb.Video(pair_blend, fps=8, caption='{}_{}{}'.format(item_name, vis_type, tag))
                videos.append(video)
        return videos

    def log_scaler(self, key, value, step=None):
        """
        Log manually defined scaler data
        """
        wandb.log({key: np.round(value, decimals=4)}, step=step)

    def log_train_loss(self, step=None):
        """
        Log train loss
        """
        data_dict = self.train_loss.pop_data(True)
        for key, value in data_dict.items():
            wandb.log({'train/{}'.format(key): value}, step=step)

    def log_val_loss(self, step=None):
        """
        Log val loss
        """
        data_dict = self.val_loss.pop_data(True)
        for key, value in data_dict.items():
            wandb.log({'val/{}'.format(key): value}, step=step)

    def log_metrics(self, step=None):
        """
        Log validation metrics as wandb.Table
        """
        df = self.metric_meter.to_df()
        wandb.log({'val/metrics': wandb.Table(dataframe=df)}, step=step)

    def log_vis(self, key, step=None, time_dim=3, tag=''):
        """
        Log training intermediate visualizations
        """
        videos = self.volume2videos(time_dim, tag)
        wandb.log({key: videos}, step=step)

    def update_val_visualization(self, time_dim=3, tag=''):
        """
        Update the validation visualization to buffer, called every step of evaluation
        """
        videos = self.volume2videos(time_dim, tag)
        self.val_table.add_data(self.name, *videos)

    def log_val_visualization(self, step=None):
        """
        Log validation visualization
        """
        wandb.log({'val/visualization': self.val_table}, step=step)
        # re-initialize the table for next logging
        del self.val_table
        self.val_table = wandb.Table(columns=['ID'] + [pair['type'] for pair in self.visual_pairs])