import jax
import jax.numpy as jnp
import flax
import flax.linen as nn
import numpy as np
from PIL import Image
import os
import functools
import scipy
from tqdm import tqdm

from . import inception
from . import utils


def compute_statistics(path, params, apply_fn, batch_size=1, img_size=None):
    if path.endswith('.npz'):
        stats = np.load(path)
        mu, sigma = stats['mu'], stats['sigma']
        return mu, sigma

    images = []
    for f in os.listdir(path):
        img = Image.open(os.path.join(path, f))
        img = jnp.array(img) / 255.0
        if img_size is not None:
            img = jax.image.resize(img, shape=(img_size[0], img_size[1], img.shape[2]), method='bilinear', antialias=False)
        images.append(img)
    images = jnp.array(images)
    
    num_batches = max(1, images.shape[0] // batch_size)
    act = []
    for i in tqdm(range(num_batches)):
        x = images[i * batch_size:i * batch_size + batch_size]
        x = 2 * x - 1
        pred = apply_fn(params, jax.lax.stop_gradient(x))
        act.append(pred.squeeze(axis=1).squeeze(axis=1))
    act = jnp.concatenate(act, axis=0)

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


def compute_frechet_distance(mu1, mu2, sigma1, sigma2, eps=1e-6):
    # Taken from: https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py
    mu1 = np.atleast_1d(mu1)
    mu2 = np.atleast_1d(mu2)
    sigma1 = np.atleast_1d(sigma1)
    sigma2 = np.atleast_1d(sigma2)

    assert mu1.shape == mu2.shape
    assert sigma1.shape == sigma2.shape

    diff = mu1 - mu2

    covmean, _ = scipy.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 = scipy.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)



