import sys
sys.path.append("..")

import os
import pathlib
import cv2
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser

import numpy as np
import torch
from scipy import linalg
from models.pose_autoencoder import EmbeddingNet

import pickle
import random

try:
    from tqdm import tqdm
except ImportError:
    # If tqdm is not available, provide a mock version of it
    def tqdm(x):
        return x


parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)

parser.add_argument('--batch_size', '-bs', type = int, default = 24)

parser.add_argument('--device', '-dev', type = int, default = 0)

parser.add_argument('--path_gt', '-pgt', type = str, default = "/home/rafael/masters/test_repro/test")

parser.add_argument('--path_generated', '-pg', type = str, default = "/home/rafael/masters/to_move/outputs_reprotest2")
parser.add_argument('--keypoint_type', '-kt', type = str, default = "ours")

parser.add_argument('--path_checkpoints', '-pck', type = str, default = "/home/rafael/masters/to_move/checkpoints_ae2/step_70300/autoencoder.pth")


def _read_instance_gt(root, instance_name):
    instance_fp = os.path.join(root, instance_name)
    with open(instance_fp, "rb") as handler:
        instance = pickle.load(handler)
    return instance

def _read_makeittalk_instance(root, instance_name):
    instance_fp = os.path.join(root, instance_name)
    instance = np.load(instance_fp)[:, :, :-1]
    return instance

def _read_ours_instance(root, instance_name):
    instance_fp = os.path.join(root, instance_name)
    instance = np.load(instance_fp)["arr_0"]
    return instance

def _read_pt_instance(baseline_root, instance_name):
    instance_fp = os.path.join(baseline_root, instance_name)
    instance = np.load(instance_fp)
    return instance

def draw_keypoints(kps):
    img = np.ones((256, 256, 3))*255

    for kp in kps:
        x = kp[0]
        y = kp[1]
        img = cv2.circle(img, (int(x*256), int(y*256)), 2, (0, 255 ,0), 2)
    return img

def get_activations(my_instances, model, batch_size=50, device='cpu'):
    """Calculates the activations of the latente layer for keypoint faces.

    Params:
    -- my_instances: List of keypoints instances [64,136]
    -- model       : Instance of autoencoder
    -- batch_size  : Batch size of instances for the model to process at once.
                     
    -- device      : Device to run calculations

    Returns:
    -- A numpy array of dimension (num instances, dims) that contains the
       latente layer activations of the given tensor when feeding autoencoder with the
       query tensor.
    """

    model.eval()

    pred_arr = np.empty((len(my_instances), 32))

    start_idx = 0
    
    for batch_start in range(0, len(my_instances), batch_size):  

        if start_idx%10 == 0:
            print ("Feeding " + str(start_idx))

        batch_final = (batch_start + batch_size) if (batch_start + batch_size) < len(my_instances) else len(my_instances)
        batch = np.asarray(my_instances[batch_start:batch_final])    
        batch_tensor = torch.from_numpy(batch).to(device=device, dtype=torch.float)

        with torch.no_grad():
            feat, a, b, recon_poses = model(batch_tensor)


        pred = feat.cpu().detach().numpy()

        pred_arr[start_idx:start_idx + pred.shape[0]] = pred

        start_idx = start_idx + pred.shape[0]

    return pred_arr


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 a layer of the
               inception net (like returned by the function 'get_predictions')
               for generated samples.
    -- mu2   : The sample mean over activations, precalculated on an
               representative data set.
    -- sigma1: The covariance matrix over activations for generated samples.
    -- sigma2: The covariance matrix over activations, precalculated on an
               representative 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
        print(msg)
        offset = np.eye(sigma1.shape[0]) * eps
        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

    
    if np.iscomplexobj(covmean):
        # Numerical error might give slight imaginary component
        #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 calculate_activation_statistics(my_instances, model, batch_size=50, device='cpu'):
    """Calculation of the statistics used by the FID.
    Params:
    -- files       : List of image files paths
    -- model       : Instance of inception model
    -- batch_size  : The images numpy array is split into batches with
                     batch size batch_size. A reasonable batch size
                     depends on the hardware.
    -- dims        : Dimensionality of features returned by Inception
    -- device      : Device to run calculations
    -- num_workers : Number of parallel dataloader workers

    Returns:
    -- mu    : The mean over samples of the activations of the pool_3 layer of
               the inception model.
    -- sigma : The covariance matrix of the activations of the pool_3 layer of
               the inception model.
    """
    act = get_activations(my_instances, model, batch_size, device)
    mu = np.mean(act, axis=0)
    sigma = np.cov(act, rowvar=False)
    return mu, sigma


def compute_statistics_of_path(path, model, batch_size, keypoint_type, device):

    my_instances = list()

    if keypoint_type == "gt":
        print ("gt")
    elif keypoint_type == "pt":
        print ("pt")
    elif keypoint_type == "mt":
        print ("mt")
    elif keypoint_type == "ours":
        print ("ours")
    elif keypoint_type == "nslp":
        print ("ours")
    else:
        print ("Keypoint type error")
        assert False


    # todo: format mt and pt keypoints
    for instance_name in sorted(os.listdir(path)):
        try:
            if keypoint_type == "gt": 
                my_instances.append(np.asarray(_read_instance_gt(path, instance_name)["kps"])[:, :-1, :].reshape(-1, 136))
            elif keypoint_type == "mt":
                my_instances.append(_read_makeittalk_instance(path, instance_name).reshape(-1, 136))
            elif keypoint_type == "pt":
                my_instances.append(_read_pt_instance(path, instance_name)[:,:-1].reshape(-1, 136))
            elif keypoint_type == "ours":
                my_instances.append(_read_ours_instance(path, instance_name).squeeze(0)[:, :-1, :].reshape(-1, 136))
            elif keypoint_type == "nslp":
                test = _read_instance_gt(path, instance_name)/256
                test = test.reshape(-1, 136).numpy()
                my_instances.append(test)
        except:
            print("found no existing instance, skipping")
            import traceback
            print(traceback.format_exc())

    m, s = calculate_activation_statistics(my_instances, model, batch_size, device)

    return m, s


def calculate_fid_given_paths(paths_gt, paths_g, batch_size, device, keypoint_type,path_checkpoints):

    """Calculates the FID of two paths"""
    if not os.path.exists(paths_gt):
        raise RuntimeError('Invalid path: %s' % paths_gt)

    if not os.path.exists(paths_g):
        raise RuntimeError('Invalid path: %s' % paths_g)

    

    # pose_dim = 68*2,n_frames

    model = EmbeddingNet(136, 64).to(device)

    try:
        model.load_state_dict(torch.load(path_checkpoints, map_location=torch.device("cpu")))
    except:
        import traceback
        print(traceback.format_exc())
        import pdb
        pdb.set_trace()


    model.eval()

    m1, s1 = compute_statistics_of_path(paths_gt, model, batch_size,
                                        "gt", device)

    m2, s2 = compute_statistics_of_path(paths_g, model, batch_size,
                                        keypoint_type, device)

    fid_value = calculate_frechet_distance(m1, s1, m2, s2)

    return fid_value


def main():

    args = parser.parse_args()
    SEED = 1234
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.backends.cudnn.deterministic = True
    #device = torch.device('cuda:{}'.format(args.device)) if args.device != -1 else torch.device('cuda')
    device = torch.device("cpu")

    
    fid_value = calculate_fid_given_paths(args.path_gt,
                                          args.path_generated,
                                          args.batch_size,
                                          device,
                                          args.keypoint_type,
                                          args.path_checkpoints)
    print('FID: ', fid_value)


if __name__ == '__main__':
    main()