"""Optional Weights & Biases logging for Ref-AVS training."""
import os

import torchvision
import wandb


class Tensorboard:
    def __init__(self, config):
        key = config.get('wandb_key') or os.environ.get('WANDB_API_KEY', '')
        if key:
            os.environ['WANDB_API_KEY'] = key
        mode = 'online' if config.get('wandb_online', False) else 'disabled'
        self.tensor_board = wandb.init(
            project=config['proj_name'],
            name=config['experiment_name'],
            config=config,
            mode=mode,
            settings=wandb.Settings(code_dir='.'),
        )
        self.restore_transform = torchvision.transforms.Compose([
            DeNormalize(config['image_mean'], config['image_std']),
            torchvision.transforms.ToPILImage(),
        ])

    def upload_wandb_info(self, info_dict):
        for key, value in info_dict.items():
            self.tensor_board.log({key: value})

    def upload_wandb_image(self, frames, pseudo_label_from_pred, pseudo_label_from_sam, img_number=4):
        n = min(pseudo_label_from_pred.shape[0], img_number)
        frames = frames[:n]
        pseudo_label_from_sam = pseudo_label_from_sam[:n].float()
        pseudo_label_from_pred = pseudo_label_from_pred[:n].float()
        pseudo_label_from_sam[pseudo_label_from_sam == 255.] = 0.5
        pseudo_label_from_pred[pseudo_label_from_pred == 255.] = 0.5
        self.tensor_board.log({
            'image': [wandb.Image(j, caption=f'id {i}') for i, j in enumerate(frames)],
            'label': [wandb.Image(j.squeeze(), caption=f'id {i}') for i, j in enumerate(pseudo_label_from_sam)],
            'logits': [wandb.Image(j.squeeze(), caption=f'id {i}') for i, j in enumerate(pseudo_label_from_pred)],
        })

    def finish(self):
        self.tensor_board.finish()


class DeNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

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