#!/usr/bin/env python3
"""Calculates the Frechet Inception Distance (FID) to evalulate GANs

The FID metric calculates the distance between two distributions of images.
Typically, we have summary statistics (mean & covariance matrix) of one
of these distributions, while the 2nd distribution is given by a GAN.

When run as a stand-alone program, it compares the distribution of
images that are stored as PNG/JPEG at a specified location with a
distribution given by summary statistics (in pickle format).

The FID is calculated by assuming that X_1 and X_2 are the activations of
the pool_3 layer of the inception net for generated samples and real world
samples respectively.

See --help to see further details.

Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
of Tensorflow

Copyright 2018 Institute of Bioinformatics, JKU Linz

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import pathlib
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

import numpy as np
import torch
from scipy import linalg
from scipy.misc import imread
from torch.nn.functional import adaptive_avg_pool2d, adaptive_max_pool2d
from scipy import misc
import random
import re
from scipy.special import softmax
from shutil import copyfile
from domain_gap.sskmean.clustering.equal_groups import EqualGroupsKMeans
import glob
import os.path as osp

import numpy as np
from sklearn.cluster import KMeans
import time

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

from domain_gap.models.inception import InceptionV3

parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('path', type=str, nargs=2,
                    help=('Path to the generated images or '
                          'to .npz statistic files'))
parser.add_argument('--batch-size', type=int, default=50,
                    help='Batch size to use')
parser.add_argument('--dims', type=int, default=2048,
                    choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
                    help=('Dimensionality of Inception features to use. '
                          'By default, uses pool3 features'))
parser.add_argument('-c', '--gpu', default='3', type=str,
                    help='GPU to use (leave blank for CPU only)')

def make_square(image, max_dim = 512):
    max_dim = max(np.shape(image)[0], np.shape(image)[1])
    h, w = image.shape[:2]
    top_pad = (max_dim - h) // 2
    bottom_pad = max_dim - h - top_pad
    left_pad = (max_dim - w) // 2
    right_pad = max_dim - w - left_pad
    padding = [(top_pad, bottom_pad), (left_pad, right_pad), (0, 0)]
    image = np.pad(image, padding, mode='constant', constant_values=0)
    window = (top_pad, left_pad, h + top_pad, w + left_pad)
    return image

def get_activations(opt, files, model, batch_size=50, dims=8192,
                    cuda=False, verbose=False):
    """Calculates the activations of the pool_3 layer for all images.

    Params:
    -- files       : List of image files paths
    -- model       : Instance of inception model
    -- batch_size  : Batch size of images for the model to process at once.
                     Make sure that the number of samples is a multiple of
                     the batch size, otherwise some samples are ignored. This
                     behavior is retained to match the original FID score
                     implementation.
    -- dims        : Dimensionality of features returned by Inception
    -- cuda        : If set to True, use GPU
    -- verbose     : If set to True and parameter out_step is given, the number
                     of calculated batches is reported.
    Returns:
    -- A numpy array of dimension (num images, dims) that contains the
       activations of the given tensor when feeding inception with the
       query tensor.
    """
    model.eval()

    # if len(files) % batch_size != 0:
    #     print(('Warning: number of images is not a multiple of the '
    #            'batch size. Some samples are going to be ignored.'))
    if batch_size > len(files):
        print(('Warning: batch size is bigger than the data size. '
               'Setting batch size to data size'))
        batch_size = len(files)

    n_batches = len(files) // batch_size
    n_remainder=  len(files) % batch_size

    print('\rnumber of batches is %d' % n_batches),
    n_used_imgs = n_batches * batch_size

    pred_arr = np.empty((n_used_imgs+n_remainder, dims))
    if n_remainder!=0:
        n_batches=n_batches+1
    for i in range(n_batches):
        if verbose:
            print('\rPropagating batch %d/%d' % (i + 1, n_batches),
                  end='', flush=True)
        start = i * batch_size
        if n_remainder!=0 and i==n_batches-1:
          end = start + n_remainder
        else:
          end = start + batch_size

        images = np.array([misc.imresize( imread(str(f)).astype(np.float32), size=[64, 64]).astype(np.float32)
                           for f in files[start:end]])

        images = images.transpose((0, 3, 1, 2))
        images /= 255

        batch = torch.from_numpy(images).type(torch.FloatTensor)
        if cuda:
            batch = batch.cuda()

        if opt.FD_model == 'inception':
            pred = model(batch)[0]
            # If model output is not scalar, apply global spatial average pooling.
            # This happens if you choose a dimensionality not equal 2048.
            if pred.shape[2] != 1 or pred.shape[3] != 1:
                pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
        if opt.FD_model == 'posenet':
            pred = model(batch)
            # print (np.shape (pred))
            pred = adaptive_max_pool2d(pred, output_size=(1, 1))
        pred_arr[start:end] = pred.cpu().data.numpy().reshape(end - start, -1)
        print('\rPropagating batch %d/%d' % (i + 1, n_batches))

    if verbose:
        print(' done')

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

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

def calculate_activation_statistics(opt, files, model, batch_size=50,
                                    dims=8192, cuda=False, verbose=False):
    """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
    -- cuda        : If set to True, use GPU
    -- verbose     : If set to True and parameter out_step is given, the
                     number of calculated batches is reported.
    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(opt, files, model, batch_size, dims, cuda, verbose)
    mu = np.mean(act, axis=0)
    sigma = np.cov(act, rowvar=False)
    #eigen_vals, eigen_vecs= np.linalg.eig(sigma)
    #sum_eigen_val=eigen_vals.sum().real
    sum_eigen_val = (sigma.diagonal()).sum()
    return mu, sigma, sum_eigen_val

def _compute_statistics_of_path(opt, path, model, batch_size, dims, cuda):
    if path.endswith('.npz'):
        f = np.load(path)
        m, s = f['mu'][:], f['sigma'][:]
        f.close()
    else:
        path = pathlib.Path(path)
        files = list(path.glob('*.jpg')) + list(path.glob('*.png'))
        #random.shuffle(files)
        #files = files[:2000]
        m, s, sum_eigen_val = calculate_activation_statistics(opt, files, model, batch_size,
                                               dims, cuda)
    return m, s, sum_eigen_val

def get_id_path_of_data (dataset_id, paths):
    img_paths = []
    dataset_ids = []
    person_ids = []
    pattern = re.compile(r'([-\d]+)_c([-\d]+)')
    did = 0
    for sub_path in paths:
        sub_path = pathlib.Path(sub_path)
        files = list(sub_path.glob('*.jpg')) + list(sub_path.glob('*.png'))
        # files=glob.glob(osp.join(sub_path, '*.png'))+glob.glob(osp.join(sub_path, '*.jpg'))
        dataset_id_list = [dataset_id[did] for n in range(len(files))]
        dataset_ids.extend(dataset_id_list)
        img_paths.extend(files)
        did += 1
    dataset = []
    ii = 0
    for img_path in img_paths:
        pid, camid = map(int, pattern.search(str(img_path)).groups())
        # if pid == -1: continue  # junk images are just ignored
        camid -= 1  # index starts from 0
        dataid = dataset_ids[ii]
        person_ids.append(pid)
        dataset.append((img_path, pid, dataid))
        ii = ii + 1

    return img_paths, person_ids, dataset_ids, dataset

def clustering_sample(tpaths, dict, dataset_id, opt, result_dir, c_num, score_name, weight, n_num):
    """clustering the ids from different datasets and sampleing"""

    # preparing dataset
    paths = [dict[i]+'bounding_box_test' for i in dataset_id]
    img_paths,  person_ids,  dataset_ids, _  = get_id_path_of_data(dataset_id, paths)

    cuda = True
    for p in paths:
        if not os.path.exists(p):
            raise RuntimeError('Invalid path: %s' % p)

    if opt.FD_model == 'inception':
        dims = 2048
        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
        model = InceptionV3([block_idx])

    if cuda:
        model.cuda()
    batch_size=256

    # caculate the various, mu sigma of target ste
    print('=========== extracting feature of target traning set ===========')
    target_path = pathlib.Path(tpaths)
    files = list(target_path.glob('*.jpg')) + list(target_path.glob('*.png'))
    # random.shuffle(files)
    # files = files[:2000]
    target_feature = get_activations(opt, files, model, batch_size, dims, cuda, verbose=False)
    m1 = np.mean(target_feature, axis=0)
    s1 = np.cov(target_feature, rowvar=False)
    sum_eigen_val1 = (s1.diagonal()).sum()

    # extracter feature for data pool
    if not os.path.exists(result_dir + '/feature.npy'):
        print('=========== extracting feature of data pool ===========')
        feature = get_activations(opt, img_paths, model, batch_size, dims, cuda, verbose=False)
        if not os.path.isdir(result_dir):
            os.mkdir(result_dir)
        np.save(result_dir + '/feature.npy', feature)
    else:
        feature = np.load(result_dir + '/feature.npy')

    person_ids_array=np.array(person_ids)
    mean_feature_per_id=[]
    pid_per_id=[]
    did_per_id=[]

    # get mean fature of perid and the fid, various of per_id with the target
    if not os.path.exists(result_dir + '/mean_feature_per_id.npy'):
        for did in dataset_id:
           ind_of_set = np.argwhere(np.array(dataset_ids) == did).squeeze()
           dataset_feature = feature[ind_of_set]
           dataset_pid = person_ids_array[ind_of_set]
           pid_of_dataset=set(dataset_pid)
           for pid in pid_of_dataset:
              ind_of_pid = np.argwhere(np.array(dataset_pid) == pid).squeeze()
              feature_per_id = dataset_feature[ind_of_pid]
              id_ave_feature=feature_per_id.mean(0)
              mean_feature_per_id.append(id_ave_feature)
              pid_per_id.append(pid)
              did_per_id.append(did)
        np.save(result_dir+ '/mean_feature_per_id.npy',mean_feature_per_id)
        pid_did_fid_var = np.c_[np.array(pid_per_id), np.array(did_per_id)]
        np.save(result_dir+ '/pid_did_fid_var.npy', pid_did_fid_var)
    else:
       mean_feature_per_id=np.load(result_dir + '/mean_feature_per_id.npy')
       pid_did_fid_var = np.load(result_dir + '/pid_did_fid_var.npy')

    #remove 0 and -1
    ori_pid_per_id = pid_did_fid_var[:, 0]
    remove_ind=np.r_[np.argwhere(ori_pid_per_id == -1), np.argwhere(ori_pid_per_id == 0)].squeeze()

    new_pid_did_fid_var = np.delete(pid_did_fid_var, remove_ind, 0)
    new_mean_feature_per_id = np.delete(mean_feature_per_id, remove_ind, 0)

    print('\r=========== clustering the data pool ===========')
    pid_per_id = new_pid_did_fid_var[:,0]
    did_per_id = new_pid_did_fid_var[:,1]
    # clustering ids based on ids' mean feature
    if not os.path.exists(result_dir + '/label_cluster_'+str(c_num)+'.npy'):
        estimator = KMeans(n_clusters=c_num)
        estimator.fit(new_mean_feature_per_id)
        label_pred = estimator.labels_
        np.save(result_dir + '/label_cluster_'+str(c_num)+'.npy',label_pred)
    else:
        label_pred = np.load('sample_data/' + '/label_cluster_'+str(c_num)+'.npy')

    print('\r=========== caculating the fid and v_gap between T and C_k ===========')
    if not os.path.exists(result_dir + '/cluster_fid_var.npy'):
        cluster_feature=[]
        cluster_fid=[]
        cluster_mmd=[]
        cluster_var_gap=[]
        for k in tqdm(range(c_num)):
            # initializatn of the first seed cluster 0
            initial_pid=pid_per_id[label_pred==k]
            initial_did=did_per_id[label_pred==k]
            initial_feature = feature[(dataset_ids == initial_did[0]) & (person_ids_array == initial_pid[0])]
            for j in range(1,len(initial_pid)):
                current_feature=feature[(dataset_ids == initial_did[j]) & (person_ids_array == initial_pid[j])]
                initial_feature=np.r_[initial_feature, current_feature]
            cluster_feature.append(initial_feature)
            mu = np.mean(initial_feature, axis=0)
            sigma = np.cov(initial_feature, rowvar=False)
            # caculating various
            current_var_gap = np.abs((sigma.diagonal()).sum() - sum_eigen_val1)
            current_fid = calculate_frechet_distance(m1, s1, mu, sigma)
            # mmd_value = polynomial_mmd_averages(torch.from_numpy(initial_feature), torch.from_numpy(target_feature))
            # current_mmd=mmd_value[0].mean()
            cluster_fid.append(current_fid)
            # cluster_mmd.append(current_mmd)
            cluster_var_gap.append(current_var_gap)
        np.save(result_dir + '/cluster_fid_var.npy', np.c_[np.array(cluster_fid), np.array(cluster_var_gap)])
        #np.save(result_dir+'/cluster_fid_var.npy', np.c_[np.array(cluster_fid),np.array(cluster_var_gap)])
    else:
        cluster_fid_var=np.load(result_dir + '/cluster_fid_var.npy')
        cluster_fid=cluster_fid_var[:,0]
        cluster_var_gap=cluster_fid_var[:,1]

#    cluster_fid=cluster_mmd
#    calculatting softmax score
    cluster_fida=np.array(cluster_fid)
    cluster_var_gapa=np.array(cluster_var_gap)
    score_fid = softmax(-cluster_fida)
    score_var_gap = softmax(-cluster_var_gapa)
    if score_name == 'fid':
        sample_rate=score_fid
    elif score_name == 'var':
        sample_rate=score_var_gap
    else:
        sample_rate = score_fid* weight + score_var_gap * (1-weight)

    c_num_len = []
    id_score = []
    for kk in range(c_num):
        initial_pid = pid_per_id[label_pred == kk]
        c_num_len.append(len(initial_pid))
    for jj in range(len(label_pred)):
        id_score.append(sample_rate[label_pred[jj]] / c_num_len[label_pred[jj]])

    selected_data_ind = np.sort(np.random.choice(range(len(id_score)), n_num, p=id_score))
    sdid = did_per_id[selected_data_ind]
    spid = pid_per_id[selected_data_ind]
    data_dir = result_dir + '/proxy_set'
    if not os.path.isdir(data_dir):
        os.mkdir(data_dir)
    print('\r=========== building proxy set ===========')
    sampled_data=np.c_[sdid,spid]
    ii = dataset_build(dict, dataset_id, sampled_data, data_dir)
    print('finished')
    return sampled_data

def dataset_build(dict, dataset_id, sampled_data,result_dir):
    pattern = re.compile(r'([-\d]+)_c([-\d]+)')
    pid=sampled_data[:, 1]
    new_pid=np.arange(len(pid))+1
    did=sampled_data[:, 0]

    for ii in range(len(pid)):
        id= pid[ii]
        id_set=did[ii]
        new_id=new_pid[ii]
        # sample images
        gallery_data_path = dict[id_set]+ 'bounding_box_test'
        for root, dirs, files in os.walk(gallery_data_path, topdown=True):
            for name in files:
                current_id, _ = map(int, pattern.search(str(name)).groups())
                if not (name[-3:] == 'png' or name[-3:] == 'jpg'):
                    continue
                if int(current_id)!= id:
                    continue
                src_path = gallery_data_path + '/' + name
                dst_path = result_dir+ '/bounding_box_test'
                dstr_path = result_dir + '/bounding_box_train'
                if not os.path.isdir(dst_path):
                    os.mkdir(dst_path)
                    os.mkdir(dstr_path)
                    # one picture in train
                    copyfile(src_path, dstr_path + '/' + '{:04}'.format(new_id) + name[4:-3] + 'jpg')
                copyfile(src_path, dst_path + '/' + '{:04}'.format(new_id) + name[4:-3]+'jpg')
        query_data_path = gallery_data_path[0:-18]+ '/query'
        for root, dirs, files in os.walk(query_data_path, topdown=True):
            for name in files:
                if not (name[-3:] == 'png' or name[-3:] == 'jpg'):
                    continue
                if int(name[0:4]) != id:
                    continue
                src_path = query_data_path + '/' + name
                dst_path = result_dir+'/query'
                if not os.path.isdir(dst_path):
                    os.mkdir(dst_path)
                copyfile(src_path, dst_path + '/' + '{:04}'.format(new_id) + name[4:-3]+'jpg')

def calculate_fd_given_paths(paths, opt):
    """Calculates the FID of two paths"""

    cuda = True
    for p in paths:
        if not os.path.exists(p):
            raise RuntimeError('Invalid path: %s' % p)

    if opt.FD_model == 'inception':
        dims = 2048
        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
        model = InceptionV3([block_idx])

    if cuda:
        model.cuda()

    m1, s1, sum_eigen_val1 = _compute_statistics_of_path(opt, paths[0], model, 256,
                                         dims, cuda)

    npz_path = None
    if not paths[0].endswith(".npz"):
        if not paths[0].endswith('/'):
            npz_path = paths[0] + ".npz"
        else:
            npz_path = paths[0][:-1] + ".npz"
        np.savez(npz_path, mu = m1, sigma = s1)
    m2, s2, sum_eigen_val2 = _compute_statistics_of_path(opt, paths[1], model, 256,
                                         dims, cuda)

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

    return fd_value, npz_path, sum_eigen_val1, sum_eigen_val2

if __name__ == '__main__':
    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    fid_value,_, sum_eigen_val1, sum_eigen_val2 = calculate_fd_given_paths(args.path,
                                          args.batch_size,
                                          args.gpu != '',
                                          8192)
    #print (fid_value)


""" This file defines utility classes and functions for algorithms. """
import numpy as np

from gps.utility.general_utils import BundleType
from gps.algorithm.policy.lin_gauss_policy import LinearGaussianPolicy

class IterationData(BundleType):
    """ Collection of iteration variables. """
    def __init__(self):
        variables = {
            'sample_list': None,  # List of samples for the current iteration.
            'traj_info': None,  # Current TrajectoryInfo object.
            'pol_info': None,  # Current PolicyInfo object.
            'traj_distr': None,  # Initial trajectory distribution.
            'new_traj_distr': None, # Updated trajectory distribution.
            'cs': None,  # Sample costs of the current iteration.
            'step_mult': 1.0,  # KL step multiplier for the current iteration.
            'eta': 1.0,  # Dual variable used in LQR backward pass.
        }
        BundleType.__init__(self, variables)

class TrajectoryInfo(BundleType):
    """ Collection of trajectory-related variables. """
    def __init__(self):
        variables = {
            'dynamics': None,  # Dynamics object for the current iteration.
            'x0mu': None,  # Mean for the initial state, used by the dynamics.
            'x0sigma': None,  # Covariance for the initial state distribution.
            'cc': None,  # Cost estimate constant term.
            'cv': None,  # Cost estimate vector term.
            'Cm': None,  # Cost estimate matrix term.
            'last_kl_step': float('inf'),  # KL step of the previous iteration.
        }
        BundleType.__init__(self, variables)

class PolicyInfo(BundleType):
    """ Collection of policy-related variables. """
    def __init__(self, hyperparams):
        T, dU, dX = hyperparams['T'], hyperparams['dU'], hyperparams['dX']
        variables = {
            'lambda_k': np.zeros((T, dU)),  # Dual variables.
            'lambda_K': np.zeros((T, dU, dX)),  # Dual variables.
            'pol_wt': hyperparams['init_pol_wt'] * np.ones(T),  # Policy weight.
            'pol_mu': None,  # Mean of the current policy output.
            'pol_sig': None,  # Covariance of the current policy output.
            'pol_K': np.zeros((T, dU, dX)),  # Policy linearization.
            'pol_k': np.zeros((T, dU)),  # Policy linearization.
            'pol_S': np.zeros((T, dU, dU)),  # Policy linearization covariance.
            'chol_pol_S': np.zeros((T, dU, dU)),  # Cholesky decomp of covar.
            'prev_kl': None,  # Previous KL divergence.
            'init_kl': None,  # The initial KL divergence, before the iteration.
            'policy_samples': [],  # List of current policy samples.
            'policy_prior': None,  # Current prior for policy linearization.
        }
        BundleType.__init__(self, variables)

    def traj_distr(self):
        """ Create a trajectory distribution object from policy info. """
        T, dU, dX = self.pol_K.shape
        # Compute inverse policy covariances.
        inv_pol_S = np.empty_like(self.chol_pol_S)
        for t in range(T):
            inv_pol_S[t, :, :] = np.linalg.solve(
                self.chol_pol_S[t, :, :],
                np.linalg.solve(self.chol_pol_S[t, :, :].T, np.eye(dU))
            )
        return LinearGaussianPolicy(self.pol_K, self.pol_k, self.pol_S,
                self.chol_pol_S, inv_pol_S)

def estimate_moments(X, mu, covar):
    """ Estimate the moments for a given linearized policy. """
    N, T, dX = X.shape
    dU = mu.shape[-1]
    if len(covar.shape) == 3:
        covar = np.tile(covar, [N, 1, 1, 1])
    Xmu = np.concatenate([X, mu], axis=2)
    ev = np.mean(Xmu, axis=0)
    em = np.zeros((N, T, dX+dU, dX+dU))
    pad1 = np.zeros((dX, dX+dU))
    pad2 = np.zeros((dU, dX))
    for n in range(N):
        for t in range(T):
            covar_pad = np.vstack([pad1, np.hstack([pad2, covar[n, t, :, :]])])
            em[n, t, :, :] = np.outer(Xmu[n, t, :], Xmu[n, t, :]) + covar_pad
    return ev, em

def gauss_fit_joint_prior(pts, mu0, Phi, m, n0, dwts, dX, dU, sig_reg):
    """ Perform Gaussian fit to data with a prior. """
    # Build weights matrix.
    D = np.diag(dwts)
    # Compute empirical mean and covariance.
    mun = np.sum((pts.T * dwts).T, axis=0)
    diff = pts - mun
    empsig = diff.T.dot(D).dot(diff)
    empsig = 0.5 * (empsig + empsig.T)
    # MAP estimate of joint distribution.
    N = dwts.shape[0]
    mu = mun
    sigma = (N * empsig + Phi + (N * m) / (N + m) *
             np.outer(mun - mu0, mun - mu0)) / (N + n0)
    sigma = 0.5 * (sigma + sigma.T)
    # Add sigma regularization.
    sigma += sig_reg
    # Conditioning to get dynamics.
    fd = np.linalg.solve(sigma[:dX, :dX], sigma[:dX, dX:dX+dU]).T
    fc = mu[dX:dX+dU] - fd.dot(mu[:dX])
    dynsig = sigma[dX:dX+dU, dX:dX+dU] - fd.dot(sigma[:dX, :dX]).dot(fd.T)
    dynsig = 0.5 * (dynsig + dynsig.T)
    return fd, fc, dynsig


"""
Agentpy Output Module
Content: DataDict class for output data
"""

import pandas as pd
import os
from os import listdir, makedirs
from os.path import getmtime, join
from SALib.analyze import sobol
from .tools import AttrDict, make_list, AgentpyError
import json
import numpy as np

class NpEncoder(json.JSONEncoder):
    """ Adds support for numpy number formats to json. """
    # By Jie Yang https://stackoverflow.com/a/57915246
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)

def _last_exp_id(name, path):
    """ Identifies existing experiment data and return highest id. """

    exp_id = 0
    output_dirs = listdir(path)
    exp_dirs = [s for s in output_dirs if name in s]
    if exp_dirs:
        ids = [int(s.split('_')[-1]) for s in exp_dirs]
        exp_id = max(ids)
    return exp_id

# TODO Create DataSubDict without methods
class DataDict(AttrDict):
    """ Nested dictionary for output data of simulations.
    Items can be accessed like attributes.
    Attributes can differ from the standard ones listed below.

    Attributes:
        info (dict):
            Metadata of the simulation.
        parameters (DataDict):
            Simulation parameters.
        variables (DataDict):
            Recorded variables, separatedper object type.
        reporters (pandas.DataFrame):
            Reported outcomes of the simulation.
        sensitivity (DataDict):
            Sensitivity data, if calculated.
    """

    def __repr__(self, indent=False):
        rep = ""
        if not indent:
            rep += "DataDict {"
        i = '    ' if indent else ''
        for k, v in self.items():
            rep += f"\n{i}'{k}': "
            if isinstance(v, (int, float, np.integer, np.floating)):
                rep += f"{v} {type(v)}"
            elif isinstance(v, str):
                x0 = f"(length {len(v)})"
                x = f"...' {x0}" if len(v) > 20 else "'"
                rep += f"'{v[:30]}{x} {type(v)}"
            elif isinstance(v, pd.DataFrame):
                lv = len(list(v.columns))
                rv = len(list(v.index))
                rep += f"DataFrame with {lv} " \
                       f"variable{'s' if lv != 1 else ''} " \
                       f"and {rv} row{'s' if rv != 1 else ''}"
            elif isinstance(v, DataDict):
                rep += f"{v.__repr__(indent=True)}"
            elif isinstance(v, dict):
                lv = len(list(v.keys()))
                rep += f"Dictionary with {lv} key{'s' if lv != 1 else ''}"
            elif isinstance(v, list):
                lv = len(v)
                rep += f"List with {lv} entr{'ies' if lv != 1 else 'y'}"
            else:
                rep += f"Object of type {type(v)}"
        if not indent:
            rep += "\n}"
        return rep

    def _short_repr(self):
        len_ = len(self.keys())
        return f"DataDict {{{len_} entr{'y' if len_ == 1 else 'ies'}}}"

    def __eq__(self, other):
        """ Check equivalence of two DataDicts."""
        if not isinstance(other, DataDict):
            return False
        for key, item in self.items():
            if key not in other:
                return False
            if isinstance(item, pd.DataFrame):
                if not self[key].equals(other[key]):
                    return False
            elif not self[key] == other[key]:
                return False
        return True

    def __ne__(self, other):
        return not self.__eq__(other)

    # Data analysis --------------------------------------------------------- #

    @staticmethod
    def _sobol_set_df_index(df, p_keys, reporter):
        df['parameter'] = p_keys
        df['reporter'] = reporter
        df.set_index(['reporter', 'parameter'], inplace=True)

    def calc_sobol(self, reporters=None, **kwargs):
        """ Calculates Sobol Sensitivity Indices
        using :func:`SALib.analyze.sobol.analyze`.
        Data must be from an :class:`Experiment` with a :class:`Sample`
        that was generated with the method 'saltelli'.
        If the experiment had more than one iteration,
        the mean value between iterations will be taken.

        Arguments:
            reporters (str or list of str, optional): The reporters that should
                be used for the analysis. If none are passed, all are used.
            **kwargs: Will be forwarded to :func:`SALib.analyze.sobol.analyze`.

        Returns:
            DataDict: The DataDict itself with an added category 'sensitivity'.
        """

        if not self.parameters.log['type'] == 'saltelli':
            raise AgentpyError("Sampling method must be 'saltelli'.")
        if self.info['iterations'] == 1:
            reporters_df = self.reporters
        else:
            reporters_df = self.reporters.groupby('sample_id').mean()

        # STEP 1 - Load salib problem from parameter log
        param_ranges_salib = self.parameters.log['salib_problem']
        calc_second_order = self.parameters.log['calc_second_order']

        # STEP 2 - Calculate Sobol Sensitivity Indices
        if reporters is None:
            reporters = reporters_df.columns
        if isinstance(reporters, str):
            reporters = [reporters]
        p_keys = self._combine_pars(sample=True, constants=False).keys()
        dfs_list = [[] for _ in range(4 if calc_second_order else 2)]

        for reporter in reporters:
            y = np.array(reporters_df[reporter])
            si = sobol.analyze(param_ranges_salib, y, calc_second_order, **kwargs)

            # Make dataframes out of S1 and ST sensitivities
            keyss = [['S1', 'ST'], ['S1_conf', 'ST_conf']]
            for keys, dfs in zip(keyss, dfs_list[0:2]):
                s = {k[0:2]: v for k, v in si.items() if k in keys}
                df = pd.DataFrame(s)
                self._sobol_set_df_index(df, p_keys, reporter)
                dfs.append(df)

            # Make dataframes out S2 sensitivities
            if calc_second_order:
                for key, dfs in zip(['S2', 'S2_conf'], dfs_list[2:4]):
                    df = pd.DataFrame(si[key])
                    self._sobol_set_df_index(df, p_keys, reporter)
                    dfs.append(df)

        # Combine dataframes for each reporter
        self['sensitivity'] = sdict = DataDict()
        sdict['sobol'] = pd.concat(dfs_list[0])
        sdict['sobol_conf'] = pd.concat(dfs_list[1])

        if calc_second_order:
            # Add Second-Order to self
            dfs_si = [sdict['sobol'], pd.concat(dfs_list[2])]
            dfs_si_conf = [sdict['sobol_conf'], pd.concat(dfs_list[3])]
            sdict['sobol'] = pd.concat(dfs_si, axis=1)
            sdict['sobol_conf'] = pd.concat(dfs_si_conf, axis=1)

            # Create Multi-Index for Columns
            arrays = [["S1", "ST"] + ["S2"] * len(p_keys), [""] * 2 + list(p_keys)]
            tuples = list(zip(*arrays))
            index = pd.MultiIndex.from_tuples(tuples, names=["order", "parameter"])
            sdict['sobol'].columns = index
            sdict['sobol_conf'].columns = index.copy()

        return self

    # Data arrangement ------------------------------------------------------ #

    def _combine_vars(self, obj_types=True, var_keys=True):
        """ Returns pandas dataframe with combined variables """

        # Retrieve variables
        if 'variables' in self:
            vs = self['variables']
        else:
            return None

        if len(vs.keys()) == 1:
            return list(vs.values())[0]  # Return df if vs has only one entry
        elif isinstance(vs, DataDict):
            df_dict = dict(vs)  # Convert to dict if vs is DataDict

        # Remove dataframes that don't include any of the selected var_keys
        if var_keys is not True:
            df_dict = {k: v for k, v in df_dict.items()
                       if any(x in v.columns for x in make_list(var_keys))}

        # Select object types
        if obj_types is not True:
            df_dict = {k: v for k, v in df_dict.items()
                       if k in make_list(obj_types)}

        # Add 'obj_id' before 't' for model df
        model_type = self.info['model_type']
        if model_type in list(df_dict.keys()):
            df = df_dict[model_type]
            df['obj_id'] = 0
            indexes = list(df.index.names)
            indexes.insert(-1, 'obj_id')
            df = df.reset_index()
            df = df.set_index(indexes)
            df_dict[model_type] = df

        # Return none if empty
        if df_dict == {}:
            return None

        # Create dataframe
        df = pd.concat(df_dict)  # Dict keys (obj_type) will be added to index
        df.index = df.index.set_names('obj_type', level=0)  # Rename new index

        # Select var_keys
        if var_keys is not True:
            # make_list prevents conversion to pd.Series for single value
            df = df[make_list(var_keys)]

        return df

    def _dict_pars_to_df(self, dict_pars):
        n = self.info['sample_size'] if 'sample_size' in self.info else 1
        d = {k: [v] * n for k, v in dict_pars.items()}
        i = pd.Index(list(range(n)), name='sample_id')
        return pd.DataFrame(d, index=i)

    def _combine_pars(self, sample=True, constants=True):
        """ Returns pandas dataframe with parameters and sample_id """
        # Cancel if there are no parameters
        if 'parameters' not in self:
            return None
        dfp = pd.DataFrame()
        if sample and 'sample' in self.parameters:
            dfp = self.parameters.sample.copy()
            if constants and 'constants' in self.parameters:
                for k, v in self.parameters.constants.items():
                    dfp[k] = v
        elif constants and 'constants' in self.parameters:
            dfp = self._dict_pars_to_df(self.parameters.constants)
        # Cancel if no parameters have been selected
        if dfp is None or dfp.empty is True:
            return None
        return dfp

    def arrange(self, variables=False, reporters=False, parameters=False,
                constants=False, obj_types=True, index=False):
        """ Combines and/or filters data based on passed arguments.

        Arguments:
            variables (bool or str or list of str, optional):
                Key or list of keys of variables to include in the dataframe.
                If True, all available variables are selected.
                If False (default), no variables are selected.
            reporters (bool or str or list of str, optional):
                Key or list of keys of reporters to include in the dataframe.
                If True, all available reporters are selected.
                If False (default), no reporters are selected.
            parameters (bool or str or list of str, optional):
                Key or list of keys of parameters to include in the dataframe.
                If True, all non-constant parameters are selected.
                If False (default), no parameters are selected.
            constants (bool, optional):
                Include constants if 'parameters' is True (default False).
            obj_types (str or list of str, optional):
                Agent and/or environment types to include in the dataframe.
                If True (default), all objects are selected.
                If False, no objects are selected.
            index (bool, optional):
                Whether to keep original multi-index structure (default False).

        Returns:
            pandas.DataFrame: The newly arranged dataframe.
        """

        dfv = dfm = dfp = df = None

        # Step 1: Variables
        if variables is not False:
            dfv = self._combine_vars(obj_types, variables)

        # Step 2: Measures
        if reporters is not False:
            dfm = self.reporters
            if reporters is not True:  # Select reporter keys
                # make_list prevents conversion to pd.Series for single value
                dfm = dfm[make_list(reporters)]

        # Step 3: Parameters
        if parameters is True:
            dfp = self._combine_pars(constants=constants)
        elif parameters is not False:
            dfp = self._combine_pars()
            dfp = dfp[make_list(parameters)]

        # Step 4: Combine dataframes
        if dfv is not None and dfm is not None:
            # Combine variables & measures
            index_keys = dfv.index.names
            dfm = dfm.reset_index()
            dfv = dfv.reset_index()
            df = pd.concat([dfm, dfv])
            df = df.set_index(index_keys)
        elif dfv is not None:
            df = dfv
        elif dfm is not None:
            df = dfm
        if dfp is not None:
            if df is None:
                df = dfp
            else:  # Combine df with parameters
                if df is not None and isinstance(df.index, pd.MultiIndex):
                    dfp = dfp.reindex(df.index, level='sample_id')
                df = pd.concat([df, dfp], axis=1)

        if df is None:
            return pd.DataFrame()

        # Step 6: Reset index
        if not index:
            df = df.reset_index()

        return df

    def arrange_reporters(self):
        """ Common use case of :obj:`DataDict.arrange`
        with `reporters=True` and `parameters=True`. """
        return self.arrange(variables=False, reporters=True, parameters=True)

    def arrange_variables(self):
        """ Common use case of :obj:`DataDict.arrange`
        with `variables=True` and `parameters=True`. """
        return self.arrange(variables=True, reporters=False, parameters=True)

    # Saving and loading data ----------------------------------------------- #

    def save(self, exp_name=None, exp_id=None, path='ap_output', display=True):
        """ Writes data to directory `{path}/{exp_name}_{exp_id}/`.

        Works only for entries that are of type :class:`DataDict`,
        :class:`pandas.DataFrame`, or serializable with JSON
        (int, float, str, dict, list). Numpy objects will be converted
        to standard objects, if possible.

        Arguments:
            exp_name (str, optional): Name of the experiment to be saved.
                If none is passed, `self.info['model_type']` is used.
            exp_id (int, optional): Number of the experiment.
                Note that passing an existing id can overwrite existing data.
                If none is passed, a new id is generated.
            path (str, optional): Target directory (default 'ap_output').
            display (bool, optional): Display saving progress (default True).
        """

        # Create output directory if it doesn't exist
        if path not in listdir():
            makedirs(path)

        # Set exp_name
        if exp_name is None:
            if 'info' in self and 'model_type' in self.info:
                exp_name = self.info['model_type']
            else:
                exp_name = 'Unnamed'

        exp_name = exp_name.replace(" ", "_")

        # Set exp_id
        if exp_id is None:
            exp_id = _last_exp_id(exp_name, path) + 1

        # Create new directory for output
        path = f'{path}/{exp_name}_{exp_id}'
        makedirs(path)

        # Save experiment data
        for key, output in self.items():

            if isinstance(output, pd.DataFrame):
                output.to_csv(f'{path}/{key}.csv')

            elif isinstance(output, DataDict):
                for k, o in output.items():

                    if isinstance(o, pd.DataFrame):
                        o.to_csv(f'{path}/{key}_{k}.csv')
                    elif isinstance(o, dict):
                        with open(f'{path}/{key}_{k}.json', 'w') as fp:
                            json.dump(o, fp, cls=NpEncoder)

            else:  # Use JSON for other object types
                try:
                    with open(f'{path}/{key}.json', 'w') as fp:
                        json.dump(output, fp, cls=NpEncoder)
                except TypeError as e:
                    print(f"Warning: Object '{key}' could not be saved. "
                          f"(Reason: {e})")
                    os.remove(f'{path}/{key}.json')

            # TODO Support grids & graphs
            # elif t == nx.Graph:
            #    nx.write_graphml(output, f'{path}/{key}.graphml')

        if display:
            print(f"Data saved to {path}")

    def _load(self, exp_name=None, exp_id=None,
              path='ap_output', display=True):

        def load_file(path, file, display):
            if display:
                print(f'Loading {file} - ', end='')
            i_cols = ['sample_id', 'iteration', 'obj_id', 't']
            ext = file.split(".")[-1]
            path = path + file
            try:
                if ext == 'csv':
                    obj = pd.read_csv(path) # Convert .csv into DataFrane
                    index = [i for i in i_cols if i in obj.columns]
                    if index:  # Set potential index columns
                        obj = obj.set_index(index)
                elif ext == 'json':
                    # Convert .json with json decoder
                    with open(path, 'r') as fp:
                        obj = json.load(fp)
                    # Convert dict to AttrDict
                    if isinstance(obj, dict):
                        obj = AttrDict(obj)
                # TODO Support grids & graphs
                # elif ext == 'graphml':
                #    self[key] = nx.read_graphml(path)
                else:
                    raise ValueError(f"File type '{ext}' not supported")
                if display:
                    print('Successful')
                return obj
            except Exception as e:
                print(f'Error: {e}')

        # Prepare for loading
        if exp_name is None:
            # Choose latest modified experiment
            exp_names = listdir(path)
            paths = [join(path, d) for d in exp_names]
            latest_exp = exp_names[paths.index(max(paths, key=getmtime))]
            exp_name = latest_exp.rsplit('_', 1)[0]

        exp_name = exp_name.replace(" ", "_")
        if not exp_id:
            exp_id = _last_exp_id(exp_name, path)
            if exp_id == 0:
                raise FileNotFoundError(f"No experiment found with "
                                        f"name '{exp_name}' in path '{path}'")
        path = f'{path}/{exp_name}_{exp_id}/'
        if display:
            print(f'Loading from directory {path}')

        # Loading data
        for file in listdir(path):
            if 'variables_' in file:
                if 'variables' not in self:
                    self['variables'] = DataDict()
                ext = file.split(".")[-1]
                key = file[:-(len(ext) + 1)].replace('variables_', '')
                self['variables'][key] = load_file(path, file, display)
            elif 'parameters_' in file:
                ext = file.split(".")[-1]
                key = file[:-(len(ext) + 1)].replace('parameters_', '')
                if 'parameters' not in self:
                    self['parameters'] = DataDict()
                self['parameters'][key] = load_file(path, file, display)
            else:
                ext = file.split(".")[-1]
                key = file[:-(len(ext) + 1)]
                self[key] = load_file(path, file, display)
        return self

    @classmethod
    def load(cls, exp_name=None, exp_id=None, path='ap_output', display=True):
        """ Reads data from directory `{path}/{exp_name}_{exp_id}/`.

            Arguments:
                exp_name (str, optional): Experiment name.
                    If none is passed, the most recent experiment is chosen.
                exp_id (int, optional): Id number of the experiment.
                    If none is passed, the highest available id used.
                path (str, optional): Target directory (default 'ap_output').
                display (bool, optional): Display loading progress (default True).

            Returns:
                DataDict: The loaded data from the chosen experiment.
        """
        return cls()._load(exp_name, exp_id, path, display)


import os
import sys
import numpy as np

from . import constants as cc
#from . import camera as cam
from . import preliminary_computations as precomp
from .local_conditions import sky_countrate
#from . import set_object as obj
from . import optics as opt
from . import photometry as phot
#from .InfoStore import add
from . import utils
#######################################

def etc_computation(info_dict):
    """ Compute either the SNR, the total exposure time or the magnitude in function of the 2 others

    Parameters
    ----------
    info_dict: dictionary
               contains all relevant information

    wavelength : array
             wavelengths in angstrom

    Returns
    ---------
    SNR: float
         Signal to noise ratio

    mag: float
         magnitude reached

    tot_exp_time: float
                  total exposure time in seconds

    """
    # display result
    verbose = info_dict['verbose']

    etc_type = info_dict['etc_type']
    SNR = info_dict['SNR']
    tot_exp_time = info_dict['Nexp']*info_dict['exptime']
    info_dict['total_exposure_time']=tot_exp_time
    Nexp = info_dict['Nexp']
    # Detector Integration Time in seconds
    if etc_type == 'snr' or etc_type == 'mag':
         DIT = info_dict['exptime'] - info_dict['T_dithering']

    #Display
    if verbose == True:
         print ('\nInformation about Passband:')
         print ('----------------------------')
         print ('Cut_on: %.f angstroms' % info_dict['Passband_cuton'])
         print ('Effective wavelength: %.f angstroms' % info_dict['effWavelength'])
         print ('Cut_off: %.f angstroms' % info_dict['Passband_cutoff'])

         print ('\nAirmass: %.2f' % info_dict['airmass'])
         print ('\nSeeing: %.2f' % info_dict['seeing_los_arcsec'])

    if info_dict['detailed_trans']==1:
         # Computes mean transmission of each components for the given passband
         mean_trans_tel = utils.mean_efficiency_passband(info_dict,opt.telescope_efficiency(info_dict))
         mean_trans_inst = utils.mean_efficiency_passband(info_dict,opt.instrument_channel_efficiency(info_dict)*info_dict['Trans_filter']*info_dict['camera_efficiency'])
         #mean_trans_optics = mean_trans_tel*mean_trans_inst
         mean_trans_filter = utils.mean_efficiency_passband(info_dict,info_dict['Trans_filter'])
         mean_eta_cam = utils.mean_efficiency_passband(info_dict,info_dict['camera_efficiency'])
         mean_trans_optics = mean_trans_tel*mean_trans_inst
         mean_trans_atm = utils.mean_efficiency_passband(info_dict,info_dict['Trans_atmosphere'])
         #mean_trans_system = mean_trans_optics*mean_trans_filter*mean_trans_atm*mean_eta_cam
         #mean_trans_system = mean_trans_optics*mean_eta_cam
         mean_trans_system = mean_trans_optics
         info_dict['trans_mean_tel']=mean_trans_tel
         info_dict['trans_mean_inst']=mean_trans_inst
         info_dict['trans_mean_optics']=mean_trans_optics
         info_dict['trans_mean_filter']=mean_trans_filter
         info_dict['trans_mean_atm']=mean_trans_atm
         info_dict['trans_mean_cam']=mean_eta_cam
         info_dict['trans_mean_system']=mean_trans_system

         if verbose == True:
              print ('\nMEAN EFFICENCIES:')
              print ('------------------')
              print ('Obscuration: %.3f' % (1.-info_dict['obstruction']))
              print ('Telescope: %.3f (+obs: %.3f)' % (mean_trans_tel,mean_trans_tel*(1.-info_dict['obstruction'])))
              print ('Instrument: %.3f' % mean_trans_inst)
              print ('Optics (tel+inst): %.3f  (+obs: %.3f)' % (mean_trans_optics,mean_trans_optics*(1.-info_dict['obstruction'])))
              print ('Filter: %.3f' % mean_trans_filter)
              print ('Atmosphere: %.3f' % mean_trans_atm)
              print ('Camera: %.3f' % mean_eta_cam)
              print ('System: %.3f (+obs: %.3f)\n' % (mean_trans_system,mean_trans_system*(1-info_dict['obstruction'])))

              #print ('Telescope alan: %.3f' % utils.mean_efficiency_passband(info_dict,opt.telescope_efficiency_alan(info_dict)))
              #print ('Instrument alan: %.3f' % utils.mean_efficiency_passband(info_dict,opt.instrument_channel_efficiency_alan(info_dict)*info_dict['Trans_filter']*info_dict['camera_efficiency']))
              #print ('System alan (+atm) (+obs+atm): %.3f (%.3f) (%.3f)' % (utils.mean_efficiency_passband(info_dict,opt.telescope_efficiency_alan(info_dict)*opt.instrument_channel_efficiency_alan(info_dict)*info_dict['Trans_filter']*info_dict['camera_efficiency']),utils.mean_efficiency_passband(info_dict,opt.telescope_efficiency_alan(info_dict)*opt.instrument_channel_efficiency_alan(info_dict)*info_dict['Trans_filter']*info_dict['camera_efficiency']*info_dict['Trans_atmosphere']),utils.mean_efficiency_passband(info_dict,opt.telescope_efficiency_alan(info_dict)*opt.instrument_channel_efficiency_alan(info_dict)*info_dict['Trans_filter']*info_dict['camera_efficiency']*info_dict['Trans_atmosphere'])*(1.-info_dict['obstruction'])))

    elif info_dict['detailed_trans'] == 0:
         mean_eta_cam = utils.mean_efficiency_passband(info_dict,info_dict['camera_efficiency'])
         mean_eta_optics=utils.mean_efficiency_passband(info_dict,phot.set_filter(info_dict))
         mean_trans_system =  mean_eta_cam * mean_eta_optics
         info_dict['trans_mean_system']=mean_trans_system
         if verbose == True:
              print ('\nMEAN EFFICENCIES:')
              print ('------------------')
              print ('Obscuration: %.3f' % (1.-info_dict['obstruction']))
              print ('System: %.2f (+obs: %.3f)\n' % (mean_trans_system,mean_trans_system*(1-info_dict['obstruction'])))

    # Number of pixels covering 1.35*FWHM of the PSF
    npix = info_dict['npix']
    # Factor when estimating the Noise from other images
    factor_ima = precomp.factor_images_averaged(info_dict)
    # Fraction of light in the brightest pixel
    f_pix = precomp.Normalisation_factor(info_dict,True)
    # Fraction of light in the PSF
    f_PSF = precomp.Normalisation_factor(info_dict,False)

    info_dict['factor_ima']=factor_ima
    info_dict['f_pix']=f_pix
    info_dict['f_PSF']=f_PSF
    #----------------------------------------------------------------------------------------

    # Background Noise countrate in e-/s/px
    #---------------------------------------
    info_dict = sky_countrate(info_dict)# e-/s/px
    BN = info_dict['Sky_CountRate']
    #print ('Sky countrate: %.2f (e-/px/s)' % BN)
    # Thermic signal (electrons/s/pixel)  <--> Dark current
    #------------------------------------------------------
    DC = info_dict['cameras'][info_dict['channel']]['DC']
    # Digitization noise   (e-/pixel)
    #----------------------------------
    # Converter analog to digital noise of 1/2 ADU (electrons/pixel)
    DigN = info_dict['dig_noise']

    # Readout noise (e-/pixel)
    #--------------------------
    RN = info_dict['cameras'][info_dict['channel']]['RN']

    # Instrument background (e-/s/pix)
    inst_bg=info_dict['Instrument_bg']

    # Object
    #---------
    # Count rate of the object in e-/s
    if etc_type == 'snr' or etc_type == 'time':
         CR, fph = info_dict['Object_fes'],info_dict['Object_fph']  # e-/s

    # Zeropoint
    #-----------
    ZP = info_dict['zeropoint']

    #Add some info in info_cit
    if verbose == True:
         print ('Zeropoint: %.2f (%s mag)' % (ZP,info_dict['photometry_system']))
    #-----------------------------------------------------------------------------------------

    # Compute the SNR
    #------------------
    if etc_type == 'snr':
         # In the case of a given magnitude to reach, the fraction of flux
         # we kept should be included so that the given mag corresponds to the
         # measured count rate.
         if info_dict['object_type'] == 'magnitude':
              CR = CR #/ f_PSF
              CR_pix = CR #/ f_pix
         else: CR_pix=CR

         # Peak SNR (Object signal/noise at the brightest pixel)
         SNR_pix = np.sqrt(Nexp) * CR_pix * f_pix * DIT / np.sqrt(CR_pix * f_pix * DIT + factor_ima * ((RN**2. + DigN**2.) + DIT * ( DC + BN + inst_bg)))

         # Total integrated noise over npix (electrons/area)
         SNR = np.sqrt(Nexp) * CR * f_PSF * DIT / np.sqrt(CR * f_PSF * DIT + factor_ima * npix*((RN**2. + DigN**2.) + DIT * ( DC + BN + inst_bg)))

         if info_dict['object_type'] == 'magnitude':
              #mag = ZP - 2.5*np.log10(CR*f_PSF)
              mag = info_dict['object_magnitude']
              mag_pix = ZP - 2.5*np.log10(CR_pix*f_pix)
              mag_pix = info_dict['object_magnitude']
         else:
              mag = ZP - 2.5*np.log10(CR)
              mag_pix = ZP - 2.5*np.log10(CR_pix)

         Ftot_el=CR*f_PSF*DIT#*np.sqrt(Nexp)
         Ftot_el_pix=CR_pix*f_pix*DIT#*np.sqrt(Nexp)
         DIT_pix=DIT

         info_dict['SNR']=SNR
         info_dict['SNR_pix']=SNR_pix
         info_dict['mag_pix']=mag_pix
         info_dict['Ftot_el_pix']=Ftot_el_pix
         info_dict['Ftot_el']=Ftot_el
         info_dict['DIT_pix']=DIT_pix

         if verbose == True:
              print ('\n\nA magnitude (%s system) of %.2f in %s band within a total exposure time of %.2f seconds splited in %d exposure(s), implies a total SNR of :\n' %(info_dict['photometry_system'],mag,info_dict['filter_band'],DIT*Nexp,Nexp))
              #print ('\t - Peak SNR at the brightest pixel: %.2f \n' % SNR_pix)
              print ('\t - Integrated SNR over %d pixels: %.2f' % (npix, SNR))
              print ('\n\nA magnitude (%s system) of %.2f in %s band within a total exposure time of %.2f seconds splited in %d exposure(s), implies a SNR for the central pixel of of :\n\n' %(info_dict['photometry_system'],mag_pix,info_dict['filter_band'],DIT_pix*Nexp,Nexp))
              print ('\t - SNR of the central pixel: %.2f \n\n' %  SNR_pix)
    #------------------------------------------------------------------------------------------

    # Compute the total exposure time
    #--------------------------------
    elif etc_type == 'time':

         # In the case of a given magnitude to reach, the fraction of flux
         # we kept should be included so that the given mag corresponds to the
         # measured count rate.
         if info_dict['object_type'] == 'magnitude': CR = CR #/ f_PSF

         # Integrated over Npixels     (solve 2nd degree equation)
         if Nexp > 1:
              SNR_1 = SNR / np.sqrt(Nexp)
         else:
              SNR_1 = SNR
         A_sys = -(CR * f_PSF)**2.
         B_sys = SNR_1**2.*(CR * f_PSF + factor_ima * npix*(DC + BN + inst_bg))
         C_sys = SNR_1**2.*factor_ima * npix*(RN**2. + DigN**2.)

         delta = B_sys*B_sys - 4.*A_sys*C_sys

         DIT = (-B_sys -np.sqrt(delta)) / (2.*A_sys)
         #mag = -2.5*np.log10(obj_janskys/3631)

         if info_dict['object_type'] == 'magnitude': mag = ZP - 2.5*np.log10(CR)#*f_PSF)
         else: mag = ZP - 2.5*np.log10(CR)

         Ftot_el=CR*f_PSF*DIT#*np.sqrt(Nexp)

         #Brightest pixel
         if info_dict['object_type'] == 'magnitude': CR_pix = CR #/ f_pix
         else: CR_pix=CR

         A_sys = -(CR_pix * f_pix)**2.
         B_sys = SNR_1**2.*(CR_pix * f_pix + factor_ima *(DC + BN + inst_bg))
         C_sys = SNR_1**2.*factor_ima *(RN**2. + DigN**2.)

         delta = B_sys*B_sys - 4.*A_sys*C_sys

         DIT_pix = (-B_sys -np.sqrt(delta)) / (2.*A_sys)

         if info_dict['object_type'] == 'magnitude': mag_pix = ZP - 2.5*np.log10(CR_pix)#*f_pix)
         else: mag_pix = ZP - 2.5*np.log10(CR)

         Ftot_el_pix=CR_pix*f_pix*DIT_pix#*np.sqrt(Nexp)

         info_dict['DIT']=DIT
         info_dict['DIT_pix']=DIT_pix
         info_dict['Ftot_el_pix']=Ftot_el_pix
         info_dict['Ftot_el']=Ftot_el
         info_dict['mag_pix']=mag_pix

         if verbose == True:
              print ('\n\nReaching a magnitude (%s system) of %.2f in %s band with a SNR of %.2f requires:\n' %(info_dict['photometry_system'],mag,info_dict['filter_band'],SNR))
              print ('\t - a Total exposure time of : %.2f s\n' % (DIT * Nexp))
              print ('\n\nReaching a magnitude (%s system) of %.2f in %s band with a SNR of %.2f for the central pixel requires:\n\n' %(info_dict['photometry_system'],mag_pix,info_dict['filter_band'],SNR))
              print ('\t - a Total exposure time of : %.2f s\n\n' % (DIT_pix * Nexp))

    #-------------------------------------------------------------------------------------

    # Compute the magnitude
    #------------------------
    elif etc_type == 'mag':
         #f_PSF=1
         #f_pix=1
         # Integrated over Npixels    (solve 2nd degree equation)
         A_sys = -( f_PSF * DIT * np.sqrt(Nexp) )**2.
         B_sys = SNR**2.* f_PSF * DIT
         C_sys = SNR**2. * (factor_ima * npix * (RN**2. + DigN**2. + DIT * (DC + BN + inst_bg)))

         delta = B_sys*B_sys - 4.*A_sys*C_sys

         CR = (-B_sys - np.sqrt(delta)) / (2.*A_sys)
         Ftot_el=CR*f_PSF*DIT#*np.sqrt(Nexp)

         mag = ZP -2.5*np.log10(CR)

         # Central pixel
         A_sys = -( f_pix * DIT * np.sqrt(Nexp) )**2.
         B_sys = SNR**2.* f_pix * DIT
         C_sys = SNR**2. * (factor_ima * (RN**2. + DigN**2. + DIT * (DC + BN + inst_bg)))

         delta = B_sys*B_sys - 4.*A_sys*C_sys

         CR_pix = (-B_sys - np.sqrt(delta)) / (2.*A_sys)
         Ftot_el_pix=CR_pix*f_pix*DIT#*np.sqrt(Nexp)
         mag_pix = ZP -2.5*np.log10(CR_pix)
         DIT_pix=DIT

         object_mag = mag*np.ones(len(info_dict['wavelength_ang']))
         fJy = phot.mag2Jy(info_dict, object_mag)  # Jy
         flam = utils.fJy_to_flambda(info_dict['wavelength_ang'], fJy)        # erg/s/cm2/A
         fph = utils.flambda_to_fph(info_dict['wavelength_ang'], flam)        # ph/s/cm2/A

         info_dict['Object_mag']=object_mag
         info_dict['object_magnitude']=mag
         info_dict['mag_pix']=mag_pix
         info_dict['DIT_pix']=DIT
         info_dict['Ftot_el_pix']=Ftot_el_pix
         info_dict['Ftot_el']=Ftot_el

         if verbose == True:
              print ('\n\nFor a total SNR=%.2f in a total exposure time of %.2f (sec) in %d exposure(s) we reach:\n' %(SNR, DIT*Nexp, Nexp))
              print ('\t - a magnitude (%s system) of: %.2f in %s band\n' % (info_dict['photometry_system'],mag,info_dict['filter_band']))
              print ('\n\nFor the central pixel a SNR=%.2f in a total exposure time of %.2f (sec) in %d exposure(s) we reach:\n\n' %(SNR, DIT_pix*Nexp, Nexp))
              print ('\t - a magnitude (%s system) of: %.2f in %s band\n\n' % (info_dict['photometry_system'],mag_pix,info_dict['filter_band']))

    info_dict['DIT']=DIT

    sigma_shot_noise = np.sqrt(Ftot_el * DIT)
    sigma_dark_current = np.sqrt(DC * npix * DIT)
    sigma_sky = np.sqrt(BN * npix * DIT)
    sigma_digitization = np.sqrt( npix*DigN**2.)
    sigma_readout_noise = np.sqrt(npix*RN**2.)

    # Total number of electrons in the brightest pixel for 1 exposure
    N_el_tot_pix1 = Ftot_el_pix + (BN + DC + inst_bg)*DIT_pix + RN + DigN
    N_el_tot_pix2 = Ftot_el*f_pix/f_PSF + (BN + DC + inst_bg)*DIT + RN + DigN

    info_dict['N_el_tot_pix1']=N_el_tot_pix1
    info_dict['N_el_tot_pix2']=N_el_tot_pix2

    if N_el_tot_pix1 > info_dict['cameras'][info_dict['channel']]['FWC']:
        info_dict['saturation']='Yes'
    else: info_dict['saturation']='No'

    if verbose == True:

         #print ('\nFull well capacity of 1 pixel: %.2f (electrons)\nInverse gain of %.2f e/ADU and %d bits implies a maximum number of electrons to be digitized of  %.2f (electrons)' % (info_dict['cameras'][info_dict['channel']]['FWC'],info_dict['cameras'][info_dict['channel']]['gain'],info_dict['cameras'][info_dict['channel']]['bits'],info_dict['cameras'][info_dict['channel']]['gain']*(2.**(info_dict['cameras'][info_dict['channel']]['bits'])-1)))
         print ('\nFull well capacity of 1 pixel: %.2f (electrons)' % (info_dict['cameras'][info_dict['channel']]['FWC']))
         print ('\n\n--------- One pixel only------------------')
         print ('\nPhoto-electrons created: central pix for %d exposure(s) of %.2f sec ' % (Nexp,DIT_pix))
         print ('\tby:')
         print ('\t- Object:         %10.2f   (electrons)' % Ftot_el_pix)
         print ('\t- Sky:            %10.2f   (electrons)' % (BN*DIT_pix))
         print ('\t- Readout:        %10.2f   (electrons)' % RN)
         print ('\t- Dark current:   %10.2f   (electrons)' % (DC*DIT_pix))
         print ('\t- Digitization:   %10.2f   (electrons)' % DigN)
         print ('\t- Instrument bg:  %10.2f   (electrons)' % (inst_bg*DIT_pix))

         print ('\nSNR: -central pixel: %.2f' % (np.sqrt(Nexp)*Ftot_el_pix/np.sqrt(Ftot_el_pix + factor_ima * ((RN**2. + DigN**2.) + DIT_pix * ( DC + BN  + inst_bg)))))

         print ('\nTotal of electrons collected in the central pixel during an exposure time of %d seconds: %.2f ' % (DIT_pix, N_el_tot_pix1))
         if N_el_tot_pix1 > info_dict['cameras'][info_dict['channel']]['FWC']:
              print ('--> Central pixel saturated: number of electrons > Full well Capacity')
         elif N_el_tot_pix1 > info_dict['cameras'][info_dict['channel']]['gain']*(2.**(info_dict['cameras'][info_dict['channel']]['bits'])-1):
              print ('--> Central pixel saturated: number of electrons > number of digitizations')
         elif N_el_tot_pix1 > 1./2*info_dict['cameras'][info_dict['channel']]['FWC']:
              print ('--> Number of electrons in central pixel > 1/2 of Full well Capacity. Risk of non-linear response.')

         else:
              print ('--> No saturation')

         print ('\n\n\n--------- Integrated over %d pixels------------------' % npix)
         print ('\nPhoto-electrons created: brightest pix |  total of %d pixels, %d exposure(s) of %.2f sec ' % (npix,Nexp,DIT))
         print ('\tby:')
         print ('\t- Object:         %10.2f   |   %10.2f   (electrons)' % (Ftot_el*f_pix/f_PSF, Ftot_el))
         print ('\t- Sky:            %10.2f   |   %10.2f   (electrons)' % (BN*DIT,(BN * npix* DIT * Nexp)))
         print ('\t- Readout:        %10.2f   |   %10.2f   (electrons)' % (RN,(RN * npix * Nexp)))
         print ('\t- Dark current:   %10.2f   |   %10.2f   (electrons)' % (DC*DIT,(DC * DIT * npix * Nexp)))
         print ('\t- Digitization:   %10.2f   |   %10.2f   (electrons)' % (DigN, (DigN * npix * Nexp)))
         print ('\t- Instrument bg:  %10.2f   |   %10.2f   (electrons)' % (inst_bg*DIT, (inst_bg *DIT * npix * Nexp)))

         #print ('\nTotal noise: %.2f ' % (np.sqrt(Ftot_el * f_PSF * DIT *Nexp + Nexp*factor_ima * npix*((RN**2. + DigN**2.) + DIT * ( DC + BN )))))

         print ('\nSNR: -Brightest pixel: %.2f' % (np.sqrt(Nexp)*Ftot_el*f_pix/f_PSF/np.sqrt(Ftot_el*f_pix/f_PSF + factor_ima * ((RN**2. + DigN**2.) + DIT * ( DC + BN + inst_bg)))))
         print ('     -integrated over %d pixels: %.2f' % (npix,np.sqrt(Nexp)*Ftot_el / np.sqrt(Ftot_el + factor_ima * npix*((RN**2. + DigN**2.) + DIT * ( DC + BN + inst_bg)))))
         print ('\nTotal of electrons collected in the brightest pixel during an exposure time of %d seconds: %.2f ' % (DIT_pix, N_el_tot_pix2))
         if N_el_tot_pix2 > info_dict['cameras'][info_dict['channel']]['FWC']:
              print ('--> Brightest pixel saturated: number of electrons > Full well Capacity')
         elif N_el_tot_pix2 > info_dict['cameras'][info_dict['channel']]['gain']*(2.**(info_dict['cameras'][info_dict['channel']]['bits'])-1):
              print ('--> Brightest pixel saturated: number of electrons > number of digitizations')
         elif N_el_tot_pix2 > 1./2*info_dict['cameras'][info_dict['channel']]['FWC']:
              print ('--> Number of electrons in brightest pixel > 1/2 of Full well Capacity. Risk of non-linear response.')

         else:
              print ('--> No saturation')

         print ('\nDead time: %.2f sec \n(%.2f sec for dithering, the %.2f sec for the readout are not taken into account)' % (info_dict['deadtime_tot'],info_dict['T_dithering'],info_dict['cameras'][info_dict['channel']]['ReadoutTime']))
    info_dict['SNR']=SNR
    info_dict['mag']=mag
    info_dict['total_exposure_time']=DIT*Nexp
    info_dict['fph']=fph

    return info_dict


# We implemented our method on top of AB3DMOT's KITTI tracking open-source code

from __future__ import print_function
import os.path, copy, numpy as np, time, sys
from numba import jit
from sklearn.utils.linear_assignment_ import linear_assignment
from filterpy.kalman import KalmanFilter
from utils import load_list_from_folder, fileparts, mkdir_if_missing
from scipy.spatial import ConvexHull
from covariance import Covariance
import json
from nuscenes import NuScenes
from nuscenes.eval.common.data_classes import EvalBoxes
from nuscenes.eval.tracking.data_classes import TrackingBox
from nuscenes.eval.detection.data_classes import DetectionBox
from pyquaternion import Quaternion
from tqdm import tqdm

from KalmanBoxTracker import KalmanBoxTracker
from AB3DMOT import AB3DMOT

@jit
def poly_area(x,y):
    return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))

@jit
def box3d_vol(corners):
    ''' corners: (8,3) no assumption on axis direction '''
    a = np.sqrt(np.sum((corners[0,:] - corners[1,:])**2))
    b = np.sqrt(np.sum((corners[1,:] - corners[2,:])**2))
    c = np.sqrt(np.sum((corners[0,:] - corners[4,:])**2))
    return a*b*c

@jit
def convex_hull_intersection(p1, p2):
    """ Compute area of two convex hull's intersection area.
        p1,p2 are a list of (x,y) tuples of hull vertices.
        return a list of (x,y) for the intersection and its volume
    """
    inter_p = polygon_clip(p1,p2)
    if inter_p is not None:
        hull_inter = ConvexHull(inter_p)
        return inter_p, hull_inter.volume
    else:
        return None, 0.0

def polygon_clip(subjectPolygon, clipPolygon):
   """ Clip a polygon with another polygon.
   Args:
     subjectPolygon: a list of (x,y) 2d points, any polygon.
     clipPolygon: a list of (x,y) 2d points, has to be *convex*
   Note:
     **points have to be counter-clockwise ordered**

   Return:
     a list of (x,y) vertex point for the intersection polygon.
   """
   def inside(p):
      return(cp2[0]-cp1[0])*(p[1]-cp1[1]) > (cp2[1]-cp1[1])*(p[0]-cp1[0])

   def computeIntersection():
      dc = [ cp1[0] - cp2[0], cp1[1] - cp2[1] ]
      dp = [ s[0] - e[0], s[1] - e[1] ]
      n1 = cp1[0] * cp2[1] - cp1[1] * cp2[0]
      n2 = s[0] * e[1] - s[1] * e[0]
      n3 = 1.0 / (dc[0] * dp[1] - dc[1] * dp[0])
      return [(n1*dp[0] - n2*dc[0]) * n3, (n1*dp[1] - n2*dc[1]) * n3]

   outputList = subjectPolygon
   cp1 = clipPolygon[-1]

   for clipVertex in clipPolygon:
      cp2 = clipVertex
      inputList = outputList
      outputList = []
      s = inputList[-1]

      for subjectVertex in inputList:
         e = subjectVertex
         if inside(e):
            if not inside(s):
               outputList.append(computeIntersection())
            outputList.append(e)
         elif inside(s):
            outputList.append(computeIntersection())
         s = e
      cp1 = cp2
      if len(outputList) == 0:
          return None
   return(outputList)

def iou3d(corners1, corners2):
    ''' Compute 3D bounding box IoU.

    Input:
        corners1: numpy array (8,3), assume up direction is negative Y
        corners2: numpy array (8,3), assume up direction is negative Y
    Output:
        iou: 3D bounding box IoU
        iou_2d: bird's eye view 2D bounding box IoU

    '''
    # corner points are in counter clockwise order
    rect1 = [(corners1[i,0], corners1[i,2]) for i in range(3,-1,-1)]
    rect2 = [(corners2[i,0], corners2[i,2]) for i in range(3,-1,-1)]
    area1 = poly_area(np.array(rect1)[:,0], np.array(rect1)[:,1])
    area2 = poly_area(np.array(rect2)[:,0], np.array(rect2)[:,1])
    inter, inter_area = convex_hull_intersection(rect1, rect2)
    iou_2d = inter_area/(area1+area2-inter_area)
    ymax = min(corners1[0,1], corners2[0,1])
    ymin = max(corners1[4,1], corners2[4,1])
    inter_vol = inter_area * max(0.0, ymax-ymin)
    vol1 = box3d_vol(corners1)
    vol2 = box3d_vol(corners2)
    iou = inter_vol / (vol1 + vol2 - inter_vol)
    return iou, iou_2d

@jit
def roty(t):
    ''' Rotation about the y-axis. '''
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c,  0,  s],
                     [0,  1,  0],
                     [-s, 0,  c]])

@jit
def rotz(t):
    ''' Rotation about the z-axis. '''
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s,  0],
                     [s,  c,  0],
                     [0,  0,  1]])

def convert_3dbox_to_8corner(bbox3d_input, nuscenes_to_kitti=False):
    ''' Takes an object and a projection matrix (P) and projects the 3d
        bounding box into the image plane.
        Returns:
            corners_2d: (8,2) array in left image coord.
            corners_3d: (8,3) array in in rect camera coord.
        Note: the output of this function will be passed to the funciton iou3d
            for calculating the 3D-IOU. But the function iou3d was written for
            kitti, so the caller needs to set nuscenes_to_kitti to True if
            the input bbox3d_input is in nuscenes format.
    '''
    # compute rotational matrix around yaw axis
    bbox3d = copy.copy(bbox3d_input)

    if nuscenes_to_kitti:
      # transform to kitti format first
      bbox3d_nuscenes = copy.copy(bbox3d)
      # kitti:    [x,  y,  z,  a, l, w, h]
      # nuscenes: [y, -z, -x, -a, w, l, h]
      bbox3d[0] =  bbox3d_nuscenes[1]
      bbox3d[1] = -bbox3d_nuscenes[2]
      bbox3d[2] = -bbox3d_nuscenes[0]
      bbox3d[3] = -bbox3d_nuscenes[3]
      bbox3d[4] =  bbox3d_nuscenes[5]
      bbox3d[5] =  bbox3d_nuscenes[4]

    R = roty(bbox3d[3])

    # 3d bounding box dimensions
    l = bbox3d[4]
    w = bbox3d[5]
    h = bbox3d[6]

    # 3d bounding box corners
    x_corners = [l/2,l/2,-l/2,-l/2,l/2,l/2,-l/2,-l/2];
    y_corners = [0,0,0,0,-h,-h,-h,-h];
    z_corners = [w/2,-w/2,-w/2,w/2,w/2,-w/2,-w/2,w/2];

    # rotate and translate 3d bounding box
    corners_3d = np.dot(R, np.vstack([x_corners,y_corners,z_corners]))
    corners_3d[0,:] = corners_3d[0,:] + bbox3d[0]
    corners_3d[1,:] = corners_3d[1,:] + bbox3d[1]
    corners_3d[2,:] = corners_3d[2,:] + bbox3d[2]

    return np.transpose(corners_3d)

def angle_in_range(angle):
  '''
  Input angle: -2pi ~ 2pi
  Output angle: -pi ~ pi
  '''
  if angle > np.pi:
    angle -= 2 * np.pi
  if angle < -np.pi:
    angle += 2 * np.pi
  return angle

def diff_orientation_correction(det, trk):
  '''
  return the angle diff = det - trk
  if angle diff > 90 or < -90, rotate trk and update the angle diff
  '''
  diff = det - trk
  diff = angle_in_range(diff)
  if diff > np.pi / 2:
    diff -= np.pi
  if diff < -np.pi / 2:
    diff += np.pi
  diff = angle_in_range(diff)
  return diff

def greedy_match(distance_matrix):
  '''
  Find the one-to-one matching using greedy allgorithm choosing small distance
  distance_matrix: (num_detections, num_tracks)
  '''
  matched_indices = []

  num_detections, num_tracks = distance_matrix.shape
  distance_1d = distance_matrix.reshape(-1)
  index_1d = np.argsort(distance_1d)
  index_2d = np.stack([index_1d // num_tracks, index_1d % num_tracks], axis=1)
  detection_id_matches_to_tracking_id = [-1] * num_detections
  tracking_id_matches_to_detection_id = [-1] * num_tracks
  for sort_i in range(index_2d.shape[0]):
    detection_id = int(index_2d[sort_i][0])
    tracking_id = int(index_2d[sort_i][1])
    if tracking_id_matches_to_detection_id[tracking_id] == -1 and detection_id_matches_to_tracking_id[detection_id] == -1:
      tracking_id_matches_to_detection_id[tracking_id] = detection_id
      detection_id_matches_to_tracking_id[detection_id] = tracking_id
      matched_indices.append([detection_id, tracking_id])

  matched_indices = np.array(matched_indices)
  return matched_indices

def associate_detections_to_trackers(detections,trackers,iou_threshold=0.1,
  use_mahalanobis=False, dets=None, trks=None, trks_S=None, mahalanobis_threshold=0.1, print_debug=False, match_algorithm='greedy'):
  """
  Assigns detections to tracked object (both represented as bounding boxes)

  detections:  N x 8 x 3
  trackers:    M x 8 x 3

  dets: N x 7
  trks: M x 7
  trks_S: N x 7 x 7

  Returns 3 lists of matches, unmatched_detections and unmatched_trackers
  """
  if(len(trackers)==0):
    return np.empty((0,2),dtype=int), np.arange(len(detections)), np.empty((0,8,3),dtype=int)
  iou_matrix = np.zeros((len(detections),len(trackers)),dtype=np.float32)
  distance_matrix = np.zeros((len(detections),len(trackers)),dtype=np.float32)

  if use_mahalanobis:
    assert(dets is not None)
    assert(trks is not None)
    assert(trks_S is not None)

  if use_mahalanobis and print_debug:
    print('dets.shape: ', dets.shape)
    print('dets: ', dets)
    print('trks.shape: ', trks.shape)
    print('trks: ', trks)
    print('trks_S.shape: ', trks_S.shape)
    print('trks_S: ', trks_S)
    S_inv = [np.linalg.inv(S_tmp) for S_tmp in trks_S]  # 7 x 7
    S_inv_diag = [S_inv_tmp.diagonal() for S_inv_tmp in S_inv]# 7
    print('S_inv_diag: ', S_inv_diag)

  for d,det in enumerate(detections):
    for t,trk in enumerate(trackers):
      if use_mahalanobis:
        S_inv = np.linalg.inv(trks_S[t]) # 7 x 7
        diff = np.expand_dims(dets[d] - trks[t], axis=1) # 7 x 1
        # manual reversed angle by 180 when diff > 90 or < -90 degree
        corrected_angle_diff = diff_orientation_correction(dets[d][3], trks[t][3])
        diff[3] = corrected_angle_diff
        distance_matrix[d, t] = np.sqrt(np.matmul(np.matmul(diff.T, S_inv), diff)[0][0])
      else:
        iou_matrix[d,t] = iou3d(det,trk)[0]             # det: 8 x 3, trk: 8 x 3
        distance_matrix = -iou_matrix

  if match_algorithm == 'greedy':
    matched_indices = greedy_match(distance_matrix)
  elif match_algorithm == 'pre_threshold':
    if use_mahalanobis:
      to_max_mask = distance_matrix > mahalanobis_threshold
      distance_matrix[to_max_mask] = mahalanobis_threshold + 1
    else:
      to_max_mask = iou_matrix < iou_threshold
      distance_matrix[to_max_mask] = 0
      iou_matrix[to_max_mask] = 0
    matched_indices = linear_assignment(distance_matrix)      # houngarian algorithm
  else:
    matched_indices = linear_assignment(distance_matrix)      # houngarian algorithm

  if print_debug:
    print('distance_matrix.shape: ', distance_matrix.shape)
    print('distance_matrix: ', distance_matrix)
    print('matched_indices: ', matched_indices)

  unmatched_detections = []
  for d,det in enumerate(detections):
    if(d not in matched_indices[:,0]):
      unmatched_detections.append(d)
  unmatched_trackers = []
  for t,trk in enumerate(trackers):
    if len(matched_indices) == 0 or (t not in matched_indices[:,1]):
      unmatched_trackers.append(t)

  #filter out matched with low IOU
  matches = []
  for m in matched_indices:
    match = True
    if use_mahalanobis:
      if distance_matrix[m[0],m[1]] > mahalanobis_threshold:
        match = False
    else:
      if(iou_matrix[m[0],m[1]]<iou_threshold):
        match = False
    if not match:
      unmatched_detections.append(m[0])
      unmatched_trackers.append(m[1])
    else:
      matches.append(m.reshape(1,2))
  if(len(matches)==0):
    matches = np.empty((0,2),dtype=int)
  else:
    matches = np.concatenate(matches,axis=0)

  if print_debug:
    print('matches: ', matches)
    print('unmatched_detections: ', unmatched_detections)
    print('unmatched_trackers: ', unmatched_trackers)

  return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

NUSCENES_TRACKING_NAMES = [
  'bicycle',
  'bus',
  'car',
  'motorcycle',
  'pedestrian',
  'trailer',
  'truck'
]

def format_sample_result(sample_token, tracking_name, tracker):
  '''
  Input:
    tracker: (9): [h, w, l, x, y, z, rot_y], tracking_id, tracking_score
  Output:
  sample_result {
    "sample_token":   <str>         -- Foreign key. Identifies the sample/keyframe for which objects are detected.
    "translation":    <float> [3]   -- Estimated bounding box location in meters in the global frame: center_x, center_y, center_z.
    "size":           <float> [3]   -- Estimated bounding box size in meters: width, length, height.
    "rotation":       <float> [4]   -- Estimated bounding box orientation as quaternion in the global frame: w, x, y, z.
    "velocity":       <float> [2]   -- Estimated bounding box velocity in m/s in the global frame: vx, vy.
    "tracking_id":    <str>         -- Unique object id that is used to identify an object track across samples.
    "tracking_name":  <str>         -- The predicted class for this sample_result, e.g. car, pedestrian.
                                       Note that the tracking_name cannot change throughout a track.
    "tracking_score": <float>       -- Object prediction score between 0 and 1 for the class identified by tracking_name.
                                       We average over frame level scores to compute the track level score.
                                       The score is used to determine positive and negative tracks via thresholding.
  }
  '''
  rotation = Quaternion(axis=[0, 0, 1], angle=tracker[6]).elements
  sample_result = {
    'sample_token': sample_token,
    'translation': [tracker[3], tracker[4], tracker[5]],
    'size': [tracker[1], tracker[2], tracker[0]],
    'rotation': [rotation[0], rotation[1], rotation[2], rotation[3]],
    'velocity': [0, 0],
    'tracking_id': str(int(tracker[7])),
    'tracking_name': tracking_name,
    'tracking_score': tracker[8]
  }

  return sample_result

def track_nuscenes(data_split, covariance_id, match_distance, match_threshold, match_algorithm, save_root, use_angular_velocity):
  '''
  submission {
    "meta": {
        "use_camera":   <bool>  -- Whether this submission uses camera data as an input.
        "use_lidar":    <bool>  -- Whether this submission uses lidar data as an input.
        "use_radar":    <bool>  -- Whether this submission uses radar data as an input.
        "use_map":      <bool>  -- Whether this submission uses map data as an input.
        "use_external": <bool>  -- Whether this submission uses external data as an input.
    },
    "results": {
        sample_token <str>: List[sample_result] -- Maps each sample_token to a list of sample_results.
    }
  }

  '''
  save_dir = os.path.join(save_root, data_split); mkdir_if_missing(save_dir)
  if 'train' in data_split:
    detection_file = '/juno/u/hkchiu/dataset/nuscenes_new/megvii_train.json'
    data_root = '/juno/u/hkchiu/dataset/nuscenes/trainval'
    version='v1.0-trainval'
    output_path = os.path.join(save_dir, 'results_train_probabilistic_tracking.json')
  elif 'val' in data_split:
    detection_file = '/juno/u/hkchiu/dataset/nuscenes_new/megvii_val.json'
    data_root = '/juno/u/hkchiu/dataset/nuscenes/trainval'
    version='v1.0-trainval'
    output_path = os.path.join(save_dir, 'results_val_probabilistic_tracking.json')
  elif 'test' in data_split:
    detection_file = '/juno/u/hkchiu/dataset/nuscenes_new/megvii_test.json'
    data_root = '/juno/u/hkchiu/dataset/nuscenes/test'
    version='v1.0-test'
    output_path = os.path.join(save_dir, 'results_test_probabilistic_tracking.json')

  nusc = NuScenes(version=version, dataroot=data_root, verbose=True)

  results = {}

  total_time = 0.0
  total_frames = 0

  with open(detection_file) as f:
    data = json.load(f)
  assert 'results' in data, 'Error: No field `results` in result file. Please note that the result format changed.' \
    'See https://www.nuscenes.org/object-detection for more information.'

  all_results = EvalBoxes.deserialize(data['results'], DetectionBox)
  meta = data['meta']
  print('meta: ', meta)
  print("Loaded results from {}. Found detections for {} samples."
    .format(detection_file, len(all_results.sample_tokens)))

  processed_scene_tokens = set()
  for sample_token_idx in tqdm(range(len(all_results.sample_tokens))):
    sample_token = all_results.sample_tokens[sample_token_idx]
    scene_token = nusc.get('sample', sample_token)['scene_token']
    if scene_token in processed_scene_tokens:
      continue
    first_sample_token = nusc.get('scene', scene_token)['first_sample_token']
    current_sample_token = first_sample_token

    mot_trackers = {tracking_name: AB3DMOT(covariance_id, tracking_name=tracking_name, use_angular_velocity=use_angular_velocity, tracking_nuscenes=True) for tracking_name in NUSCENES_TRACKING_NAMES}

    while current_sample_token != '':
      results[current_sample_token] = []
      dets = {tracking_name: [] for tracking_name in NUSCENES_TRACKING_NAMES}
      info = {tracking_name: [] for tracking_name in NUSCENES_TRACKING_NAMES}
      for box in all_results.boxes[current_sample_token]:
        if box.detection_name not in NUSCENES_TRACKING_NAMES:
          continue
        q = Quaternion(box.rotation)
        angle = q.angle if q.axis[2] > 0 else -q.angle
        #print('box.rotation,  angle, axis: ', box.rotation, q.angle, q.axis)
        #print('box.rotation,  angle, axis: ', q.angle, q.axis)
        #[h, w, l, x, y, z, rot_y]
        detection = np.array([
          box.size[2], box.size[0], box.size[1],
          box.translation[0],  box.translation[1], box.translation[2],
          angle])
        #print('detection: ', detection)
        information = np.array([box.detection_score])
        dets[box.detection_name].append(detection)
        info[box.detection_name].append(information)

      dets_all = {tracking_name: {'dets': np.array(dets[tracking_name]), 'info': np.array(info[tracking_name])}
        for tracking_name in NUSCENES_TRACKING_NAMES}

      total_frames += 1
      start_time = time.time()
      for tracking_name in NUSCENES_TRACKING_NAMES:
        if dets_all[tracking_name]['dets'].shape[0] > 0:
          trackers = mot_trackers[tracking_name].update(dets_all[tracking_name], match_distance, match_threshold, match_algorithm, scene_token)
          # (N, 9)
          # (h, w, l, x, y, z, rot_y), tracking_id, tracking_score
          # print('trackers: ', trackers)
          for i in range(trackers.shape[0]):
            sample_result = format_sample_result(current_sample_token, tracking_name, trackers[i])
            results[current_sample_token].append(sample_result)
      cycle_time = time.time() - start_time
      total_time += cycle_time

      # get next frame and continue the while loop
      current_sample_token = nusc.get('sample', current_sample_token)['next']

    # left while loop and mark this scene as processed
    processed_scene_tokens.add(scene_token)

  # finished tracking all scenes, write output data
  output_data = {'meta': meta, 'results': results}
  with open(output_path, 'w') as outfile:
    json.dump(output_data, outfile)

  print("Total Tracking took: %.3f for %d frames or %.1f FPS"%(total_time,total_frames,total_frames/total_time))

if __name__ == '__main__':
  if len(sys.argv)!=9:
    print("Usage: python main.py data_split(train, val, test) covariance_id(0, 1, 2) match_distance(iou or m) match_threshold match_algorithm(greedy or h) use_angular_velocity(true or false) dataset save_root")
    sys.exit(1)

  data_split = sys.argv[1]
  covariance_id = int(sys.argv[2])
  match_distance = sys.argv[3]
  match_threshold = float(sys.argv[4])
  match_algorithm = sys.argv[5]
  use_angular_velocity = sys.argv[6] == 'True' or sys.argv[6] == 'true'
  dataset = sys.argv[7]
  save_root = os.path.join('./' + sys.argv[8])

  if dataset == 'kitti':
    print('track kitti not supported')
  elif dataset == 'nuscenes':
    print('track nuscenes')
    track_nuscenes(data_split, covariance_id, match_distance, match_threshold, match_algorithm, save_root, use_angular_velocity)