# Minhyuk Sung (mhsung@cs.stanford.edu)
# April 2018

import os, sys
BASE_DIR = os.path.normpath(
        os.path.join(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(BASE_DIR, '..'))

from datasets import *
from generate_outputs import *
from scipy.optimize import linear_sum_assignment
#import matplotlib.pyplot as plt
import numpy as np


def compute_all_keypoints(sess, net, data):
    P = data.point_clouds
    assert(P.shape[0] == data.n_data)
    assert(P.shape[1] == data.n_points)

    KP = data.keypoints
    assert(KP.shape[0] == data.n_data)
    assert(KP.shape[1] == data.n_labels)

    A = predict_A(P, sess, net)
    assert(A.shape[0] == data.n_data)
    assert(A.shape[1] == data.n_points)
    assert(A.shape[2] == net.K)

    pred_KP = np.argmax(A, axis=1)

    return P, KP, pred_KP


def evaluate_PCK(P, KP, pred_KP):
    n_data = P.shape[0]
    n_points = P.shape[1]
    n_labels = KP.shape[1]
    K = pred_KP.shape[1]

    # dists_info: (point_cloud_index, label, basis_index, distance)
    dists_info = []

    for k in range(n_data):
        # NOTE:
        # Skip if the keypoint does not exist.
        labels = [i for i in range(n_labels) if KP[k,i] >= 0]

        # Find the closest prediction (w/o matching).
        for i, label in enumerate(labels):
            all_dists = np.zeros(K)

            idx_i = KP[k,label]
            assert(idx_i < n_points)
            p_i = P[k,idx_i]

            for j in range(K):
                idx_j = pred_KP[k,j]
                assert(idx_j < n_points)
                p_j = P[k,idx_j]

                all_dists[j] = np.linalg.norm(p_i - p_j)

            j = np.argmin(all_dists)
            dists_info.append((k, i, j, all_dists[j]))

    dists_info = np.array(dists_info)

    return dists_info


def evaluate_PCK_after_label_basis_matching(P, KP, pred_KP):
    n_data = P.shape[0]
    n_points = P.shape[1]
    n_labels = KP.shape[1]
    K = pred_KP.shape[1]

    # Find the best mapping from labels to bases.
    all_dists = np.zeros((n_data, n_labels, K))
    label_counts = np.zeros(n_labels)

    for k in range(n_data):
        for i in range(n_labels):

            # NOTE:
            # Skip if the keypoint does not exist.
            if KP[k,i] < 0: continue

            idx_i = KP[k,i]
            assert(idx_i < n_points)
            p_i = P[k,idx_i]

            label_counts[i] += 1.

            for j in range(K):
                idx_j = pred_KP[k,j]
                assert(idx_j < n_points)
                p_j = P[k,idx_j]

                all_dists[k,i,j] += np.linalg.norm(p_i - p_j)

    mean_dists = np.sum(all_dists, axis=0) / \
            np.expand_dims(label_counts, axis=-1)
    row_ind, col_ind = linear_sum_assignment(mean_dists)


    # dists_info: (point_cloud_index, label, basis_index, distance)
    dists_info = []

    for k in range(n_data):
        for (i, j) in zip(row_ind, col_ind):
            if KP[k,i] < 0: continue
            dists_info.append((k, i, j, all_dists[k,i,j]))

    dists_info = np.array(dists_info)

    return dists_info


def save_results(dists_info, out_dir, postfix=None):
    # dists_info: (point_cloud_index, label, basis_index, distance)
    dists = dists_info[:,3]

    if postfix is not None:
        out_file = os.path.join(out_dir, 'distances_{}.npy'.format(postfix))
    else:
        out_file = os.path.join(out_dir, 'distances.npy')

    np.save(out_file, dists)
    print("Saved '{}'.".format(out_file))

    '''
    # Draw plot.
    n_matches = dists.size

    x_list = np.linspace(0.0, 0.1, 20 + 1)
    counts = np.zeros(x_list.size, dtype=int)

    for i in range(x_list.size):
        counts[i] = np.sum(dists <= x_list[i])

    y_list = counts.astype(x_list.dtype) / float(n_matches)

    plt.clf()
    plt.plot(x_list, y_list)
    plt.ylim(0., 1.)
    plt.yticks(np.linspace(0., 1., 10 + 1))

    if postfix is not None:
        out_file = os.path.join(out_dir, 'pck_{}.png'.format(postfix))
    else:
        out_file = os.path.join(out_dir, 'pck.png')

    plt.savefig(out_file)
    print("Saved '{}'.".format(out_file))
    '''


def evaluate(sess, net, data, out_dir):
    if not os.path.exists(out_dir): os.makedirs(out_dir)

    P, KP, pred_KP = compute_all_keypoints(sess, net, data)

    dists = evaluate_PCK(P, KP, pred_KP)
    save_results(dists, out_dir)

    dists_after_matching = evaluate_PCK_after_label_basis_matching(
            P, KP, pred_KP)
    save_results(dists_after_matching, out_dir, postfix='after_matching')


[((7, 16, 7, 44), 'os.path.join', 'os.path.join', ({(7, 29, 7, 37): 'BASE_DIR', (7, 39, 7, 43): '""".."""'}, {}), "(BASE_DIR, '..')", False, 'import os, sys\n'), ((30, 14, 30, 34), 'numpy.argmax', 'np.argmax', (), '', True, 'import numpy as np\n'), ((67, 17, 67, 37), 'numpy.array', 'np.array', ({(67, 26, 67, 36): 'dists_info'}, {}), '(dists_info)', True, 'import numpy as np\n'), ((79, 16, 79, 47), 'numpy.zeros', 'np.zeros', ({(79, 25, 79, 46): '(n_data, n_labels, K)'}, {}), '((n_data, n_labels, K))', True, 'import numpy as np\n'), ((80, 19, 80, 37), 'numpy.zeros', 'np.zeros', ({(80, 28, 80, 36): 'n_labels'}, {}), '(n_labels)', True, 'import numpy as np\n'), ((104, 23, 104, 56), 'scipy.optimize.linear_sum_assignment', 'linear_sum_assignment', ({(104, 45, 104, 55): 'mean_dists'}, {}), '(mean_dists)', False, 'from scipy.optimize import linear_sum_assignment\n'), ((115, 17, 115, 37), 'numpy.array', 'np.array', ({(115, 26, 115, 36): 'dists_info'}, {}), '(dists_info)', True, 'import numpy as np\n'), ((129, 4, 129, 28), 'numpy.save', 'np.save', ({(129, 12, 129, 20): 'out_file', (129, 22, 129, 27): 'dists'}, {}), '(out_file, dists)', True, 'import numpy as np\n'), ((102, 17, 102, 42), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((103, 12, 103, 49), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((127, 19, 127, 57), 'os.path.join', 'os.path.join', ({(127, 32, 127, 39): 'out_dir', (127, 41, 127, 56): '"""distances.npy"""'}, {}), "(out_dir, 'distances.npy')", False, 'import os, sys\n'), ((160, 11, 160, 34), 'os.path.exists', 'os.path.exists', ({(160, 26, 160, 33): 'out_dir'}, {}), '(out_dir)', False, 'import os, sys\n'), ((160, 36, 160, 56), 'os.makedirs', 'os.makedirs', ({(160, 48, 160, 55): 'out_dir'}, {}), '(out_dir)', False, 'import os, sys\n'), ((6, 37, 6, 62), 'os.path.abspath', 'os.path.abspath', ({(6, 53, 6, 61): '__file__'}, {}), '(__file__)', False, 'import os, sys\n'), ((51, 24, 51, 35), 'numpy.zeros', 'np.zeros', ({(51, 33, 51, 34): 'K'}, {}), '(K)', True, 'import numpy as np\n'), ((64, 16, 64, 36), 'numpy.argmin', 'np.argmin', ({(64, 26, 64, 35): 'all_dists'}, {}), '(all_dists)', True, 'import numpy as np\n'), ((62, 31, 62, 56), 'numpy.linalg.norm', 'np.linalg.norm', ({(62, 46, 62, 55): 'p_i - p_j'}, {}), '(p_i - p_j)', True, 'import numpy as np\n'), ((100, 36, 100, 61), 'numpy.linalg.norm', 'np.linalg.norm', ({(100, 51, 100, 60): '(p_i - p_j)'}, {}), '(p_i - p_j)', True, 'import numpy as np\n')]