from model import MVAE

import torch
import os
import sys
import shutil
from torchvision import transforms

def elbo_loss(recon_image, image, recon_text, text, mu, logvar,
              lambda_mnist=1.0, lambda_svhn=1.0, annealing_factor=1):
    """Bimodal ELBO loss function.

    @param recon_image: torch.Tensor
                        reconstructed image
    @param image: torch.Tensor
                  input image
    @param recon_text: torch.Tensor
                       reconstructed text probabilities
    @param text: torch.Tensor
                 input text (one-hot)
    @param mu: torch.Tensor
               mean of latent distribution
    @param logvar: torch.Tensor
                   log-variance of latent distribution
    @param lambda_mnist: float [default: 1.0]
                         weight for mnist BCE
    @param lambda_svhn: float [default: 1.0]
                       weight for svhn BCE
    @param annealing_factor: integer [default: 1]
                             multiplier for KL divergence term
    @return ELBO: torch.Tensor
                  evidence lower bound
    """
    image_bce, text_bce = 0, 0  # default params
    if recon_image is not None and image is not None:
        image_bce = torch.sum(binary_cross_entropy_with_logits(
            recon_image.view(-1, 1 * 28 * 28),
            image.view(-1, 1 * 28 * 28)), dim=1)

    if recon_text is not None and text is not None:
        text_bce = torch.sum(binary_cross_entropy_with_logits(
            recon_text.view(-1, 3 * 32 * 32),
            text.view(-1, 3 * 32 * 32)), dim=1)

    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
    ELBO = torch.mean(lambda_mnist * image_bce + lambda_svhn * text_bce
                      + annealing_factor * KLD)
    return ELBO


def binary_cross_entropy_with_logits(input, target):
    """Sigmoid Activation + Binary Cross Entropy

    @param input: torch.Tensor (size N)
    @param target: torch.Tensor (size N)
    @return loss: torch.Tensor (size N)
    """
    if not (target.size() == input.size()):
        raise ValueError("Target size ({}) must be the same as input size ({})".format(
            target.size(), input.size()))

    return (torch.clamp(input, 0) - input * target
            + torch.log(1 + torch.exp(-torch.abs(input))))


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def is_multidata(dataB):
    return isinstance(dataB, list) or isinstance(dataB, tuple)


def unpack_data(dataB, device='cuda'):
    # dataB :: (Tensor, Idx) | [(Tensor, Idx)]
    """ Unpacks the data batch object in an appropriate manner to extract data """
    if is_multidata(dataB):
        if torch.is_tensor(dataB[0]):
            return dataB[0].to(device)
        elif isinstance(dataB[0], tuple) or isinstance(dataB[0], list):
            return [d.to(device) for d in list(zip(*dataB))[0]]
        else:
            raise RuntimeError('Unknown type for {}'.format(dataB[0]))
    else:
        raise RuntimeError('Unknown type for {}'.format(dataB))


def save_checkpoint(state, is_best, folder='./', filename='checkpoint.pth.tar'):
    if not os.path.isdir(folder):
        os.mkdir(folder)
    torch.save(state, os.path.join(folder, filename))
    if is_best:
        shutil.copyfile(os.path.join(folder, filename),
                        os.path.join(folder, 'model_best.pth.tar'))


def load_checkpoint(file_path, use_cuda=False):
    checkpoint = torch.load(file_path) if use_cuda else \
        torch.load(file_path, map_location=lambda storage, location: storage)
    model = MVAE(checkpoint['n_latents'])
    model.load_state_dict(checkpoint['state_dict'])
    return model

def resize_img(img, ref):
    tx = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize(ref.size(-1)),  # use last dim assuming square image
        transforms.ToTensor()
    ])
    return torch.stack([tx(_img) for _img in img.cpu()]).to(img.device).expand_as(ref)

# https://stackoverflow.com/questions/14906764/how-to-redirect-stdout-to-both-file-and-console-with-scripting
class Logger(object):
    def __init__(self, filename):
        self.terminal = sys.stdout
        self.log = open(filename, "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        # this flush method is needed for python 3 compatibility.
        # this handles the flush command by doing nothing.
        # you might want to specify some extra behavior here.
        pass

def unpack_data_mlp(dataB, option='both'):
    if len(dataB[0])==2:
        if option == 'both':
            return dataB[0][0], dataB[1][0], dataB[1][1]
        elif option == 'svhn':
            return dataB[1][0], dataB[1][1]
        elif option == 'mnist':
            return dataB[0][0], dataB[0][1]
    else:
        return dataB
