import numpy as np
from scipy import linalg
from packaging import version
import scipy

from .inception import InceptionV3
import torch.nn.functional as F

class RepeatIterator:
    def __init__(self, iterator, n):
        self.iterator = iterator
        self.n = n
        
    def __iter__(self):
        for _ in range(self.n):
            yield from self.iterator
            
    def __len__(self):
        return len(self.iterator) * self.n

def init_inception():
    block_dim = 2048
    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[block_dim]
    print('Initializing Inception network...')
    scipy_version = version.parse(scipy.__version__)
    if scipy_version > version.parse('1.3.3') and scipy_version < version.parse('1.5.0'):
        print('[Performance warning] If this step takes too long, downgrade SciPy to version 1.3.3 or update it to 1.5+, e.g. '
              'pip install scipy==1.5.2')
    inception_model = InceptionV3([block_idx])
    return inception_model

def forward_inception_batch(inception_model, images):
    pred = inception_model(images)[0]
    if pred.shape[2] != 1 or pred.shape[3] != 1:
        pred = F.adaptive_avg_pool2d(pred, output_size=(1, 1))
    return pred.data.cpu().numpy().reshape(images.shape[0], -1)

def calculate_stats(act):
    mu = np.mean(act, axis=0)
    sigma = np.cov(act, rowvar=False)
    return mu, sigma

# Borrowed from https://github.com/mseitzer/pytorch-fid
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
    """Numpy implementation of the Frechet Distance.
    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
    and X_2 ~ N(mu_2, C_2) is
            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).

    Stable version by Dougal J. Sutherland.

    Params:
    -- mu1 : Numpy array containing the activations of the pool_3 layer of the
             inception net ( like returned by the function 'get_predictions')
             for generated samples.
    -- mu2   : The sample mean over activations of the pool_3 layer, precalcualted
               on an representive data set.
    -- sigma1: The covariance matrix over activations of the pool_3 layer for
               generated samples.
    -- sigma2: The covariance matrix over activations of the pool_3 layer,
               precalcualted on an representive data set.

    Returns:
    --   : The Frechet Distance.
    """

    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
        warnings.warn(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)

    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean