import torch
from scipy.stats import entropy
from scipy import linalg
import numpy as np
import prdc
import pickle
import os
import tqdm


class GenerativeModelScore:
    def __init__(self):
        
        self.inception_model = torch.hub.load('pytorch/vision:v0.9.0', 'inception_v3', pretrained=True)
        self.inception_model.forward = self._forward
        self.inception_model.eval()

        self.real_images = None
        self.fake_images = None
        
        
        self.real_predict_softmax = None
        self.real_feature = None
        
        self.fake_predict_softmax = None
        self.fake_feature = None
        
        self.lazy = False
        
    def _forward(self, x):
        import torchvision

        if x.size(1) != 3 : 
            x = self.inception_model._transform_input(x)

        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)

        # N x 3 x 299 x 299
        x = self.inception_model.Conv2d_1a_3x3(x)
        # N x 32 x 149 x 149
        x = self.inception_model.Conv2d_2a_3x3(x)
        # N x 32 x 147 x 147
        x = self.inception_model.Conv2d_2b_3x3(x)
        # N x 64 x 147 x 147
        x = self.inception_model.maxpool1(x)
        # N x 64 x 73 x 73
        x = self.inception_model.Conv2d_3b_1x1(x)
        # N x 80 x 73 x 73
        x = self.inception_model.Conv2d_4a_3x3(x)
        # N x 192 x 71 x 71
        x = self.inception_model.maxpool2(x)
        # N x 192 x 35 x 35
        x = self.inception_model.Mixed_5b(x)
        # N x 256 x 35 x 35
        x = self.inception_model.Mixed_5c(x)
        # N x 288 x 35 x 35
        x = self.inception_model.Mixed_5d(x)
        # N x 288 x 35 x 35
        x = self.inception_model.Mixed_6a(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6b(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6c(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6d(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6e(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_7a(x)
        # N x 1280 x 8 x 8
        x = self.inception_model.Mixed_7b(x)
        # N x 2048 x 8 x 8
        x = self.inception_model.Mixed_7c(x)
        # N x 2048 x 8 x 8
        # Adaptive average pooling
        x = self.inception_model.avgpool(x)
        # N x 2048 x 1 x 1
        feature = x.detach()
        x = self.inception_model.dropout(x)
        # N x 2048 x 1 x 1
        x = torch.flatten(x, 1)
        # N x 2048
        x = self.inception_model.fc(x)
        # N x 1000 (num_classes)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()

        # Now compute the mean kl-div
        split_scores = []

        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))

        return np.mean(split_scores), np.std(split_scores)
    
    
    def trainloaderinfo_to_hashedname(self, train_loader) : 
        import hashlib
        dataset_info = str(train_loader.dataset).split('\n')
        name, datapoints, split = dataset_info[0:3]
        transform = dataset_info[4:]
        except_root_info = [name, datapoints, split] + transform
        name = hashlib.md5(str(except_root_info).encode()).hexdigest() + '.pickle'
        return name
        

    def load_or_make(self, train_loader, inception_model_info_path='../inception_model_info/', device='cuda:0', run_test=False, force_make=False) : 
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path + real_images_info_file_name) and not force_make: 
            print("Using generated real image info.")
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path + real_images_info_file_name)

        else : 
            self.model_to(device)

            #put real image
            for each_batch in tqdm.tqdm(train_loader, desc='insert real dataset') : 
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test : break

            #generate real images info
            self.lazy_forward(batch_size=64, device=device, real_forward=True, shuffle=False)
            
            self.calculate_real_image_statistics()
            #save real images info for next experiments
            self.save_real_images_info(inception_model_info_path + real_images_info_file_name)
            #offload inception_model
            self.model_to('cpu')
        
        
    def feature_to_mu_sig(self, act):
        mu = np.mean(act, axis=0)
        sigma = np.cov(act, rowvar=False)
        return mu, sigma

    def calculate_frechet_distance(self, 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)).
        """

        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

        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))

        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)
    
    def model_to(self, device) : 
        self.inception_model = self.inception_model.to(device)
    
    def lazy_mode(self, tf) :
        self.lazy = tf
    
    def put_real(self, real_images) : 
        if self.lazy : 
            if self.real_images is None : 
                self.real_images = real_images
            else : 
                self.real_images = torch.cat([self.real_images, real_images])
        else : 
            self.real_forward(real_images)
            
            
    def real_forward(self, real_images):
        real_predict_softmax, real_feature = self.analysis_softmax_and_feature(real_images)
        if self.real_predict_softmax is None : 
            self.real_predict_softmax = real_predict_softmax.detach().cpu()
            self.real_feature = real_feature.detach().cpu()
        else : 
            self.real_predict_softmax = torch.cat([self.real_predict_softmax, real_predict_softmax.detach().cpu()])
            self.real_feature = torch.cat([self.real_feature, real_feature.detach().cpu()])    
                
  
        
    def put_fake(self, fake_images) : 
        if self.lazy : 
            if self.fake_images is None : 
                self.fake_images = fake_images
            else : 
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)
            
    
    def fake_forward(self, fake_images) : 
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(fake_images)
        if self.fake_predict_softmax is None : 
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else : 
            self.fake_predict_softmax = torch.cat([self.fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.detach().cpu()])
    
    
                  
    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4, real_forward=False, fake_forward=False, device='cpu') : 
        
        assert self.lazy, "lazy_forward only run in lazy mode. call lazy_mode() first."
        
        from torch.utils.data import  TensorDataset, DataLoader
        import tqdm 
        
        if real_forward : 
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
            print("generate real images info")
            for real_images in tqdm.tqdm(real_loader) : 
                self.real_forward(real_images[0].to(device))
        
        if fake_forward : 
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
            print("generate fake images info")
            for fake_images in tqdm.tqdm(fake_loader) :
                self.fake_forward(fake_images[0].to(device))
            
    def save_real_images_info(self, file_name='real_images_info.pickle') : 
        with open(file_name, 'wb') as f :
            pickle.dump((self.real_inception_score, self.real_feature_np, (self.real_mu, self.real_sigma)), f)
    
    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            (self.real_inception_score, self.real_feature_np, (self.real_mu, self.real_sigma)) = \
                pickle.load(f)
            
    def calculate_real_image_statistics(self) : 
        self.real_inception_score = self.predict_to_inception_score(self.real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.real_feature_np)        
    
    def calculate_fake_image_statistics(self) : 
        self.fake_inception_score = self.predict_to_inception_score(self.fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.fake_feature_np)
    
    def clear_fake(self) : 
        self.fake_images = None
        self.fake_predict_softmax = None
        self.fake_feature = None
        self.fake_mu, self.fake_sigma = None, None
        
    def calculate_generative_score(self):
        fid = self.calculate_frechet_distance(self.real_mu, self.real_sigma, self.fake_mu, self.fake_sigma)
        real_pick = np.random.permutation(self.real_feature_np)[:10000]
        fake_pick = np.random.permutation(self.fake_feature_np)[:10000]
        metrics = prdc.compute_prdc(real_features=real_pick, fake_features=fake_pick, nearest_k=5)
        metrics['fid'] = fid
        metrics['real_is'] = self.real_inception_score
        metrics['fake_is'] = self.fake_inception_score
        return metrics
    

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)

8fc96468afc32a0e329ad48d6f3ddf8692d0c1bf
<mask token>


class GenerativeModelScore:
    <mask token>

    def _forward(self, x):
        import torchvision
        if x.size(1) != 3:
            x = self.inception_model._transform_input(x)
        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)
        x = self.inception_model.Conv2d_1a_3x3(x)
        x = self.inception_model.Conv2d_2a_3x3(x)
        x = self.inception_model.Conv2d_2b_3x3(x)
        x = self.inception_model.maxpool1(x)
        x = self.inception_model.Conv2d_3b_1x1(x)
        x = self.inception_model.Conv2d_4a_3x3(x)
        x = self.inception_model.maxpool2(x)
        x = self.inception_model.Mixed_5b(x)
        x = self.inception_model.Mixed_5c(x)
        x = self.inception_model.Mixed_5d(x)
        x = self.inception_model.Mixed_6a(x)
        x = self.inception_model.Mixed_6b(x)
        x = self.inception_model.Mixed_6c(x)
        x = self.inception_model.Mixed_6d(x)
        x = self.inception_model.Mixed_6e(x)
        x = self.inception_model.Mixed_7a(x)
        x = self.inception_model.Mixed_7b(x)
        x = self.inception_model.Mixed_7c(x)
        x = self.inception_model.avgpool(x)
        feature = x.detach()
        x = self.inception_model.dropout(x)
        x = torch.flatten(x, 1)
        x = self.inception_model.fc(x)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()
        split_scores = []
        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))
        return np.mean(split_scores), np.std(split_scores)
    <mask token>

    def load_or_make(self, train_loader, inception_model_info_path=
        '../inception_model_info/', device='cuda:0', run_test=False,
        force_make=False):
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(
            train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path +
            real_images_info_file_name) and not force_make:
            print('Using generated real image info.')
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path +
                real_images_info_file_name)
        else:
            self.model_to(device)
            for each_batch in tqdm.tqdm(train_loader, desc=
                'insert real dataset'):
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test:
                    break
            self.lazy_forward(batch_size=64, device=device, real_forward=
                True, shuffle=False)
            self.calculate_real_image_statistics()
            self.save_real_images_info(inception_model_info_path +
                real_images_info_file_name)
            self.model_to('cpu')
    <mask token>

    def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-06):
        """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)).
        """
        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
        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))
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
                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
    <mask token>

    def lazy_mode(self, tf):
        self.lazy = tf
    <mask token>
    <mask token>

    def put_fake(self, fake_images):
        if self.lazy:
            if self.fake_images is None:
                self.fake_images = fake_images
            else:
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)

    def fake_forward(self, fake_images):
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(
            fake_images)
        if self.fake_predict_softmax is None:
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else:
            self.fake_predict_softmax = torch.cat([self.
                fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.
                detach().cpu()])

    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4,
        real_forward=False, fake_forward=False, device='cpu'):
        assert self.lazy, 'lazy_forward only run in lazy mode. call lazy_mode() first.'
        from torch.utils.data import TensorDataset, DataLoader
        import tqdm
        if real_forward:
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate real images info')
            for real_images in tqdm.tqdm(real_loader):
                self.real_forward(real_images[0].to(device))
        if fake_forward:
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate fake images info')
            for fake_images in tqdm.tqdm(fake_loader):
                self.fake_forward(fake_images[0].to(device))
    <mask token>

    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            self.real_inception_score, self.real_feature_np, (self.real_mu,
                self.real_sigma) = pickle.load(f)

    def calculate_real_image_statistics(self):
        self.real_inception_score = self.predict_to_inception_score(self.
            real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.
            real_feature_np)

    def calculate_fake_image_statistics(self):
        self.fake_inception_score = self.predict_to_inception_score(self.
            fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.
            fake_feature_np)
    <mask token>
    <mask token>

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)

<mask token>


class GenerativeModelScore:
    <mask token>

    def _forward(self, x):
        import torchvision
        if x.size(1) != 3:
            x = self.inception_model._transform_input(x)
        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)
        x = self.inception_model.Conv2d_1a_3x3(x)
        x = self.inception_model.Conv2d_2a_3x3(x)
        x = self.inception_model.Conv2d_2b_3x3(x)
        x = self.inception_model.maxpool1(x)
        x = self.inception_model.Conv2d_3b_1x1(x)
        x = self.inception_model.Conv2d_4a_3x3(x)
        x = self.inception_model.maxpool2(x)
        x = self.inception_model.Mixed_5b(x)
        x = self.inception_model.Mixed_5c(x)
        x = self.inception_model.Mixed_5d(x)
        x = self.inception_model.Mixed_6a(x)
        x = self.inception_model.Mixed_6b(x)
        x = self.inception_model.Mixed_6c(x)
        x = self.inception_model.Mixed_6d(x)
        x = self.inception_model.Mixed_6e(x)
        x = self.inception_model.Mixed_7a(x)
        x = self.inception_model.Mixed_7b(x)
        x = self.inception_model.Mixed_7c(x)
        x = self.inception_model.avgpool(x)
        feature = x.detach()
        x = self.inception_model.dropout(x)
        x = torch.flatten(x, 1)
        x = self.inception_model.fc(x)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()
        split_scores = []
        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))
        return np.mean(split_scores), np.std(split_scores)
    <mask token>

    def load_or_make(self, train_loader, inception_model_info_path=
        '../inception_model_info/', device='cuda:0', run_test=False,
        force_make=False):
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(
            train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path +
            real_images_info_file_name) and not force_make:
            print('Using generated real image info.')
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path +
                real_images_info_file_name)
        else:
            self.model_to(device)
            for each_batch in tqdm.tqdm(train_loader, desc=
                'insert real dataset'):
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test:
                    break
            self.lazy_forward(batch_size=64, device=device, real_forward=
                True, shuffle=False)
            self.calculate_real_image_statistics()
            self.save_real_images_info(inception_model_info_path +
                real_images_info_file_name)
            self.model_to('cpu')

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

    def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-06):
        """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)).
        """
        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
        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))
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
                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

    def model_to(self, device):
        self.inception_model = self.inception_model.to(device)

    def lazy_mode(self, tf):
        self.lazy = tf
    <mask token>

    def real_forward(self, real_images):
        real_predict_softmax, real_feature = self.analysis_softmax_and_feature(
            real_images)
        if self.real_predict_softmax is None:
            self.real_predict_softmax = real_predict_softmax.detach().cpu()
            self.real_feature = real_feature.detach().cpu()
        else:
            self.real_predict_softmax = torch.cat([self.
                real_predict_softmax, real_predict_softmax.detach().cpu()])
            self.real_feature = torch.cat([self.real_feature, real_feature.
                detach().cpu()])

    def put_fake(self, fake_images):
        if self.lazy:
            if self.fake_images is None:
                self.fake_images = fake_images
            else:
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)

    def fake_forward(self, fake_images):
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(
            fake_images)
        if self.fake_predict_softmax is None:
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else:
            self.fake_predict_softmax = torch.cat([self.
                fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.
                detach().cpu()])

    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4,
        real_forward=False, fake_forward=False, device='cpu'):
        assert self.lazy, 'lazy_forward only run in lazy mode. call lazy_mode() first.'
        from torch.utils.data import TensorDataset, DataLoader
        import tqdm
        if real_forward:
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate real images info')
            for real_images in tqdm.tqdm(real_loader):
                self.real_forward(real_images[0].to(device))
        if fake_forward:
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate fake images info')
            for fake_images in tqdm.tqdm(fake_loader):
                self.fake_forward(fake_images[0].to(device))
    <mask token>

    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            self.real_inception_score, self.real_feature_np, (self.real_mu,
                self.real_sigma) = pickle.load(f)

    def calculate_real_image_statistics(self):
        self.real_inception_score = self.predict_to_inception_score(self.
            real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.
            real_feature_np)

    def calculate_fake_image_statistics(self):
        self.fake_inception_score = self.predict_to_inception_score(self.
            fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.
            fake_feature_np)
    <mask token>

    def calculate_generative_score(self):
        fid = self.calculate_frechet_distance(self.real_mu, self.real_sigma,
            self.fake_mu, self.fake_sigma)
        real_pick = np.random.permutation(self.real_feature_np)[:10000]
        fake_pick = np.random.permutation(self.fake_feature_np)[:10000]
        metrics = prdc.compute_prdc(real_features=real_pick, fake_features=
            fake_pick, nearest_k=5)
        metrics['fid'] = fid
        metrics['real_is'] = self.real_inception_score
        metrics['fake_is'] = self.fake_inception_score
        return metrics

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)

<mask token>


class GenerativeModelScore:

    def __init__(self):
        self.inception_model = torch.hub.load('pytorch/vision:v0.9.0',
            'inception_v3', pretrained=True)
        self.inception_model.forward = self._forward
        self.inception_model.eval()
        self.real_images = None
        self.fake_images = None
        self.real_predict_softmax = None
        self.real_feature = None
        self.fake_predict_softmax = None
        self.fake_feature = None
        self.lazy = False

    def _forward(self, x):
        import torchvision
        if x.size(1) != 3:
            x = self.inception_model._transform_input(x)
        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)
        x = self.inception_model.Conv2d_1a_3x3(x)
        x = self.inception_model.Conv2d_2a_3x3(x)
        x = self.inception_model.Conv2d_2b_3x3(x)
        x = self.inception_model.maxpool1(x)
        x = self.inception_model.Conv2d_3b_1x1(x)
        x = self.inception_model.Conv2d_4a_3x3(x)
        x = self.inception_model.maxpool2(x)
        x = self.inception_model.Mixed_5b(x)
        x = self.inception_model.Mixed_5c(x)
        x = self.inception_model.Mixed_5d(x)
        x = self.inception_model.Mixed_6a(x)
        x = self.inception_model.Mixed_6b(x)
        x = self.inception_model.Mixed_6c(x)
        x = self.inception_model.Mixed_6d(x)
        x = self.inception_model.Mixed_6e(x)
        x = self.inception_model.Mixed_7a(x)
        x = self.inception_model.Mixed_7b(x)
        x = self.inception_model.Mixed_7c(x)
        x = self.inception_model.avgpool(x)
        feature = x.detach()
        x = self.inception_model.dropout(x)
        x = torch.flatten(x, 1)
        x = self.inception_model.fc(x)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()
        split_scores = []
        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))
        return np.mean(split_scores), np.std(split_scores)
    <mask token>

    def load_or_make(self, train_loader, inception_model_info_path=
        '../inception_model_info/', device='cuda:0', run_test=False,
        force_make=False):
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(
            train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path +
            real_images_info_file_name) and not force_make:
            print('Using generated real image info.')
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path +
                real_images_info_file_name)
        else:
            self.model_to(device)
            for each_batch in tqdm.tqdm(train_loader, desc=
                'insert real dataset'):
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test:
                    break
            self.lazy_forward(batch_size=64, device=device, real_forward=
                True, shuffle=False)
            self.calculate_real_image_statistics()
            self.save_real_images_info(inception_model_info_path +
                real_images_info_file_name)
            self.model_to('cpu')

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

    def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-06):
        """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)).
        """
        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
        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))
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
                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

    def model_to(self, device):
        self.inception_model = self.inception_model.to(device)

    def lazy_mode(self, tf):
        self.lazy = tf
    <mask token>

    def real_forward(self, real_images):
        real_predict_softmax, real_feature = self.analysis_softmax_and_feature(
            real_images)
        if self.real_predict_softmax is None:
            self.real_predict_softmax = real_predict_softmax.detach().cpu()
            self.real_feature = real_feature.detach().cpu()
        else:
            self.real_predict_softmax = torch.cat([self.
                real_predict_softmax, real_predict_softmax.detach().cpu()])
            self.real_feature = torch.cat([self.real_feature, real_feature.
                detach().cpu()])

    def put_fake(self, fake_images):
        if self.lazy:
            if self.fake_images is None:
                self.fake_images = fake_images
            else:
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)

    def fake_forward(self, fake_images):
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(
            fake_images)
        if self.fake_predict_softmax is None:
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else:
            self.fake_predict_softmax = torch.cat([self.
                fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.
                detach().cpu()])

    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4,
        real_forward=False, fake_forward=False, device='cpu'):
        assert self.lazy, 'lazy_forward only run in lazy mode. call lazy_mode() first.'
        from torch.utils.data import TensorDataset, DataLoader
        import tqdm
        if real_forward:
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate real images info')
            for real_images in tqdm.tqdm(real_loader):
                self.real_forward(real_images[0].to(device))
        if fake_forward:
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate fake images info')
            for fake_images in tqdm.tqdm(fake_loader):
                self.fake_forward(fake_images[0].to(device))

    def save_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'wb') as f:
            pickle.dump((self.real_inception_score, self.real_feature_np, (
                self.real_mu, self.real_sigma)), f)

    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            self.real_inception_score, self.real_feature_np, (self.real_mu,
                self.real_sigma) = pickle.load(f)

    def calculate_real_image_statistics(self):
        self.real_inception_score = self.predict_to_inception_score(self.
            real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.
            real_feature_np)

    def calculate_fake_image_statistics(self):
        self.fake_inception_score = self.predict_to_inception_score(self.
            fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.
            fake_feature_np)
    <mask token>

    def calculate_generative_score(self):
        fid = self.calculate_frechet_distance(self.real_mu, self.real_sigma,
            self.fake_mu, self.fake_sigma)
        real_pick = np.random.permutation(self.real_feature_np)[:10000]
        fake_pick = np.random.permutation(self.fake_feature_np)[:10000]
        metrics = prdc.compute_prdc(real_features=real_pick, fake_features=
            fake_pick, nearest_k=5)
        metrics['fid'] = fid
        metrics['real_is'] = self.real_inception_score
        metrics['fake_is'] = self.fake_inception_score
        return metrics

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)

import torch
from scipy.stats import entropy
from scipy import linalg
import numpy as np
import prdc
import pickle
import os
import tqdm


class GenerativeModelScore:

    def __init__(self):
        self.inception_model = torch.hub.load('pytorch/vision:v0.9.0',
            'inception_v3', pretrained=True)
        self.inception_model.forward = self._forward
        self.inception_model.eval()
        self.real_images = None
        self.fake_images = None
        self.real_predict_softmax = None
        self.real_feature = None
        self.fake_predict_softmax = None
        self.fake_feature = None
        self.lazy = False

    def _forward(self, x):
        import torchvision
        if x.size(1) != 3:
            x = self.inception_model._transform_input(x)
        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)
        x = self.inception_model.Conv2d_1a_3x3(x)
        x = self.inception_model.Conv2d_2a_3x3(x)
        x = self.inception_model.Conv2d_2b_3x3(x)
        x = self.inception_model.maxpool1(x)
        x = self.inception_model.Conv2d_3b_1x1(x)
        x = self.inception_model.Conv2d_4a_3x3(x)
        x = self.inception_model.maxpool2(x)
        x = self.inception_model.Mixed_5b(x)
        x = self.inception_model.Mixed_5c(x)
        x = self.inception_model.Mixed_5d(x)
        x = self.inception_model.Mixed_6a(x)
        x = self.inception_model.Mixed_6b(x)
        x = self.inception_model.Mixed_6c(x)
        x = self.inception_model.Mixed_6d(x)
        x = self.inception_model.Mixed_6e(x)
        x = self.inception_model.Mixed_7a(x)
        x = self.inception_model.Mixed_7b(x)
        x = self.inception_model.Mixed_7c(x)
        x = self.inception_model.avgpool(x)
        feature = x.detach()
        x = self.inception_model.dropout(x)
        x = torch.flatten(x, 1)
        x = self.inception_model.fc(x)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()
        split_scores = []
        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))
        return np.mean(split_scores), np.std(split_scores)

    def trainloaderinfo_to_hashedname(self, train_loader):
        import hashlib
        dataset_info = str(train_loader.dataset).split('\n')
        name, datapoints, split = dataset_info[0:3]
        transform = dataset_info[4:]
        except_root_info = [name, datapoints, split] + transform
        name = hashlib.md5(str(except_root_info).encode()).hexdigest(
            ) + '.pickle'
        return name

    def load_or_make(self, train_loader, inception_model_info_path=
        '../inception_model_info/', device='cuda:0', run_test=False,
        force_make=False):
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(
            train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path +
            real_images_info_file_name) and not force_make:
            print('Using generated real image info.')
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path +
                real_images_info_file_name)
        else:
            self.model_to(device)
            for each_batch in tqdm.tqdm(train_loader, desc=
                'insert real dataset'):
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test:
                    break
            self.lazy_forward(batch_size=64, device=device, real_forward=
                True, shuffle=False)
            self.calculate_real_image_statistics()
            self.save_real_images_info(inception_model_info_path +
                real_images_info_file_name)
            self.model_to('cpu')

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

    def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-06):
        """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)).
        """
        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
        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))
        if np.iscomplexobj(covmean):
            if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
                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

    def model_to(self, device):
        self.inception_model = self.inception_model.to(device)

    def lazy_mode(self, tf):
        self.lazy = tf

    def put_real(self, real_images):
        if self.lazy:
            if self.real_images is None:
                self.real_images = real_images
            else:
                self.real_images = torch.cat([self.real_images, real_images])
        else:
            self.real_forward(real_images)

    def real_forward(self, real_images):
        real_predict_softmax, real_feature = self.analysis_softmax_and_feature(
            real_images)
        if self.real_predict_softmax is None:
            self.real_predict_softmax = real_predict_softmax.detach().cpu()
            self.real_feature = real_feature.detach().cpu()
        else:
            self.real_predict_softmax = torch.cat([self.
                real_predict_softmax, real_predict_softmax.detach().cpu()])
            self.real_feature = torch.cat([self.real_feature, real_feature.
                detach().cpu()])

    def put_fake(self, fake_images):
        if self.lazy:
            if self.fake_images is None:
                self.fake_images = fake_images
            else:
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)

    def fake_forward(self, fake_images):
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(
            fake_images)
        if self.fake_predict_softmax is None:
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else:
            self.fake_predict_softmax = torch.cat([self.
                fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.
                detach().cpu()])

    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4,
        real_forward=False, fake_forward=False, device='cpu'):
        assert self.lazy, 'lazy_forward only run in lazy mode. call lazy_mode() first.'
        from torch.utils.data import TensorDataset, DataLoader
        import tqdm
        if real_forward:
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate real images info')
            for real_images in tqdm.tqdm(real_loader):
                self.real_forward(real_images[0].to(device))
        if fake_forward:
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size,
                shuffle=shuffle, num_workers=num_workers)
            print('generate fake images info')
            for fake_images in tqdm.tqdm(fake_loader):
                self.fake_forward(fake_images[0].to(device))

    def save_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'wb') as f:
            pickle.dump((self.real_inception_score, self.real_feature_np, (
                self.real_mu, self.real_sigma)), f)

    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            self.real_inception_score, self.real_feature_np, (self.real_mu,
                self.real_sigma) = pickle.load(f)

    def calculate_real_image_statistics(self):
        self.real_inception_score = self.predict_to_inception_score(self.
            real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.
            real_feature_np)

    def calculate_fake_image_statistics(self):
        self.fake_inception_score = self.predict_to_inception_score(self.
            fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.
            fake_feature_np)

    def clear_fake(self):
        self.fake_images = None
        self.fake_predict_softmax = None
        self.fake_feature = None
        self.fake_mu, self.fake_sigma = None, None

    def calculate_generative_score(self):
        fid = self.calculate_frechet_distance(self.real_mu, self.real_sigma,
            self.fake_mu, self.fake_sigma)
        real_pick = np.random.permutation(self.real_feature_np)[:10000]
        fake_pick = np.random.permutation(self.fake_feature_np)[:10000]
        metrics = prdc.compute_prdc(real_features=real_pick, fake_features=
            fake_pick, nearest_k=5)
        metrics['fid'] = fid
        metrics['real_is'] = self.real_inception_score
        metrics['fake_is'] = self.fake_inception_score
        return metrics

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)

import torch
from scipy.stats import entropy
from scipy import linalg
import numpy as np
import prdc
import pickle
import os
import tqdm


class GenerativeModelScore:
    def __init__(self):
        
        self.inception_model = torch.hub.load('pytorch/vision:v0.9.0', 'inception_v3', pretrained=True)
        self.inception_model.forward = self._forward
        self.inception_model.eval()

        self.real_images = None
        self.fake_images = None
        
        
        self.real_predict_softmax = None
        self.real_feature = None
        
        self.fake_predict_softmax = None
        self.fake_feature = None
        
        self.lazy = False
        
    def _forward(self, x):
        import torchvision

        if x.size(1) != 3 : 
            x = self.inception_model._transform_input(x)

        resize = torchvision.transforms.Resize((299, 299))
        x = resize(x)

        # N x 3 x 299 x 299
        x = self.inception_model.Conv2d_1a_3x3(x)
        # N x 32 x 149 x 149
        x = self.inception_model.Conv2d_2a_3x3(x)
        # N x 32 x 147 x 147
        x = self.inception_model.Conv2d_2b_3x3(x)
        # N x 64 x 147 x 147
        x = self.inception_model.maxpool1(x)
        # N x 64 x 73 x 73
        x = self.inception_model.Conv2d_3b_1x1(x)
        # N x 80 x 73 x 73
        x = self.inception_model.Conv2d_4a_3x3(x)
        # N x 192 x 71 x 71
        x = self.inception_model.maxpool2(x)
        # N x 192 x 35 x 35
        x = self.inception_model.Mixed_5b(x)
        # N x 256 x 35 x 35
        x = self.inception_model.Mixed_5c(x)
        # N x 288 x 35 x 35
        x = self.inception_model.Mixed_5d(x)
        # N x 288 x 35 x 35
        x = self.inception_model.Mixed_6a(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6b(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6c(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6d(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_6e(x)
        # N x 768 x 17 x 17
        x = self.inception_model.Mixed_7a(x)
        # N x 1280 x 8 x 8
        x = self.inception_model.Mixed_7b(x)
        # N x 2048 x 8 x 8
        x = self.inception_model.Mixed_7c(x)
        # N x 2048 x 8 x 8
        # Adaptive average pooling
        x = self.inception_model.avgpool(x)
        # N x 2048 x 1 x 1
        feature = x.detach()
        x = self.inception_model.dropout(x)
        # N x 2048 x 1 x 1
        x = torch.flatten(x, 1)
        # N x 2048
        x = self.inception_model.fc(x)
        # N x 1000 (num_classes)
        return x, feature

    def predict_to_inception_score(self, predict, splits=1):
        preds = torch.softmax(predict, dim=1).numpy()

        # Now compute the mean kl-div
        split_scores = []

        for k in range(splits):
            part = preds
            py = np.mean(part, axis=0)
            scores = []
            for i in range(part.shape[0]):
                pyx = part[i, :]
                scores.append(entropy(pyx, py))
            split_scores.append(np.exp(np.mean(scores)))

        return np.mean(split_scores), np.std(split_scores)
    
    
    def trainloaderinfo_to_hashedname(self, train_loader) : 
        import hashlib
        dataset_info = str(train_loader.dataset).split('\n')
        name, datapoints, split = dataset_info[0:3]
        transform = dataset_info[4:]
        except_root_info = [name, datapoints, split] + transform
        name = hashlib.md5(str(except_root_info).encode()).hexdigest() + '.pickle'
        return name
        

    def load_or_make(self, train_loader, inception_model_info_path='../inception_model_info/', device='cuda:0', run_test=False, force_make=False) : 
        real_images_info_file_name = self.trainloaderinfo_to_hashedname(train_loader)
        os.makedirs(inception_model_info_path, exist_ok=True)
        if os.path.exists(inception_model_info_path + real_images_info_file_name) and not force_make: 
            print("Using generated real image info.")
            print(train_loader.dataset)
            self.load_real_images_info(inception_model_info_path + real_images_info_file_name)

        else : 
            self.model_to(device)

            #put real image
            for each_batch in tqdm.tqdm(train_loader, desc='insert real dataset') : 
                X_train_batch = each_batch[0]
                self.put_real(X_train_batch)
                if run_test : break

            #generate real images info
            self.lazy_forward(batch_size=64, device=device, real_forward=True, shuffle=False)
            
            self.calculate_real_image_statistics()
            #save real images info for next experiments
            self.save_real_images_info(inception_model_info_path + real_images_info_file_name)
            #offload inception_model
            self.model_to('cpu')
        
        
    def feature_to_mu_sig(self, act):
        mu = np.mean(act, axis=0)
        sigma = np.cov(act, rowvar=False)
        return mu, sigma

    def calculate_frechet_distance(self, 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)).
        """

        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

        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))

        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)
    
    def model_to(self, device) : 
        self.inception_model = self.inception_model.to(device)
    
    def lazy_mode(self, tf) :
        self.lazy = tf
    
    def put_real(self, real_images) : 
        if self.lazy : 
            if self.real_images is None : 
                self.real_images = real_images
            else : 
                self.real_images = torch.cat([self.real_images, real_images])
        else : 
            self.real_forward(real_images)
            
            
    def real_forward(self, real_images):
        real_predict_softmax, real_feature = self.analysis_softmax_and_feature(real_images)
        if self.real_predict_softmax is None : 
            self.real_predict_softmax = real_predict_softmax.detach().cpu()
            self.real_feature = real_feature.detach().cpu()
        else : 
            self.real_predict_softmax = torch.cat([self.real_predict_softmax, real_predict_softmax.detach().cpu()])
            self.real_feature = torch.cat([self.real_feature, real_feature.detach().cpu()])    
                
  
        
    def put_fake(self, fake_images) : 
        if self.lazy : 
            if self.fake_images is None : 
                self.fake_images = fake_images
            else : 
                self.fake_images = torch.cat([self.fake_images, fake_images])
        else:
            self.fake_forward(fake_images)
            
    
    def fake_forward(self, fake_images) : 
        fake_predict_softmax, fake_feature = self.analysis_softmax_and_feature(fake_images)
        if self.fake_predict_softmax is None : 
            self.fake_predict_softmax = fake_predict_softmax.detach().cpu()
            self.fake_feature = fake_feature.detach().cpu()
        else : 
            self.fake_predict_softmax = torch.cat([self.fake_predict_softmax, fake_predict_softmax.detach().cpu()])
            self.fake_feature = torch.cat([self.fake_feature, fake_feature.detach().cpu()])
    
    
                  
    def lazy_forward(self, batch_size=64, shuffle=True, num_workers=4, real_forward=False, fake_forward=False, device='cpu') : 
        
        assert self.lazy, "lazy_forward only run in lazy mode. call lazy_mode() first."
        
        from torch.utils.data import  TensorDataset, DataLoader
        import tqdm 
        
        if real_forward : 
            real_dataset = TensorDataset(self.real_images)
            real_loader = DataLoader(real_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
            print("generate real images info")
            for real_images in tqdm.tqdm(real_loader) : 
                self.real_forward(real_images[0].to(device))
        
        if fake_forward : 
            fake_dataset = TensorDataset(self.fake_images)
            fake_loader = DataLoader(fake_dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
            print("generate fake images info")
            for fake_images in tqdm.tqdm(fake_loader) :
                self.fake_forward(fake_images[0].to(device))
            
    def save_real_images_info(self, file_name='real_images_info.pickle') : 
        with open(file_name, 'wb') as f :
            pickle.dump((self.real_inception_score, self.real_feature_np, (self.real_mu, self.real_sigma)), f)
    
    def load_real_images_info(self, file_name='real_images_info.pickle'):
        with open(file_name, 'rb') as f:
            (self.real_inception_score, self.real_feature_np, (self.real_mu, self.real_sigma)) = \
                pickle.load(f)
            
    def calculate_real_image_statistics(self) : 
        self.real_inception_score = self.predict_to_inception_score(self.real_predict_softmax)[0]
        self.real_feature_np = self.real_feature.view(-1, 2048).numpy()
        self.real_mu, self.real_sigma = self.feature_to_mu_sig(self.real_feature_np)        
    
    def calculate_fake_image_statistics(self) : 
        self.fake_inception_score = self.predict_to_inception_score(self.fake_predict_softmax)[0]
        self.fake_feature_np = self.fake_feature.view(-1, 2048).numpy()
        self.fake_mu, self.fake_sigma = self.feature_to_mu_sig(self.fake_feature_np)
    
    def clear_fake(self) : 
        self.fake_images = None
        self.fake_predict_softmax = None
        self.fake_feature = None
        self.fake_mu, self.fake_sigma = None, None
        
    def calculate_generative_score(self):
        fid = self.calculate_frechet_distance(self.real_mu, self.real_sigma, self.fake_mu, self.fake_sigma)
        real_pick = np.random.permutation(self.real_feature_np)[:10000]
        fake_pick = np.random.permutation(self.fake_feature_np)[:10000]
        metrics = prdc.compute_prdc(real_features=real_pick, fake_features=fake_pick, nearest_k=5)
        metrics['fid'] = fid
        metrics['real_is'] = self.real_inception_score
        metrics['fake_is'] = self.fake_inception_score
        return metrics
    

    def analysis_softmax_and_feature(self, images):
        return self.inception_model(images)
