"""The core implementation of Inception Score and FID."""

from typing import List, Union, Tuple

import numpy as np
import torch
from scipy import linalg
from tqdm.auto import tqdm
from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
from torch.utils.data.distributed import DistributedSampler

from . import districuted as dist
from .inception import InceptionV3


# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


def get_inception_feature(
    images: Union[torch.FloatTensor, DataLoader],
    dims: List[int],
    batch_size: int = 50,
    use_torch: bool = False,
    verbose: bool = False,
    device: torch.device = None,
) -> Union[torch.FloatTensor, np.ndarray]:
    """Calculate Inception Score and FID.

    For each image, only a forward propagation is required to calculating
    features for FID and Inception Score.

    Args:
        images: tensor or torch.utils.data.Dataloader. The images
                must be float tensor of range [0, 1].
        dims: List of int, see InceptionV3.BLOCK_INDEX_BY_DIM for
              available dimension.
        batch_size: int, The batch size for calculating activations. If
                    `images` is torch.utils.data.Dataloader, this argument is
                    ignored.
        use_torch: When True, use torch to calculate FID. Otherwise, use numpy.
        verbose: Set verbose to False for disabling progress bar. Otherwise,
                 the progress bar is showing when calculating activations.
        device: the torch device which is used to calculate inception feature
    Returns:
        inception_features: a list of extracted inception features
            corresponding to given dims.
    """
    assert all(dim in InceptionV3.BLOCK_INDEX_BY_DIM for dim in dims)
    if device is None:
        device = dist.device()

    if not isinstance(images, DataLoader):
        num_images = len(images)
        if dist.world_size() > 1:
            sampler = DistributedSampler(
                TensorDataset(images), shuffle=False)
        else:
            sampler = SequentialSampler(TensorDataset(images))
        # print(sampler)
        loader = DataLoader(
            images,
            batch_size=batch_size,
            sampler=sampler,
            drop_last=False,
        )
    else:
        num_images = len(images.dataset)
        loader = images

    block_idxs = [InceptionV3.BLOCK_INDEX_BY_DIM[dim] for dim in dims]
    # Initialize InceptionV3 model
    if dist.rank() == 0:
        # Only rank 0 initialize download the model
        model = InceptionV3(block_idxs).to(device)
        dist.barrier()
    else:
        # Other ranks wait until rank 0 download the model
        dist.barrier()
        model = InceptionV3(block_idxs).to(device)
    model.eval()

    if dist.rank() == 0:
        if use_torch:
            features = [
                torch.empty((num_images, dim)).to(device)
                for dim in dims]
        else:
            features = [
                np.empty((num_images, dim))
                for dim in dims]

    pbar = tqdm(
        total=num_images, dynamic_ncols=True, leave=False,
        disable=(dist.rank() != 0 or not verbose),
        desc="get_inception_feature")
    start = 0
    for batch_images in loader:
        batch_images = batch_images.to(device)
        # calculate inception feature
        end = min(start + len(batch_images) * dist.world_size(), num_images)
        with torch.no_grad():
            outputs = model(batch_images)
            if end == num_images:
                # This is the last batch, so we need to remove the padding.
                if num_images % dist.world_size() != 0:
                    is_padded = dist.rank() >= num_images % dist.world_size()
                else:
                    is_padded = False
                if is_padded:
                    # Remove the padding
                    outputs = [output[: -1] for output in outputs]
                    for output in outputs:
                        print(output.shape)
            outputs = [dist.gather(output) for output in outputs]
            if dist.rank() == 0:
                for feature, output, dim in zip(features, outputs, dims):
                    if use_torch:
                        feature[start: end] = output.view(-1, dim)
                    else:
                        feature[start: end] = output.view(-1, dim).cpu().numpy()
                dist.barrier()
            else:
                dist.barrier()
        pbar.update(end - start)
        start = end
    pbar.close()
    assert start == num_images
    if dist.rank() == 0:
        return features
    else:
        return [None for _ in range(len(dims))]


def torch_cov(m, rowvar=False):
    """Estimate a covariance matrix given data.

    Covariance indicates the level to which two variables vary together.
    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
    then the covariance matrix element `C_{ij}` is the covariance of
    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.

    Args:
        m: A 1-D or 2-D array containing multiple variables and observations.
           Each row of `m` represents a variable, and each column a single
           observation of all those variables.
        rowvar: If `rowvar` is True, then each row represents a
                variable, with observations in the columns. Otherwise, the
                relationship is transposed: each column represents a variable,
                while the rows contain observations.

    Returns:
        The covariance matrix of the variables.
    """
    if m.dim() > 2:
        raise ValueError('m has more than 2 dimensions')
    if m.dim() < 2:
        m = m.view(1, -1)
    if not rowvar and m.size(0) != 1:
        m = m.t()
    # m = m.type(torch.double)  # uncomment this line if desired
    fact = 1.0 / (m.size(1) - 1)
    m -= torch.mean(m, dim=1, keepdim=True)
    mt = m.t()  # if complex: mt = m.t().conj()
    return fact * m.matmul(mt).squeeze()


# Pytorch implementation of matrix sqrt, from Tsung-Yu Lin, and Subhransu Maji
# https://github.com/msubhransu/matrix-sqrt
def sqrt_newton_schulz(A, numIters, dtype=None): # noqa
    with torch.no_grad():
        if dtype is None:
            dtype = A.type()
        batchSize = A.shape[0]
        dim = A.shape[1]
        normA = A.mul(A).sum(dim=1).sum(dim=1).sqrt()
        Y = A.div(normA.view(batchSize, 1, 1).expand_as(A))
        K = torch.eye(dim, dim).view(1, dim, dim).repeat(batchSize, 1, 1)
        Z = torch.eye(dim, dim).view(1, dim, dim).repeat(batchSize, 1, 1)
        K = K.type(dtype)
        Z = Z.type(dtype)
        for i in range(numIters):
            T = 0.5 * (3.0 * K - Z.bmm(Y))
            Y = Y.bmm(T)
            Z = T.bmm(Z)
        sA = Y * torch.sqrt(normA).view(batchSize, 1, 1).expand_as(A)
    return sA


def calculate_frechet_inception_distance(
    acts: Union[torch.FloatTensor, np.ndarray],
    mu: np.ndarray,
    sigma: np.ndarray,
    use_torch: bool = False,
    eps: float = 1e-6,
    device: torch.device = torch.device('cuda:0'),
) -> float: # noqa
    if use_torch:
        m1 = torch.mean(acts, axis=0)
        s1 = torch_cov(acts, rowvar=False)
        mu = torch.tensor(mu).to(m1.dtype).to(device)
        sigma = torch.tensor(sigma).to(s1.dtype).to(device)
    else:
        m1 = np.mean(acts, axis=0)
        s1 = np.cov(acts, rowvar=False)
    return calculate_frechet_distance(m1, s1, mu, sigma, use_torch, eps)


def calculate_frechet_distance(
    mu1: Union[torch.FloatTensor, np.ndarray],
    sigma1: Union[torch.FloatTensor, np.ndarray],
    mu2: Union[torch.FloatTensor, np.ndarray],
    sigma2: Union[torch.FloatTensor, np.ndarray],
    use_torch: bool = False,
    eps: float = 1e-6,
) -> float:
    """Calculate Frechet Distance.

    Args:
        mu1: The sample mean over activations for a set of samples.
        sigma1: The covariance matrix over activations for a set of samples.
        mu2: The sample mean over activations for another set of samples.
        sigma2: The covariance matrix over activations for another set of
                samples.
        use_torch: When True, use torch to calculate FID. Otherwise, use numpy.
        eps: prevent covmean from being singular matrix

    Returns:
        The Frechet Distance.
    """
    if use_torch:
        assert mu1.shape == mu2.shape, \
            'Training and test mean vectors have different lengths'
        assert sigma1.shape == sigma2.shape, \
            'Training and test covariances have different dimensions'

        diff = mu1 - mu2
        # Run 50 itrs of newton-schulz to get the matrix sqrt of
        # sigma1 dot sigma2
        covmean = sqrt_newton_schulz(sigma1.mm(sigma2).unsqueeze(0), 50)
        if torch.any(torch.isnan(covmean)):
            return float('nan')
        covmean = covmean.squeeze()
        out = (diff.dot(diff) +                         # noqa: W504
               torch.trace(sigma1) +                    # noqa: W504
               torch.trace(sigma2) -                    # noqa: W504
               2 * torch.trace(covmean)).cpu().item()
    else:
        mu1 = np.atleast_1d(mu1)
        mu2 = np.atleast_1d(mu2)

        sigma1 = np.atleast_2d(sigma1)
        sigma2 = np.atleast_2d(sigma2)

        assert mu1.shape == mu2.shape, \
            'Training and test mean vectors have different lengths'
        assert sigma1.shape == sigma2.shape, \
            'Training and test covariances have different dimensions'

        diff = mu1 - mu2

        # Product might be almost singular
        covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
        if not np.isfinite(covmean).all():
            msg = ('fid calculation produces singular product; '
                   'adding %s to diagonal of cov estimates') % eps
            print(msg)
            offset = np.eye(sigma1.shape[0]) * eps
            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

        # Numerical error might give slight imaginary component
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
                m = np.max(np.abs(covmean.imag))
                raise ValueError('Imaginary component {}'.format(m))
            covmean = covmean.real

        tr_covmean = np.trace(covmean)

        out = (diff.dot(diff) +         # noqa: W504
               np.trace(sigma1) +       # noqa: W504
               np.trace(sigma2) -       # noqa: W504
               2 * tr_covmean).item()
    return out


def calculate_inception_score(
    probs: Union[torch.FloatTensor, np.ndarray],
    splits: int = 10,
    use_torch: bool = False,
) -> Tuple[float, float]: # noqa
    # Inception Score
    scores = []
    for i in range(splits):
        part = probs[
            (i * probs.shape[0] // splits):
            ((i + 1) * probs.shape[0] // splits), :]
        if use_torch:
            kl = part * (
                torch.log(part) -                                   # noqa: W504
                torch.log(torch.unsqueeze(torch.mean(part, 0), 0)))
            kl = torch.mean(torch.sum(kl, 1))
            scores.append(torch.exp(kl))
        else:
            kl = part * (
                np.log(part) -                                      # noqa: W504
                np.log(np.expand_dims(np.mean(part, 0), 0)))
            kl = np.mean(np.sum(kl, 1))
            scores.append(np.exp(kl))
    if use_torch:
        scores = torch.stack(scores)
        inception_score = torch.mean(scores).cpu().item()
        std = torch.std(scores).cpu().item()
    else:
        inception_score, std = (np.mean(scores).item(), np.std(scores).item())
    del probs, scores
    return inception_score, std
