#!/usr/bin/env python
import sys
import numpy as np
import argparse

Res31 = {'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',
         'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',
         'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',
         'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',
         'ASX': 'N', 'GLX': 'Q', 'UNK': 'X', 'INI': 'K', 'AAR': 'R',
         'ACE': 'X', 'ACY': 'G', 'AEI': 'T', 'AGM': 'R', 'ASQ': 'D',
         'AYA': 'A', 'BHD': 'D', 'CAS': 'C', 'CAY': 'C', 'CEA': 'C',
         'CGU': 'E', 'CME': 'C', 'CMT': 'C', 'CSB': 'C', 'CSD': 'C',
         'CSE': 'C', 'CSO': 'C', 'CSP': 'C', 'CSS': 'C', 'CSW': 'C',
         'CSX': 'C', 'CXM': 'M', 'CYG': 'C', 'CYM': 'C', 'DOH': 'D',
         'EHP': 'F', 'FME': 'M', 'FTR': 'W', 'GL3': 'G', 'H2P': 'H',
         'HIC': 'H', 'HIP': 'H', 'HTR': 'W', 'HYP': 'P', 'KCX': 'K',
         'LLP': 'K', 'LLY': 'K', 'LYZ': 'K', 'M3L': 'K', 'MEN': 'N',
         'MGN': 'Q', 'MHO': 'M', 'MHS': 'H', 'MIS': 'S', 'MLY': 'K',
         'MLZ': 'K', 'MSE': 'M', 'NEP': 'H', 'NPH': 'C', 'OCS': 'C',
         'OCY': 'C', 'OMT': 'M', 'OPR': 'R', 'PAQ': 'Y', 'PCA': 'Q',
         'PHD': 'D', 'PRS': 'P', 'PTH': 'Y', 'PYX': 'C', 'SEP': 'S',
         'SMC': 'C', 'SME': 'M', 'SNC': 'C', 'SNN': 'D', 'SVA': 'S',
         'TPO': 'T', 'TPQ': 'Y', 'TRF': 'W', 'TRN': 'W', 'TRO': 'W',
         'TYI': 'Y', 'TYN': 'Y', 'TYQ': 'Y', 'TYS': 'Y', 'TYY': 'Y',
         'YOF': 'Y', 'FOR': 'X', '---': '-', 'PTR': 'Y', 'LCX': 'K',
         'SEC': 'D', 'MCL': 'K', 'LDH': 'K'}

typecode = "ARNDCQEGHILKMFPSTWYVX"
Ntype = 21
code_dict = dict()

def read_pdb(pdb_file, hydrogen=False):

    main_chain_atoms = ['N', 'CA', 'C', 'O']
    fp_file = open(pdb_file)
    lines = fp_file.readlines()
    fp_file.close()

    chain_dict = dict()
    for line in lines:
        if line[:6] == 'ATOM  ':
            chain_id = line[21]
            res_idx = line[22:27]  # not integer
            atom_idx = int(line[6:11].strip())
            atomname = line[12:16]
            atomname2 = atomname.strip()
            res_name = line[17:20]
            altloc = line[16]
            coor = np.array([float(line[30:38]), float(line[38:46]),
                             float(line[46:54])], dtype=np.float32)
            if altloc != ' ' and altloc != 'A':
                continue
            if atomname2[0] == 'H' and not hydrogen:
                continue
            if chain_id not in chain_dict:
                chain_dict[chain_id] = dict()
            if res_idx not in chain_dict[chain_id]:
                chain_dict[chain_id][res_idx] = {
                    'res_name': res_name, 'atom': dict()}
            chain_dict[chain_id][res_idx]['atom'][atom_idx] = {
                'atomname': atomname, 'coor': coor,
                'line': line}

    chain_keys = sorted(chain_dict.keys())
    for chain_id in chain_keys:
        chain = chain_dict[chain_id]
        res_idx_keys = sorted(chain.keys())
        for res_idx in res_idx_keys:
            residue = chain[res_idx]
            res_name = residue['res_name']
            atom_dict = residue['atom']
            atom_keys = sorted(atom_dict.keys())
            sg_coor = []
            for atom_idx in atom_keys:
                atom = atom_dict[atom_idx]
                atomname = atom['atomname']
                coor = atom['coor']
                if atomname == ' CA ':
                    chain_dict[chain_id][res_idx]['CA'] = coor
                    if res_name != 'GLY':
                        chain_dict[chain_id][res_idx]['CB'] = coor
                if atomname == ' CB ':
                    chain_dict[chain_id][res_idx]['CB'] = coor
                atomname2 = atomname.strip()
                if res_name != 'GLY':
                    if (atomname2 not in main_chain_atoms) and (atomname2[0] != 'H'):
                        sg_coor += [coor]
            if res_name != 'GLY':
                sg_coor = np.concatenate([sg_coor]).mean(axis=0)
            else:
                sg_coor = chain_dict[chain_id][res_idx]['CA']
            chain_dict[chain_id][res_idx]['SG'] = sg_coor
    return chain_dict

def read_pdb_ligand(file_name, ligand_name, hydrogen=False):
    fp = open(file_name)
    lines = fp.readlines()
    fp.close()

    coor_list = []
    for line in lines:
        if line[0:6] != "HETATM" and line[0:6] != "ATOM  ":
            continue
        lig_name = line[17:20].strip()
        if lig_name != ligand_name and ligand_name != '':
            continue
        if line[13] == 'H' and not hydrogen:
            continue
        coor = [float(line[30:38]), float(line[38:46]), float(line[46:54])]
        coor_list += [coor]
    coor = np.array(coor_list)

    return coor

def read_mol2_ligand(file_name, ligand_name, hydrogen=False):
    fp = open(file_name)
    lines = fp.readlines()
    fp.close()

    coor_list = []
    line = lines[3]

    state = ''
    line_count = 0
    for line in lines:
        if line.startswith('@<TRIPOS>'):
            state = line[9:].strip()
            line_count = 0
            continue
        if state == "MOLECULE":
            line_count += 1
            if line_count == 2:
                num_atoms = int(line[0:5])
                num_bonds = int(line[5:11])
            continue
        if state == "ATOM":
            lig_name = line[59:62].strip()
            if lig_name != ligand_name and ligand_name != '':
                continue
            if line[47] == 'H' and not hydrogen:
                continue
            coor = [float(line[16:26]), float(line[26:36]), float(line[36:46])]
            coor_list += [coor]
    coor = np.array(coor_list)

    return coor

def read_mol_ligand(file_name, ligand_name, hydrogen=False):
    fp = open(file_name)
    lines = fp.readlines()
    fp.close()

    coor_list = []
    line = lines[3]
    num_atoms = int(line[0:3])
    num_bonds = int(line[3:6])
    mol_version = line[34:39]

    for line in lines[4:4+num_atoms]:
        if line[31] == 'H' and not hydrogen:
            continue
        coor = [float(line[0:10]), float(line[10:20]), float(line[20:30])]
        coor_list += [coor]
    coor = np.array(coor_list)

    return coor

def find_pocket(chain_dict, ligand_coor, dist_cutoff=5):

    cmin = ligand_coor.min(axis=0) - dist_cutoff
    cmax = ligand_coor.max(axis=0) + dist_cutoff
    num_ligand_atoms = ligand_coor.shape[0]
    chain_pocket = dict()
    pocket_res_list = list()
    chain_keys = sorted(chain_dict.keys())
    for chain_id in chain_keys:
        chain = chain_dict[chain_id]
        res_idx_keys = sorted(chain.keys())
        for res_idx in res_idx_keys:
            pocket_residue = False
            residue = chain[res_idx]
            res_name = residue['res_name']
            atom_dict = residue['atom']
            atom_keys = sorted(atom_dict.keys())
            for atom_idx in atom_keys:
                atom = atom_dict[atom_idx]
                atomname = atom['atomname']
                atomname2 = atomname.strip()
                if atomname2[0] == 'H':
                    continue
                coor = atom['coor']
                if (coor < cmin).any() or (coor > cmax).any():
                    continue
                dist_min = 100
                for j in range(0, num_ligand_atoms):
                    dist = np.linalg.norm(coor-ligand_coor[j])
                    if dist < dist_min:
                        dist_min = dist
                        lc = ligand_coor[j]
                    if dist < dist_cutoff:
                        pocket_residue = True
                        break
                if pocket_residue:
                    break
            if pocket_residue:
                if chain_id not in chain_pocket:
                    chain_pocket[chain_id] = dict()
                chain_pocket[chain_id][res_idx] = residue
                pocket_res = '%s:%s' % (chain_id,
                                        Res31[res_name] + res_idx.strip())
                pocket_res_list += [pocket_res]
    return chain_pocket, pocket_res_list

def write_pocket(chain_dict, pocket_file):

    fp_out = open(pocket_file, 'w')
    chain_keys = sorted(chain_dict.keys())
    for chain_id in chain_keys:
        chain = chain_dict[chain_id]
        res_idx_keys = sorted(chain.keys())
        pocket_residue = False
        for res_idx in res_idx_keys:
            residue = chain[res_idx]
            res_name = residue['res_name']
            atom_dict = residue['atom']
            atom_keys = sorted(atom_dict.keys())
            for atom_idx in atom_keys:
                atom = atom_dict[atom_idx]
                line_out = atom['line']
                fp_out.write(line_out)
    fp_out.close()
    return

def write_binding_residue_file(pocket_res_list, binding_file):
    fp_out = open(binding_file, 'w')
    binding_chain = dict()
    for pocket_res in pocket_res_list:
        line_out = pocket_res + '\n'
        fp_out.write(line_out)
    fp_out.close()
    return

def main():
    parser = argparse.ArgumentParser(description='find_pocket')
    parser.add_argument('-p', '--protein_file', type=str,
                        required=True, help='protein file, ' +
                        'ex: -p protein.pdb')
    parser.add_argument('-o', '--output_file', type=str,
                        required=True, help='output protein file, ' +
                        'ex: -o pocket.pdb')
    parser.add_argument('-b', '--residue_file', type=str,
                        required=False, default='binding_residue.txt',
                        help='output file for binding_residue' +
                        ', ex: -o binding_residue.txt')

    parser.add_argument('-l', '--ligand_file', type=str,
                        required=True, help='ligand file, ' +
                        'ex: -l ligand.pdb, only pdb, mol2, mol, sdf')

    parser.add_argument('-n', '--ligand_name', type=str, default='',
                        required=False, help='ligand name in PDB, ' +
                        'ex: -n HYZ, default: -n ""')

    parser.add_argument('-c', '--dist_cutoff', type=float, default=8.0,
                        required=False, help='distant cutoff, '
                        'unit:A, default: 8.0 ')

    parser.add_argument('--hydrogen', action='store_true',
                        required=False,
                        help='--hydrogen: include hydrogen atoms or' +
                        ' exclude hydrogen atoms ')

    args = parser.parse_args()
    protein_file = args.protein_file
    ligand_file = args.ligand_file
    ligand_name = args.ligand_name.strip()
    pocket_file = args.output_file
    binding_residue_file = args.residue_file

    dist_cutoff = args.dist_cutoff
    hydrogen = args.hydrogen

    chain_dict = read_pdb(protein_file, hydrogen=hydrogen)

    file_format = ligand_file.strip().split('.')[-1].lower()
    if file_format == 'pdb':
        ligand_coor = read_pdb_ligand(ligand_file, ligand_name, hydrogen)
    elif file_format == 'mol' or file_format == 'sdf':
        ligand_coor = read_mol_ligand(ligand_file, ligand_name, hydrogen)
    elif file_format == 'mol2':
        ligand_coor = read_mol2_ligand(ligand_file, ligand_name, hydrogen)
    else:
        print('file format ', file_format, 'is not supported')

    chain_pocket, pocket_res_list = find_pocket(
        chain_dict, ligand_coor, dist_cutoff)
#    print(chain_pocket)
#    print(pocket_res_list)
    write_pocket(chain_pocket, pocket_file)
    write_binding_residue_file(pocket_res_list, binding_residue_file)

    binding_chain = dict()
    for pocket_res in pocket_res_list:
        chain_id = pocket_res[0]
        if chain_id not in binding_chain:
            binding_chain[chain_id] = 0
        binding_chain[chain_id] += 1
#    print(binding_chain)
    line_out = ''
    for chain_id in binding_chain:
        cnum = binding_chain[chain_id]
        line_out += '%s:%d,' % (chain_id, cnum)
    line_out = line_out.strip(',')
    print(line_out)

if __name__ == "__main__":
    main()


import os
import unittest
from unittest.mock import MagicMock, patch

import numpy as np

from nuscenes_radar_devkit.nuscenes import NuScenes
from nuscenes_radar_devkit.nuscenes.eval.prediction import metrics
from nuscenes_radar_devkit.nuscenes.eval.prediction.data_classes import Prediction
from nuscenes_radar_devkit.nuscenes.prediction import PredictHelper

class TestFunctions(unittest.TestCase):

    def setUp(self):
        self.x_one_mode = np.ones((1, 5, 2))
        self.y_one_mode = np.expand_dims(np.arange(10).reshape(5, 2), 0)
        self.p_one_mode = np.array([[1.]])

        x_many_modes = np.repeat(self.x_one_mode, 3, axis=0)
        x_many_modes[0, :] = 0
        x_many_modes[2, :] = 2
        self.x_many_modes = x_many_modes
        self.y_many_modes = np.repeat(self.y_one_mode, 3, axis=0)
        self.p_many_modes = np.array([[0.2, 0.5, 0.3]])

        self.x_many_batches_and_modes = np.repeat(np.expand_dims(self.x_many_modes, 0), 5, axis=0)
        self.y_many_batches_and_modes = np.repeat(np.expand_dims(self.y_many_modes, 0), 5, axis=0)
        self.p_many_batches_and_modes = np.array([[0.2, 0.5, 0.3],
                                                  [0.5, 0.3, 0.2],
                                                  [0.2, 0.3, 0.5],
                                                  [0.3, 0.2, 0.5],
                                                  [0.3, 0.5, 0.2]])

    def test_returns_2d_array_float(self):

        func = lambda x: 2
        value = metrics.returns_2d_array(func)(2)
        np.testing.assert_equal(value, np.array([[2]]))

        func = lambda x: 3.
        value = metrics.returns_2d_array(func)(np.ones((10, 1)))
        np.testing.assert_equal(value, np.array([[3]]))

    def test_returns_2d_array_one_dim(self):

        func = lambda x: np.ones(10)
        value = metrics.returns_2d_array(func)(1)
        np.testing.assert_equal(value, np.ones((1, 10)))

    def test_mean_distances_one_mode(self):

        value = metrics.mean_distances(self.x_one_mode, self.y_one_mode)
        np.testing.assert_allclose(value, np.array([[5.33529]]), atol=1e-4, rtol=1e-4)

    def test_mean_distances_many_modes(self):
        value = metrics.mean_distances(self.x_many_modes, self.y_many_modes)
        np.testing.assert_allclose(value, np.array([[6.45396, 5.33529, 4.49286]]), atol=1e-4, rtol=1e-4)

    def test_mean_distances_many_batches_and_modes(self):
        value = metrics.mean_distances(self.x_many_batches_and_modes, self.y_many_batches_and_modes)
        np.testing.assert_allclose(value, np.array(5*[[6.45396, 5.33529, 4.49286]]), atol=1e-4, rtol=1e-4)

    def test_max_distances_one_mode(self):
        value = metrics.max_distances(self.x_one_mode, self.y_one_mode)
        np.testing.assert_allclose(value, np.array([[10.63014]]), atol=1e-4, rtol=1e-4)

    def test_max_distances_many_modes(self):
        value = metrics.max_distances(self.x_many_modes, self.y_many_modes)
        np.testing.assert_allclose(value, np.array([[12.04159, 10.63014, 9.21954]]), atol=1e-4, rtol=1e-4)

    def test_max_distances_many_batches_and_modes(self):
        value = metrics.max_distances(self.x_many_batches_and_modes, self.y_many_batches_and_modes)
        np.testing.assert_allclose(value, np.array(5*[[12.04159, 10.63014, 9.21954]]), atol=1e-4, rtol=1e-4)

    def test_final_distances_one_mode(self):
        value = metrics.max_distances(self.x_one_mode, self.y_one_mode)
        np.testing.assert_allclose(value, np.array([[10.63014]]), atol=1e-4, rtol=1e-4)

    def test_final_distances_many_modes(self):
        value = metrics.max_distances(self.x_many_modes, self.y_many_modes)
        np.testing.assert_allclose(value, np.array([[12.04159, 10.63014, 9.21954]]), atol=1e-4, rtol=1e-4)

    def test_final_distances_many_batches_and_modes(self):
        value = metrics.max_distances(self.x_many_batches_and_modes, self.y_many_batches_and_modes)
        np.testing.assert_allclose(value, np.array(5*[[12.04159, 10.63014, 9.21954]]), atol=1e-4, rtol=1e-4)

    def test_miss_max_distance_one_mode(self):
        value = metrics.miss_max_distances(self.x_one_mode, self.y_one_mode, 1)
        np.testing.assert_equal(value, np.array([[True]]))

        value = metrics.miss_max_distances(self.x_one_mode, self.y_one_mode, 15)
        np.testing.assert_equal(value, np.array([[False]]))

    def test_miss_max_distances_many_modes(self):
        value = metrics.miss_max_distances(self.x_many_modes, self.y_many_modes, 10)
        np.testing.assert_equal(value, np.array([[True, True, False]]))

    def test_miss_max_distances_many_batches_and_modes(self):
        value = metrics.miss_max_distances(self.x_many_batches_and_modes, self.y_many_batches_and_modes, 10)
        np.testing.assert_equal(value, np.array(5*[[True, True, False]]))

    def test_miss_rate_top_k_one_mode(self):
        value = metrics.miss_rate_top_k(self.x_one_mode, self.y_one_mode, self.p_one_mode, 2)
        np.testing.assert_equal(value, np.array([[True]]))

    def test_miss_rate_top_k_many_modes(self):
        value = metrics.miss_rate_top_k(self.x_many_modes, self.y_many_modes, self.p_many_modes, 10)
        np.testing.assert_equal(value, np.array([[True, False, False]]))

    def test_miss_rate_top_k_many_batches_and_modes(self):
        value = metrics.miss_rate_top_k(self.x_many_batches_and_modes,
                                        self.y_many_batches_and_modes, self.p_many_batches_and_modes, 10)
        np.testing.assert_equal(value, np.array([[True, False, False],
                                                 [True, True, False],
                                                 [False, False, False],
                                                 [False, False, False],
                                                 [True, True, False]]))

    def test_min_ade_k_one_mode(self):
        value = metrics.min_ade_k(self.x_one_mode, self.y_one_mode, self.p_one_mode)
        np.testing.assert_allclose(value, np.array([[5.33529]]), atol=1e-4, rtol=1e-4)

    def test_min_ade_k_many_modes(self):
        value = metrics.min_ade_k(self.x_many_modes, self.y_many_modes, self.p_many_modes)
        np.testing.assert_allclose(value, np.array([[5.33529, 4.49286, 4.49286]]), atol=1e-4, rtol=1e-4)

    def test_min_ade_k_many_batches_and_modes(self):
        value = metrics.min_ade_k(self.x_many_batches_and_modes, self.y_many_batches_and_modes,
                                  self.p_many_batches_and_modes)
        np.testing.assert_allclose(value, np.array([[5.33529, 4.49286, 4.49286],
                                                    [6.45396, 5.33529, 4.49286],
                                                    [4.49286, 4.49286, 4.49286],
                                                    [4.49286, 4.49286, 4.49286],
                                                    [5.33529, 5.33529, 4.49286]
                                                    ]), atol=1e-4, rtol=1e-4)

    def test_min_fde_k_one_mode(self):
        value = metrics.min_fde_k(self.x_one_mode, self.y_one_mode, self.p_one_mode)
        np.testing.assert_allclose(value, np.array([[10.63014]]), atol=1e-4, rtol=1e-4)

    def test_min_fde_k_many_modes(self):
        value = metrics.min_fde_k(self.x_many_modes, self.y_many_modes, self.p_many_modes)
        np.testing.assert_allclose(value, np.array([[10.63014, 9.21954, 9.21954]]), atol=1e-4, rtol=1e-4)

    def test_min_fde_k_many_batches_and_modes(self):
        value = metrics.min_fde_k(self.x_many_batches_and_modes, self.y_many_batches_and_modes,
                                  self.p_many_batches_and_modes)
        np.testing.assert_allclose(value, np.array([[10.63014, 9.21954, 9.21954],
                                                    [12.04159, 10.63014, 9.21954],
                                                    [9.21954, 9.21954, 9.21954],
                                                    [9.21954, 9.21954, 9.21954],
                                                    [10.63014, 10.64014, 9.21954]]), atol=1e-3, rtol=1e-3)

    def test_stack_ground_truth(self):
        value = metrics.stack_ground_truth(np.ones((5, 2)), 10)
        np.testing.assert_equal(value, np.ones((10, 5, 2)))

    def test_desired_number_of_modes_one_mode(self):
        results = np.ones((10, 1))
        value = metrics.desired_number_of_modes(results, [1, 5, 15, 25])
        np.testing.assert_equal(value, np.ones((10, 4)))

    def test_desired_number_of_modes_enough_data(self):
        results = np.arange(75).reshape(3, 25)
        value = metrics.desired_number_of_modes(results, [1, 5, 15, 25])
        np.testing.assert_equal(value, np.array([[0, 4, 14, 24],
                                                 [25, 29, 39, 49],
                                                 [50, 54, 64, 74]]))

    def test_desired_number_of_modes_not_enough(self):
        results = np.arange(30).reshape(2, 15)
        value = metrics.desired_number_of_modes(results, [1, 5, 15, 25])
        np.testing.assert_equal(value, np.array([[0, 4, 14, 14],
                                                 [15, 19, 29, 29]]))

class TestAggregators(unittest.TestCase):

    def test_RowMean(self):
        rm = metrics.RowMean()
        value = rm(np.arange(20).reshape(2, 10))
        self.assertListEqual(list(value), [5, 6, 7, 8, 9, 10, 11, 12, 13, 14])

        self.assertDictEqual(rm.serialize(), {'name': 'RowMean'})

class TestMetrics(unittest.TestCase):

    def test_MinADEK(self):
        min_ade = metrics.MinADEK([1, 5, 10], [metrics.RowMean()])
        self.assertDictEqual(min_ade.serialize(), {'name': 'MinADEK',
                                                   'k_to_report': [1, 5, 10],
                                                   'aggregators': [{'name': 'RowMean'}]})

    def test_MinFDEK(self):
        min_fde = metrics.MinFDEK([1, 5, 10], [metrics.RowMean()])
        self.assertDictEqual(min_fde.serialize(), {'name': 'MinFDEK',
                                                   'k_to_report': [1, 5, 10],
                                                   'aggregators': [{'name': 'RowMean'}]})

    def test_MissRateTopK(self):
        hit_rate = metrics.MissRateTopK([1, 5, 10], [metrics.RowMean()], 2)
        self.assertDictEqual(hit_rate.serialize(), {'k_to_report': [1, 5, 10],
                                                    'name': 'MissRateTopK',
                                                    'aggregators': [{'name': 'RowMean'}],
                                                    'tolerance': 2})

    def test_OffRoadRate(self):
        with patch.object(metrics.OffRoadRate, 'load_drivable_area_masks'):
            helper = MagicMock(spec=PredictHelper)
            off_road_rate = metrics.OffRoadRate(helper, [metrics.RowMean()])
            self.assertDictEqual(off_road_rate.serialize(), {'name': 'OffRoadRate',
                                                             'aggregators': [{'name': 'RowMean'}]})

    def test_deserialize_metric(self):

        config = {'name': 'MinADEK',
                  'k_to_report': [1, 5, 10],
                  'aggregators': [{'name': 'RowMean'}]}

        helper = MagicMock(spec=PredictHelper)
        m = metrics.deserialize_metric(config, helper)
        self.assertEqual(m.name, 'MinADEK')
        self.assertListEqual(m.k_to_report, [1, 5, 10])
        self.assertEqual(m.aggregators[0].name, 'RowMean')

        config = {'name': 'MinFDEK',
                  'k_to_report': [1, 5, 10],
                  'aggregators': [{'name': 'RowMean'}]}

        m = metrics.deserialize_metric(config, helper)
        self.assertEqual(m.name, 'MinFDEK')
        self.assertListEqual(m.k_to_report, [1, 5, 10])
        self.assertEqual(m.aggregators[0].name, 'RowMean')

        config = {'name': 'MissRateTopK',
                  'k_to_report': [1, 5, 10],
                  'tolerance': 2,
                  'aggregators': [{'name': 'RowMean'}]}

        m = metrics.deserialize_metric(config, helper)
        self.assertEqual(m.name, 'MissRateTopK_2')
        self.assertListEqual(m.k_to_report, [1, 5, 10])
        self.assertEqual(m.aggregators[0].name, 'RowMean')

        with patch.object(metrics.OffRoadRate, 'load_drivable_area_masks'):
            config = {'name': 'OffRoadRate',
                      'aggregators': [{'name': 'RowMean'}]}

            m = metrics.deserialize_metric(config, helper)
            self.assertEqual(m.name, 'OffRoadRate')
            self.assertEqual(m.aggregators[0].name, 'RowMean')

    def test_flatten_metrics(self):
        results = {"MinFDEK": {"RowMean": [5.92, 6.1, 7.2]},
                   "MinADEK": {"RowMean": [2.48, 3.29, 3.79]},
                   "MissRateTopK_2": {"RowMean": [0.37, 0.45, 0.55]}}

        metric_functions = [metrics.MinFDEK([1, 5, 10], aggregators=[metrics.RowMean()]),
                            metrics.MinADEK([1, 5, 10], aggregators=[metrics.RowMean()]),
                            metrics.MissRateTopK([1, 5, 10], tolerance=2, aggregators=[metrics.RowMean()])]

        flattened = metrics.flatten_metrics(results, metric_functions)

        answer = {'MinFDEK_1': 5.92, 'MinFDEK_5': 6.1, 'MinFDEK_10': 7.2,
                  'MinADEK_1': 2.48, 'MinADEK_5': 3.29, 'MinADEK_10': 3.79,
                  'MissRateTopK_2_1': 0.37, 'MissRateTopK_2_5': 0.45, 'MissRateTopK_2_10': 0.55}

        self.assertDictEqual(flattened, answer)

class TestOffRoadRate(unittest.TestCase):

    def _do_test(self, map_name, predictions, answer):
        with patch.object(PredictHelper, 'get_map_name_from_sample_token') as get_map_name:
            get_map_name.return_value = map_name
            nusc = NuScenes('v1.0-mini', dataroot=os.environ['NUSCENES'], verbose=False)
            helper = PredictHelper(nusc)

            off_road_rate = metrics.OffRoadRate(helper, [metrics.RowMean()])

            probabilities = np.array([1/3] * predictions.shape[0])
            prediction = Prediction('foo-instance', 'foo-sample', predictions, probabilities)

            # Two violations out of three trajectories
            np.testing.assert_allclose(off_road_rate(np.array([]), prediction), np.array([answer]))

    def test_boston(self):
        predictions = np.array([[(486.91778944573264, 812.8782745377198),
                                 (487.3648565923963, 813.7269620253566),
                                 (487.811923719944, 814.5756495230632),
                                 (488.2589908474917, 815.4243370207698)],
                                [(486.91778944573264, 812.8782745377198),
                                 (487.3648565923963, 813.7269620253566),
                                 (487.811923719944, 814.5756495230632),
                                 (0, 0)],
                                [(0, 0), (0, 1), (0, 2), (0, 3)]])
        self._do_test('boston-seaport', predictions, 2/3)

    def test_one_north(self):
        predictions = np.array([[[965.8515334916171, 535.711518726687],
                                 [963.6475430050381, 532.9713854167148],
                                 [961.4435525191437, 530.231252106192],
                                 [959.239560587773, 527.4911199583674]],
                                [[508.8742570078554, 875.3458194583762],
                                 [505.2029816111618, 877.7929160023881],
                                 [501.5317062144682, 880.2400125464],
                                 [497.86043081777467, 882.6871090904118]],
                                [[0, 0], [0, 1], [0, 2], [0, 3]]])
        self._do_test('singapore-onenorth', predictions, 1/3)

    def test_queenstown(self):
        predictions = np.array([[[744.8769428947988, 2508.398411382534],
                                 [747.7808552527478, 2507.131371270205],
                                 [750.7893530020073, 2506.1385301483474],
                                 [751, 2506]],
                                [[-100, 0], [-10, 100], [0, 2], [-20, 70]]])
        self._do_test('singapore-queenstown', predictions, 1/2)

    def test_hollandvillage(self):
        predictions = np.array([[(1150.811356677105, 1598.0397224872172),
                                (1158.783061670897, 1595.5210995059333),
                                (1166.7543904812692, 1593.0012894706226),
                                (1174.6895821186222, 1590.3704726754975)],
                               [(1263.841977478558, 943.4546342496925),
                                (1262.3235250519404, 944.6782247770625),
                                (1260.8163412684773, 945.9156425437817),
                                (1259.3272449205788, 947.1747683330505)]])
        self._do_test('singapore-hollandvillage', predictions, 0)


import numpy as np
import random

from railrl.data_management.env_replay_buffer import EnvReplayBuffer
from railrl.data_management.segment_tree import SumSegmentTree, MinSegmentTree
from railrl.misc.np_util import truncated_geometric

class HerReplayBuffer(EnvReplayBuffer):
    """
    Save goals from the same trajectory into the replay buffer.
    Only add_path is implemented.

    Implementation details:
     - Every sample from [0, self._size] will be valid.
    """
    def __init__(
            self,
            max_size,
            env,
            num_goals_to_sample=4,
            fraction_goals_are_rollout_goals=None,
            resampling_strategy='uniform',
            truncated_geom_factor=1.,
    ):
        """

        :param max_size:
        :param observation_dim:
        :param action_dim:
        :param num_goals_to_sample:
        :param fraction_goals_are_rollout_goals:
        :param resampling_strategy: How to resample states from the rest of
        the trajectory?
        - 'uniform': Sample them uniformly
        - 'truncated_geometric': Used a truncated geometric distribution
        """
        raise DeprecationWarning(
            "Are you sure you're not supposed to use ObsDictRelabelingBuffer?"
        )
        assert resampling_strategy in [
            'uniform',
            'truncated_geometric',
        ]
        super().__init__(max_size, env)
        self.num_goals_to_sample = num_goals_to_sample
        self._goals = np.zeros((max_size, self.env.goal_dim))
        self._num_steps_left = np.zeros((max_size, 1))
        if fraction_goals_are_rollout_goals is None:
            fraction_goals_are_rollout_goals = (
                1. / num_goals_to_sample
            )
        self.fraction_goals_are_rollout_goals = (
            fraction_goals_are_rollout_goals
        )
        self.truncated_geom_factor = float(truncated_geom_factor)
        self.resampling_strategy = resampling_strategy

        # Let j be any index in self._idx_to_future_obs_idx[i]
        # Then self._next_obs[j] is a valid next observation for observation i
        self._idx_to_future_obs_idx = [None] * max_size

    def add_sample(self, observation, action, reward, terminal,
                   next_observation, **kwargs):
        raise NotImplementedError("Only use add_path")

    def add_path(self, path):
        obs = path["observations"]
        actions = path["actions"]
        rewards = path["rewards"]
        next_obs = path["next_observations"]
        terminals = path["terminals"]
        goals = path["goals"]
        num_steps_left = path["num_steps_left"]
        path_len = len(rewards)

        actions = flatten_n(actions)
        obs = flatten_n(obs)
        next_obs = flatten_n(next_obs)

        if self._top + path_len >= self._max_replay_buffer_size:
            num_pre_wrap_steps = self._max_replay_buffer_size - self._top
            # numpy slice
            pre_wrap_buffer_slice = np.s_[
                self._top:self._top + num_pre_wrap_steps, :
            ]
            pre_wrap_path_slice = np.s_[0:num_pre_wrap_steps, :]

            num_post_wrap_steps = path_len - num_pre_wrap_steps
            post_wrap_buffer_slice = slice(0, num_post_wrap_steps)
            post_wrap_path_slice = slice(num_pre_wrap_steps, path_len)
            for buffer_slice, path_slice in [
                (pre_wrap_buffer_slice, pre_wrap_path_slice),
                (post_wrap_buffer_slice, post_wrap_path_slice),
            ]:
                self._observations[buffer_slice] = obs[path_slice]
                self._actions[buffer_slice] = actions[path_slice]
                self._rewards[buffer_slice] = rewards[path_slice]
                self._next_obs[buffer_slice] = next_obs[path_slice]
                self._terminals[buffer_slice] = terminals[path_slice]
                self._goals[buffer_slice] = goals[path_slice]
                self._num_steps_left[buffer_slice] = num_steps_left[path_slice]
            # Pointers from before the wrap
            for i in range(self._top, self._max_replay_buffer_size):
                self._idx_to_future_obs_idx[i] = np.hstack((
                    # Pre-wrap indices
                    np.arange(i, self._max_replay_buffer_size),
                    # Post-wrap indices
                    np.arange(0, num_post_wrap_steps)
                ))
            # Pointers after the wrap
            for i in range(0, num_post_wrap_steps):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i,
                    num_post_wrap_steps,
                )
        else:
            slc = np.s_[self._top:self._top + path_len, :]
            self._observations[slc] = obs
            self._actions[slc] = actions
            self._rewards[slc] = rewards
            self._next_obs[slc] = next_obs
            self._terminals[slc] = terminals
            self._goals[slc] = goals
            self._num_steps_left[slc] = num_steps_left
            for i in range(self._top, self._top + path_len):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i, self._top + path_len
                )
        self._top = (self._top + path_len) % self._max_replay_buffer_size
        self._size = min(self._size + path_len, self._max_replay_buffer_size)

    def _sample_indices(self, batch_size):
        return np.random.randint(0, self._size, batch_size)

    def random_batch(self, batch_size):
        indices = self._sample_indices(batch_size)
        next_obs_idxs = []
        for i in indices:
            possible_next_obs_idxs = self._idx_to_future_obs_idx[i]
            # This is generally faster than random.choice. Makes you wonder what
            # random.choice is doing
            num_options = len(possible_next_obs_idxs)
            if num_options == 1:
                next_obs_i = 0
            else:
                if self.resampling_strategy == 'uniform':
                    next_obs_i = int(np.random.randint(0, num_options))
                elif self.resampling_strategy == 'truncated_geometric':
                    next_obs_i = int(truncated_geometric(
                        p=self.truncated_geom_factor/num_options,
                        truncate_threshold=num_options-1,
                        size=1,
                        new_value=0
                    ))
                else:
                    raise ValueError("Invalid resampling strategy: {}".format(
                        self.resampling_strategy
                    ))
            next_obs_idxs.append(possible_next_obs_idxs[next_obs_i])
        next_obs_idxs = np.array(next_obs_idxs)
        resampled_goals = self.env.convert_obs_to_goals(
            self._next_obs[next_obs_idxs]
        )
        num_goals_are_from_rollout = int(
            batch_size * self.fraction_goals_are_rollout_goals
        )
        if num_goals_are_from_rollout > 0:
            resampled_goals[:num_goals_are_from_rollout] = self._goals[
                indices[:num_goals_are_from_rollout]
            ]
        return dict(
            observations=self._observations[indices],
            actions=self._actions[indices],
            rewards=self._rewards[indices],
            terminals=self._terminals[indices],
            next_observations=self._next_obs[indices],
            goals_used_for_rollout=self._goals[indices],
            resampled_goals=resampled_goals,
            num_steps_left=self._num_steps_left[indices],
            indices=np.array(indices).reshape(-1, 1),
        )

class PrioritizedHerReplayBuffer(HerReplayBuffer):
    """
    Building off of openai baselines code
    """
    def __init__(
            self,
            max_size,
            env,
            alpha=0.6,
            eps=1e-6,
            **kwargs
    ):
        super().__init__(max_size, env, **kwargs)
        assert alpha > 0
        self._alpha = alpha
        self._eps = eps

        it_capacity = 1
        while it_capacity < max_size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def add_path(self, path):
        """
        This implementation is the same, but I added a few:
        ```
        self._it_sum[i] = self._max_priority ** self._alpha
        self._it_min[i] = self._max_priority ** self._alpha
        ```
        """
        obs = path["observations"]
        actions = path["actions"]
        rewards = path["rewards"]
        next_obs = path["next_observations"]
        terminals = path["terminals"]
        goals = path["goals"]
        num_steps_left = path["num_steps_left"]
        path_len = len(rewards)

        actions = flatten_n(actions)
        obs = flatten_n(obs)
        next_obs = flatten_n(next_obs)

        if self._top + path_len >= self._max_replay_buffer_size:
            num_pre_wrap_steps = self._max_replay_buffer_size - self._top
            # numpy slice
            pre_wrap_buffer_slice = np.s_[
                                    self._top:self._top + num_pre_wrap_steps, :
                                    ]
            pre_wrap_path_slice = np.s_[0:num_pre_wrap_steps, :]

            num_post_wrap_steps = path_len - num_pre_wrap_steps
            post_wrap_buffer_slice = slice(0, num_post_wrap_steps)
            post_wrap_path_slice = slice(num_pre_wrap_steps, path_len)
            for buffer_slice, path_slice in [
                (pre_wrap_buffer_slice, pre_wrap_path_slice),
                (post_wrap_buffer_slice, post_wrap_path_slice),
            ]:
                self._observations[buffer_slice] = obs[path_slice]
                self._actions[buffer_slice] = actions[path_slice]
                self._rewards[buffer_slice] = rewards[path_slice]
                self._next_obs[buffer_slice] = next_obs[path_slice]
                self._terminals[buffer_slice] = terminals[path_slice]
                self._goals[buffer_slice] = goals[path_slice]
                self._num_steps_left[buffer_slice] = num_steps_left[path_slice]
            # Pointers from before the wrap
            for i in range(self._top, self._max_replay_buffer_size):
                self._idx_to_future_obs_idx[i] = np.hstack((
                    # Pre-wrap indices
                    np.arange(i, self._max_replay_buffer_size),
                    # Post-wrap indices
                    np.arange(0, num_post_wrap_steps)
                ))
                self._it_sum[i] = self._max_priority ** self._alpha
                self._it_min[i] = self._max_priority ** self._alpha
            # Pointers after the wrap
            for i in range(0, num_post_wrap_steps):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i,
                    num_post_wrap_steps,
                )
                self._it_sum[i] = self._max_priority ** self._alpha
                self._it_min[i] = self._max_priority ** self._alpha
        else:
            slc = np.s_[self._top:self._top + path_len, :]
            self._observations[slc] = obs
            self._actions[slc] = actions
            self._rewards[slc] = rewards
            self._next_obs[slc] = next_obs
            self._terminals[slc] = terminals
            self._goals[slc] = goals
            self._num_steps_left[slc] = num_steps_left
            for i in range(self._top, self._top + path_len):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i, self._top + path_len
                )
                self._it_sum[i] = self._max_priority ** self._alpha
                self._it_min[i] = self._max_priority ** self._alpha
        self._top = (self._top + path_len) % self._max_replay_buffer_size
        self._size = min(self._size + path_len, self._max_replay_buffer_size)

    def random_batch(self, batch_size, beta=1):
        batch = super().random_batch(batch_size)
        weights = self._get_weights(batch['indices'], beta)
        batch['is_weights'] = weights
        return batch

    def _get_weights(self, idxes, beta):
        weights = []
        p_min = self._it_min.min() / self._it_sum.sum()
        max_weight = (p_min * self._size) ** (-beta)

        for idx in idxes:
            idx = int(idx)
            p_sample = self._it_sum[idx] / self._it_sum.sum()
            weight = (p_sample * self._size) ** (-beta)
            weights.append(weight / max_weight)
        return np.array(weights).reshape(-1, 1)

    def _sample_indices(self, batch_size):
        res = []
        for _ in range(batch_size):
            mass = random.random() * self._it_sum.sum(0, self._size - 1)
            idx = self._it_sum.find_prefixsum_idx(mass)
            res.append(idx)
        return res

    def update_priorities(self, idxes, priorities):
        """Update priorities of sampled transitions.
        sets priority of transition at index idxes[i] in buffer
        to priorities[i].
        Parameters
        ----------
        idxes: [int]
            List of idxes of sampled transitions
        priorities: [float]
            List of updated priorities corresponding to
            transitions at the sampled idxes denoted by
            variable `idxes`.
        """
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            idx = int(idx)
            priority = priority + self._eps
            assert priority > 0
            assert 0 <= idx < self._size
            self._it_sum[idx] = priority ** self._alpha
            self._it_min[idx] = priority ** self._alpha

            self._max_priority = max(self._max_priority, priority)

class SimplePrioritizedHerReplayBuffer(PrioritizedHerReplayBuffer):
    def __init__(
            self,
            max_size,
            env,
            max_time_to_next_goal=None,
            **kwargs
    ):
        super().__init__(max_size, env, **kwargs)
        if max_time_to_next_goal is not None:
            assert max_time_to_next_goal >= self.num_goals_to_sample
        self.max_time_to_next_goal = max_time_to_next_goal
        self.last_path_start_idx = None

    def add_sample(self, observation, action, reward, terminal,
                   next_observation, goal, num_steps_left, **kwargs):
        self._observations[self._top] = observation
        self._actions[self._top] = action
        self._rewards[self._top] = reward
        self._terminals[self._top] = terminal
        self._next_obs[self._top] = next_observation
        self._goals[self._top] = goal
        self._num_steps_left[self._top] = num_steps_left
        self._it_sum[self._top] = self._max_priority ** self._alpha
        self._it_min[self._top] = self._max_priority ** self._alpha
        self._advance()

    def add_path(self, path):
        self.last_path_start_idx = self._top
        obs = path["observations"]
        actions = path["actions"]
        rewards = path["rewards"]
        next_obs = path["next_observations"]
        terminals = path["terminals"]
        goals = path["goals"]
        num_steps_left = path["num_steps_left"]
        path_len = len(terminals)

        actions = flatten_n(actions)
        obs = flatten_n(obs)
        next_obs = flatten_n(next_obs)

        for i, (
            observation,
            action,
            reward,
            terminal,
            next_observation,
            goal,
            num_steps_left,
        ) in enumerate(zip(
            obs,
            actions,
            rewards,
            terminals,
            next_obs,
            goals,
            num_steps_left,
        )):
            num_goals_to_sample = min(self.num_goals_to_sample, path_len-i)
            if self.max_time_to_next_goal is None:
                max_i = path_len
            else:
                # Don't add +1 because we index into next_obs and not ob
                max_i = min(i + self.max_time_to_next_goal, path_len)
            if self.max_time_to_next_goal == self.num_goals_to_sample:
                goal_idxs = np.arange(
                    i,
                    min(i + self.num_goals_to_sample, path_len)
                )
            else:
                goal_idxs = np.random.randint(i, max_i, num_goals_to_sample)
            for goal_i in goal_idxs:
                # Add both to the replay buffer to keep it balanced between
                # HER data and real data.
                self.add_sample(observation, action, reward, terminal,
                                next_observation, goal, np.array([goal_i - i]))
                self.add_sample(observation, action, reward, terminal,
                                next_observation, next_obs[goal_i],
                                np.array([goal_i - i]))

    def random_batch(self, batch_size, beta=1.0):
        indices = self._sample_indices(batch_size)
        weights = self._get_weights(indices, beta)
        return dict(
            observations=self._observations[indices],
            actions=self._actions[indices],
            rewards=self._rewards[indices],
            terminals=self._terminals[indices],
            next_observations=self._next_obs[indices],
            goals=self._goals[indices],
            num_steps_left=self._num_steps_left[indices],
            indices=np.array(indices).reshape(-1, 1),
            is_weights=weights.reshape(-1, 1),
        )

    def most_recent_path_batch(self, beta=1.0):
        last_path_start_idx = self.last_path_start_idx
        # We just looped around the replay buffer
        if last_path_start_idx > self._top:
            last_path_start_idx = 0

        indices = np.arange(last_path_start_idx, self._top)
        weights = self._get_weights(indices, beta)
        return dict(
            observations=self._observations[indices],
            actions=self._actions[indices],
            rewards=self._rewards[indices],
            terminals=self._terminals[indices],
            next_observations=self._next_obs[indices],
            goals=self._goals[indices],
            num_steps_left=self._num_steps_left[indices],
            indices=np.array(indices).reshape(-1, 1),
            is_weights=weights.reshape(-1, 1),
        )

class SimpleHerReplayBuffer(EnvReplayBuffer):
    """
    Only add relabel goals when putting paths into the replay buffer,
    like in the orignial HER paper.
    """
    def __init__(
            self,
            max_size,
            env,
            num_goals_to_sample=4,
            **kwargs
    ):
        super().__init__(max_size, env, **kwargs)
        self.num_goals_to_sample = num_goals_to_sample
        self._goals = np.zeros((max_size, self.env.goal_dim))

    def add_path(self, path):
        self.last_path_start_idx = self._top
        obs = path["observations"]
        actions = path["actions"]
        rewards = path["rewards"]
        next_obs = path["next_observations"]
        terminals = path["terminals"]
        goals = path["goals"]
        env_infos = path["env_infos"]
        path_len = len(terminals)

        for i, (
                observation,
                action,
                reward,
                terminal,
                next_observation,
                goal,
                env_info,
        ) in enumerate(zip(
            obs,
            actions,
            rewards,
            terminals,
            next_obs,
            goals,
            env_infos,
        )):
            # It's not really necessary to recompute the reward, but just to
            # make sure I always use the same reward as in HER
            new_reward = self.env.compute_her_reward_np(
                observation,
                action,
                next_observation,
                goal,
                env_info
            )
            self.add_sample(
                observation, action, new_reward, terminal,
                next_observation, goal, env_info=env_info
            )
            num_goals_to_sample = min(self.num_goals_to_sample, path_len-i)
            max_i = path_len
            goal_idxs = np.random.randint(i, max_i, num_goals_to_sample)
            for goal_i in goal_idxs:
                # Add both to the replay buffer to keep it balanced between
                # HER data and real data.
                new_goal = self.env.convert_ob_to_goal(next_obs[goal_i])
                new_reward = self.env.compute_her_reward_np(
                    observation,
                    action,
                    next_observation,
                    new_goal,
                    env_info
                )
                self.add_sample(
                    observation, action, new_reward, terminal,
                    next_observation, new_goal, env_info=env_info
                )

    def add_sample(self, observation, action, reward, terminal,
                   next_observation, goal, **kwargs):
        self._goals[self._top] = goal
        super().add_sample(
            observation, action, reward, terminal, next_observation, **kwargs
        )

    def random_batch(self, batch_size):
        indices = np.random.randint(0, self._size, batch_size)
        batch = dict(
            observations=self._observations[indices],
            actions=self._actions[indices],
            rewards=self._rewards[indices],
            terminals=self._terminals[indices],
            next_observations=self._next_obs[indices],
            goals=self._goals[indices],
        )
        for key in self._env_info_keys:
            assert key not in batch.keys()
            batch[key] = self._env_infos[key][indices]
        return batch

class RelabelingReplayBuffer(EnvReplayBuffer):
    """
    Save goals from the same trajectory into the replay buffer.
    Only add_path is implemented.
    Implementation details:
     - Every sample from [0, self._size] will be valid.
    """
    def __init__(
            self,
            max_size,
            env,
            fraction_goals_are_rollout_goals=1.0, # default, no HER
            fraction_resampled_goals_are_env_goals=0.0, # this many goals are just sampled from environment directly
            resampling_strategy='uniform', # 'uniform' is the HER 'future' strategy
            truncated_geom_factor=1.,
            **kwargs
    ):
        """
        :param resampling_strategy: How to resample states from the rest of
        the trajectory?
        - 'uniform': Sample them uniformly
        - 'truncated_geometric': Used a truncated geometric distribution
        """
        assert resampling_strategy in [
            'uniform',
            'truncated_geometric',
        ]
        super().__init__(max_size, env, **kwargs)
        self._goals = np.zeros((max_size, self.env.goal_dim))
        self._num_steps_left = np.zeros((max_size, 1))
        self.fraction_goals_are_rollout_goals = fraction_goals_are_rollout_goals
        self.fraction_resampled_goals_are_env_goals = fraction_resampled_goals_are_env_goals
        self.truncated_geom_factor = float(truncated_geom_factor)
        self.resampling_strategy = resampling_strategy

        # Let j be any index in self._idx_to_future_obs_idx[i]
        # Then self._next_obs[j] is a valid next observation for observation i
        self._idx_to_future_obs_idx = [None] * max_size

    def add_sample(self, observation, action, reward, terminal,
                   next_observation, **kwargs):
        raise NotImplementedError("Only use add_path")

    def add_path(self, path):
        obs = path["observations"]
        actions = path["actions"]
        rewards = path["rewards"]
        next_obs = path["next_observations"]
        terminals = path["terminals"]
        goals = path["goals"]
        num_steps_left = path["rewards"].copy() # path["num_steps_left"] # irrelevant for non-TDM
        env_infos = np.array(path["env_infos"])
        path_len = len(rewards)

        actions = flatten_n(actions)
        obs = flatten_n(obs)
        next_obs = flatten_n(next_obs)
        env_infos = flatten_env_info(env_infos, self._env_info_keys)

        if self._top + path_len >= self._max_replay_buffer_size:
            num_pre_wrap_steps = self._max_replay_buffer_size - self._top
            # numpy slice
            pre_wrap_buffer_slice = np.s_[
                self._top:self._top + num_pre_wrap_steps, :
            ]
            pre_wrap_path_slice = np.s_[0:num_pre_wrap_steps, :]

            num_post_wrap_steps = path_len - num_pre_wrap_steps
            post_wrap_buffer_slice = slice(0, num_post_wrap_steps)
            post_wrap_path_slice = slice(num_pre_wrap_steps, path_len)
            for buffer_slice, path_slice in [
                (pre_wrap_buffer_slice, pre_wrap_path_slice),
                (post_wrap_buffer_slice, post_wrap_path_slice),
            ]:
                self._observations[buffer_slice] = obs[path_slice]
                self._actions[buffer_slice] = actions[path_slice]
                self._rewards[buffer_slice] = rewards[path_slice]
                self._next_obs[buffer_slice] = next_obs[path_slice]
                self._terminals[buffer_slice] = terminals[path_slice]
                self._goals[buffer_slice] = goals[path_slice]
                self._num_steps_left[buffer_slice] = num_steps_left[path_slice]
                for key in self._env_info_keys:
                    self._env_infos[key][buffer_slice] = env_infos[key][path_slice]
            # Pointers from before the wrap
            for i in range(self._top, self._max_replay_buffer_size):
                self._idx_to_future_obs_idx[i] = np.hstack((
                    # Pre-wrap indices
                    np.arange(i, self._max_replay_buffer_size),
                    # Post-wrap indices
                    np.arange(0, num_post_wrap_steps)
                ))
            # Pointers after the wrap
            for i in range(0, num_post_wrap_steps):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i,
                    num_post_wrap_steps,
                )
        else:
            slc = np.s_[self._top:self._top + path_len, :]
            self._observations[slc] = obs
            self._actions[slc] = actions
            self._rewards[slc] = rewards
            self._next_obs[slc] = next_obs
            self._terminals[slc] = terminals
            self._goals[slc] = goals
            self._num_steps_left[slc] = num_steps_left
            for key in self._env_info_keys:
                self._env_infos[key][slc] = env_infos[key]
            for i in range(self._top, self._top + path_len):
                self._idx_to_future_obs_idx[i] = np.arange(
                    i, self._top + path_len
                )
        self._top = (self._top + path_len) % self._max_replay_buffer_size
        self._size = min(self._size + path_len, self._max_replay_buffer_size)

    def _sample_indices(self, batch_size):
        return np.random.randint(0, self._size, batch_size)

    def random_batch(self, batch_size):
        indices = self._sample_indices(batch_size)
        next_obs_idxs = []
        for i in indices:
            possible_next_obs_idxs = self._idx_to_future_obs_idx[i]
            # This is generally faster than random.choice. Makes you wonder what
            # random.choice is doing
            num_options = len(possible_next_obs_idxs)
            if num_options == 1:
                next_obs_i = 0
            else:
                if self.resampling_strategy == 'uniform':
                    next_obs_i = int(np.random.randint(0, num_options))
                elif self.resampling_strategy == 'truncated_geometric':
                    next_obs_i = int(truncated_geometric(
                        p=self.truncated_geom_factor/num_options,
                        truncate_threshold=num_options-1,
                        size=1,
                        new_value=0
                    ))
                else:
                    raise ValueError("Invalid resampling strategy: {}".format(
                        self.resampling_strategy
                    ))
            next_obs_idxs.append(possible_next_obs_idxs[next_obs_i])
        next_obs_idxs = np.array(next_obs_idxs)
        resampled_goals = self.env.convert_obs_to_goals(
            self._next_obs[next_obs_idxs]
        )
        num_goals_are_from_rollout = int(
            batch_size * self.fraction_goals_are_rollout_goals
        )
        if num_goals_are_from_rollout > 0:
            resampled_goals[:num_goals_are_from_rollout] = self._goals[
                indices[:num_goals_are_from_rollout]
            ]
        # recompute rewards
        new_obs = self._observations[indices]
        new_next_obs = self._next_obs[indices]
        new_actions = self._actions[indices]
        new_rewards = self._rewards[indices].copy() # needs to be recomputed
        env_info_dicts = [self.rebuild_env_info_dict(idx) for idx in indices]
        random_numbers = np.random.rand(batch_size)
        for i in range(batch_size):
            if random_numbers[i] < self.fraction_resampled_goals_are_env_goals:
                resampled_goals[i, :] = self.env.sample_goal_for_rollout() # env_goals[i, :]

            new_reward = self.env.compute_her_reward_np(
                new_obs[i, :],
                new_actions[i, :],
                new_next_obs[i, :],
                resampled_goals[i, :],
                env_info_dicts[i],
            )
            new_rewards[i] = new_reward
        # new_rewards = self.env.computer_her_reward_np_batch(
        #     new_obs,
        #     new_actions,
        #     new_next_obs,
        #     resampled_goals,
        #     env_infos,
        # )

        batch = dict(
            observations=new_obs,
            actions=new_actions,
            rewards=new_rewards,
            terminals=self._terminals[indices],
            next_observations=new_next_obs,
            goals_used_for_rollout=self._goals[indices],
            resampled_goals=resampled_goals,
            num_steps_left=self._num_steps_left[indices],
            indices=np.array(indices).reshape(-1, 1),
            goals=resampled_goals,
        )
        for key in self._env_info_keys:
            assert key not in batch.keys()
            batch[key] = self._env_infos[key][indices]
        return batch

def flatten_n(xs):
    xs = np.asarray(xs)
    return xs.reshape((xs.shape[0], -1))

def flatten_env_info(env_infos, env_info_keys):
# Turns list of env_info dicts into env_info dict of 2D np arrays
    return {
        key: flatten_n(
                [env_infos[i][key] for i in range(len(env_infos))]
        )
        for key in env_info_keys
    }


import numpy as np
import pytest as pt
from collections import OrderedDict

def test_morganfp():
    pt.importorskip("rdkit")
    from chemreps.fingerprints import morganfp
    rep = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    fname = 'data/sdf/butane.sdf'
    mfp = morganfp(fname, radius=2, nBits=1024)

    assert np.all(np.abs(mfp-rep) <= 1e-4) == True

    with pt.raises(NotImplementedError):
        rep = morganfp('data/xyz/butane.xyz')

    with pt.raises(NotImplementedError):
        rep = morganfp('data/cml/butane.cml')

    with pt.raises(NotImplementedError):
        rep = morganfp('data/cclib/butane.cclib')


# Video handling utilities.
#
# Two things in this file
# 1. Kelvinator - simple method of extracting key frames for a video.
# 2. Resize - simple transcoder to create smaller versions of videos.
#
# Requirements:
# 1. ffmpeg
# 2. x264 for transcoding
# 3. Python Imaging: Install PIL
# 4. Python linear algebra: numpy
#
# For more info on transcoding, see:
# - http://ffmpeg.org/trac/ffmpeg/wiki/x264EncodingGuide
# size abbreviations:
# - http://linuxers.org/tutorial/how-extract-images-video-using-ffmpeg

try:
    import Image
except ImportError, msg:
    Image = False
try:
    import numpy as np
except ImportError, msg:
    np = False
import math
import os
import operator
import shutil
import subprocess
import sys

from django.core.files.storage import default_storage
from django.conf import settings
from celery import task

from utility import *

##  KELVINATOR
##
def extract(notify_buf, working_dir, jpeg_dir, video_file, start_offset, extraction_frame_rate):
    infoLog(notify_buf, "Kicking off ffmpeg, hold onto your hats")
    cmdline = [ ffmpeg_cmd() ]
    cmdline += \
        [ '-i', working_dir + "/" + video_file, # input
          '-ss', str(start_offset),             # start a few seconds late
          '-r', str(extraction_frame_rate),     # thumbs per second to extract
          '-f', 'image2',                       # thumb format
          jpeg_dir + '/img%5d.jpeg',            # thumb filename template
        ]
    infoLog(notify_buf, "EXTRACT: " + " ".join(cmdline))
    returncode = subprocess.call(cmdline)

    if returncode == 0:
        infoLog(notify_buf, "ffmpeg completed, returncode %d" % returncode)
    else:
        errorLog(notify_buf, "ffmpeg completed, returncode %d" % returncode)
        cleanup_working_dir(notify_buf, working_dir)
        raise VideoError("ffmpeg error %d" % returncode)

def difference(notify_buf, working_dir, jpeg_dir, extraction_frame_rate, frames_per_minute_target):
    """Select targeted number of slides from candidate set of frames"""
    # It is an error to call this method without having numpy installed
    if not np:
        raise VideoError("Could not extract keyframes from video file without numpy")
    if not Image:
        raise VideoError("Could not extract keyframes from video file without PIL (Image)")

    # from http://mail.python.org/pipermail/image-sig/1997-March/000223.html
    def computeDiff(file1, file2):
        h1 = Image.open(file1).histogram()
        h2 = Image.open(file2).histogram()
        rms = math.sqrt(reduce(operator.add, map(lambda a,b: (a-b)**2, h1, h2))/len(h1))
        return rms

    # get local maximum values from an array, which are also larger the threshold.
    # local maximum is extracted by comparing current with three neighbors in both directions.
    def localMaximum(candidates, threshold):
        cuts = [];
        for i in range(3, len(candidates)-4):
            cur_score = candidates[i][0]
            if cur_score > threshold and cur_score > candidates[i-1][0] and cur_score > candidates[i-2][0] and cur_score > candidates[i-3][0] and cur_score > candidates[i+1][0] and cur_score > candidates[i+2][0] and cur_score > candidates[i+3][0] :
                cuts.append(candidates[i])
        return cuts

    image_list = os.listdir(jpeg_dir)
    if len(image_list) == 0:
        cleanup_working_dir(notify_buf, working_dir)
        raise VideoError("Failed to extract keyframes from video file")

    image_list.sort()
    infoLog(notify_buf, "Extraction frame rate: %d fps" % extraction_frame_rate)
    duration = len(image_list)/extraction_frame_rate    # in seconds
    infoLog(notify_buf, "Video duration: %d seconds" % duration)
    infoLog(notify_buf, "Initial keyframes: %d" % len(image_list))
    infoLog(notify_buf, "Target keyframes per minute: %d" % frames_per_minute_target)
    max_keyframes = int(math.ceil(frames_per_minute_target * duration/60.0))
    infoLog(notify_buf, "Upper bound on number of keyframes kelvinator will output: %d" % max_keyframes)
    infoLog(notify_buf, "Internal differencing threshold: average of all scores")
    image_num = len(image_list)

    # window size, all the frames in the window centered at
    # current frame are used to determine shot boundary.
    # this value doesn't need to be changed for different videos.
    k = 5

    # calculate difference matrix
    difference_matrix = np.zeros((image_num,image_num))
    for i in range(0, image_num-1-k):
        for j in range(i+1, i+k):
            difference_matrix[i, j] = computeDiff(jpeg_dir+"/"+image_list[i], jpeg_dir+"/"+image_list[j])
    difference_matrix = difference_matrix + difference_matrix.transpose()

    # callate shot boundary scores for each frames, score = cut(A,B)/associate(A) + cut(A,B)/associate(B)
    candidates = []
    for i in range(k, image_num-1-k):
        cutAB = np.sum(difference_matrix[i-k:i, i:i+k])
        assocA = np.sum(difference_matrix[i-k:i, i-k:i])
        assocB = np.sum(difference_matrix[i:i+k, i:i+k])
        if (assocA!=0) and (assocB!=0):
            score = cutAB/assocA + cutAB/assocB
        else:
            score = 0
        candidates.append((score, i))

    # extract local maximum as the shot boundaries.
    # the threshold is assigned to be the mean of all the scores.
    # [important] we may want to change the threshold to control the number of key frames.
    # higher threshold generates fewer shot boundaries.
    threshold = np.mean([pair[0] for pair in candidates]);
    cuts = localMaximum(candidates, threshold)

    # limit shot boundary number fewer than max_keyframes
    if len(cuts) >= max_keyframes :
        # sort key frames by score
        cuts.sort(reverse=True)
        cuts = cuts[:max_keyframes];

    # select the 3nd frame after each shot boundary as the key frame.
    # alternatively, we can also select middle frame between two shot boundaries as the key frame.
    cut_offset = 2

    # below code is used to output keyframes.
    keep_frames = []
    keep_times = []
    jpeg_dir_parent = os.path.abspath(os.path.join(jpeg_dir, os.path.pardir))
    jpeg_dir_result = jpeg_dir_parent + "/tmp"
    if os.path.exists(jpeg_dir_result):
        pass
    else:
        os.mkdir(jpeg_dir_result)
    # sort key frames by index
    sorted(cuts, key=lambda x: x[1])

    # move key frames into a tmp folder, then move back.
    for i in range(len(cuts)):
        index = min(cuts[i][1] + cut_offset, len(image_list) - 1)
        shutil.move(jpeg_dir+"/"+image_list[index], jpeg_dir_result)
        keep_frames.append(image_list[index])
        keep_times.append(index)

    cleanup_working_dir(notify_buf, jpeg_dir)
    os.rename(jpeg_dir_result, jpeg_dir)

    return (keep_frames, keep_times)

def write_manifest(notify_buf, jpeg_dir, keep_frames, keep_times):
    outfile_name = jpeg_dir + "/manifest.txt"
    outfile = open(outfile_name, 'w')
    index = 0
    outfile.write("{")

    while(index < len(keep_frames)):
        outfile.write("\"")
        toWrite = str(keep_times[index])
        outfile.write(toWrite)
        outfile.write("\":{\"imgsrc\":\"")
        outfile.write(keep_frames[index])
        outfile.write("\"}")
        if index < len(keep_frames)-1:
            outfile.write(",")
        index += 1
    outfile.write("}")
    outfile.close()

def put_thumbs(notify_buf, jpeg_dir, prefix, suffix, video_id, store_loc):
    # I wish Filesystem API worked the same for local and remote, but it don't
    if store_loc == 'local':
        root = getattr(settings, 'MEDIA_ROOT')
        store_path = root + "/" + prefix + "/" + suffix + "/videos/" + str(video_id) + "/jpegs"
        if default_storage.exists(store_path):
            infoLog(notify_buf, "Found prior directory, removing: %s" % store_path)
            dirRemove(store_path)
        os.mkdir(store_path)
    else:
        store_path = prefix + "/" + suffix + "/videos/" + str(video_id) + "/jpegs"
        default_storage.delete(store_path)

    # not doing write to tmp and then mv because the file storage API limitation
    image_list = os.listdir(jpeg_dir)
    image_list.sort()
    for fname in image_list:
        infoLog(notify_buf, "Uploading: %s" % fname)
        local_file = open(jpeg_dir + "/" + fname, 'rb')
        store_file = default_storage.open(store_path + "/" + fname, 'wb')
        file_data = local_file.read()
        store_file.write(file_data)
        local_file.close()
        store_file.close()
    infoLog(notify_buf, "Uploaded: %s files" % str(len(image_list)))

# Main Kelvinator Task (CELERY)
@task()
def kelvinate(store_path_raw, frames_per_minute_target=2, notify_addr=None):
    """
    Given a path to a video in a readable S3 bucket, extract the frames and
    upload back to S3.

    store_path must be the full path to the video file, not just to its parent folder.
    """

    notify_buf = []

    if not np:
        sys.stderr.write("WARNING\nThe python library 'numpy' is not installed, so kelvinate function cannot run.\n")
        infoLog(notify_buf, "Kelvinate terminated: numpy not installed.")
        return False

    infoLog(notify_buf, "Kelvinate: extracting %s" % store_path_raw)

    extraction_frame_rate = 1   # seconds
    start_offset = 3            # seconds

    frames_per_minute_target = float(frames_per_minute_target)

    (store_path, course_prefix, course_suffix, video_id, video_filename) = splitpath(store_path_raw)

    store_loc = "remote"
    if getattr(settings, 'AWS_ACCESS_KEY_ID') == "local":
        store_loc = "local"

    work_dir = None
    try:
        (work_dir, jpegs) = create_working_dirs("kelvinator", notify_buf, "jpegs")
        get_video(notify_buf, work_dir, video_filename, store_path)
        extract(notify_buf, work_dir, jpegs, video_filename, start_offset, extraction_frame_rate)
        (thumbs, times) = difference(notify_buf, work_dir, jpegs, extraction_frame_rate, frames_per_minute_target)
        write_manifest(notify_buf, jpegs, thumbs, times)
        put_thumbs(notify_buf, jpegs, course_prefix, course_suffix, video_id, store_loc)
    except:
        if work_dir: cleanup_working_dir(notify_buf, work_dir)
        notify("Kelvinator", notify_buf, notify_addr,
                course_prefix, course_suffix, video_filename, store_path) # in utilities.py
        raise

    cleanup_working_dir(notify_buf, work_dir)
    notify("Kelvinator", notify_buf, notify_addr, course_prefix, course_suffix, video_filename, store_path)

##
##  RESIZE

# video sizes we support: key is size name (used for target subdirectory) and value
# are the parameters (as a list) that we'll pass to ffmpeg.
sizes = { "large":  [ "-crf", "23", "-s", "1280x720" ],   # original size, compressed
          "medium": [ "-crf", "27", "-s", "852x480" ],    # wvga at 16:9
          "small":  [ "-crf", "30", "-s", "640x360" ],
          "tiny":   [ "-crf", "40", "-s", "320x180" ],
        }

# Actually transcode the video down
def do_resize(notify_buf, working_dir, target_dir, video_file, target_size):
    cmdline = [ ffmpeg_cmd() ]
    cmdline += [ "-i", working_dir + "/" + video_file,  # infile
                 "-c:v", "libx264",          # video codec
                 "-profile:v", "baseline",   # most compatible
                 "-strict", "-2",            # magic to allow aac audio enc
               ]
    cmdline += sizes[target_size]
    cmdline += [ target_dir + "/" + video_file ]  # outfile

    infoLog(notify_buf, "RESIZE: " + " ".join(cmdline))
    returncode = subprocess.call(cmdline)

    if returncode == 0:
        infoLog(notify_buf, "completed with returncode %d" % returncode)
    else:
        errorLog(notify_buf, "completed with returncode %d" % returncode)
        cleanup_working_dir(notify_buf, working_dir)
        raise VideoError("do_resize error %d" % returncode)

def upload(notify_buf, target_dir, target_part, prefix, suffix, video_id, video_file, store_loc):
    # I wish Filesystem API worked the same for local and remote, but it don't
    if store_loc == 'local':
        root = getattr(settings, 'MEDIA_ROOT')
        store_path = root + "/" + prefix + "/" + suffix + "/videos/" + str(video_id) + "/" + target_part
        if default_storage.exists(store_path):
            infoLog(notify_buf, "Found prior directory, removing: %s" % store_path)
            dirRemove(store_path)
        os.mkdir(store_path)
    else:
        store_path = prefix + "/" + suffix + "/videos/" + str(video_id) + "/" + target_part
        default_storage.delete(store_path)

    statinfo = os.stat(target_dir + "/" + video_file)
    infoLog(notify_buf, "Final file size: %s" % str(statinfo.st_size))

    local_file = open(target_dir + "/" + video_file, 'rb')
    store_file = default_storage.open(store_path + "/" + video_file, 'wb')
    store_file.write(local_file.read())
    local_file.close()
    store_file.close()

# Main Resize Task (CELERY)
@task()
def resize(store_path_raw, target_raw, notify_addr=None):
    """
    Given a video path, scale it down and save the result alongside the original
    video.  So we can provide different download options.

    Preset is either "small" or "large".
    """

    notify_buf = []
    infoLog(notify_buf, "Resize: converting %s version of %s" % (target_raw, store_path_raw))

    target = target_raw.lower()
    if target not in sizes.keys():
        VideoError("Target size \"%s\" not supported" % target)

    (store_path, course_prefix, course_suffix, video_id, video_file) = splitpath(store_path_raw)

    store_loc = 'remote'
    if getattr(settings, 'AWS_ACCESS_KEY_ID') == 'local':
        store_loc = 'local'

    work_dir = None
    try:
        (work_dir, smaller_dir) = create_working_dirs("resize", notify_buf, target)
        get_video(notify_buf, work_dir, video_file, store_path)
        do_resize(notify_buf, work_dir, smaller_dir, video_file, target)
        upload(notify_buf, smaller_dir, target, course_prefix, course_suffix, video_id, video_file, store_loc)
    except:
        if work_dir: cleanup_working_dir(notify_buf, work_dir)
        notify("Resize (%s)" % target, notify_buf, notify_addr, course_prefix, course_suffix,
               video_file, store_path)
        raise

    cleanup_working_dir(notify_buf, work_dir)
    notify("Resize (%s)" % target, notify_buf, notify_addr, course_prefix, course_suffix,
           video_file, store_path)


#!/usr/bin/env python3

import argparse
import os
import sys
import time
import glob
import json
import cv2
from pprint import pprint
import numpy as np

import tensorflow as tf
import tensorlayer as tl

## COCO imports
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from pycocotools.cocoanalyze import COCOanalyze

sys.path.append('../../')

from openpose_plus.inference.common import measure, plot_humans, read_imgfile, CocoPart
from openpose_plus.inference.estimator import TfPoseEstimator
from openpose_plus.models import get_model
from pycocotools.coco import COCO

tf.logging.set_verbosity(tf.logging.INFO)
tl.logging.set_verbosity(tl.logging.INFO)

def round_int(val):
    return int(round(val))

def write_json(img_name, humans, json_out):

    coco_ids = [0, 15, 14, 17, 16, 5, 2, 6, 3, 7, 4, 11, 8, 12, 9, 13, 10]

    image_h, image_w = cv2.imread(img_name).shape[:2]

    image_id = os.path.basename(img_name).replace('.jpg', '').replace('COCO_val2014_', '').replace('COCO_train2014_', '')
    image_id = int(image_id.lstrip('0'))

    for h in humans:

        new_keypoints = []

        for i in coco_ids:
            if i not in h.body_parts.keys():
                new_keypoints.append(0)
                new_keypoints.append(0)
                new_keypoints.append(0)
            else:
                new_keypoints.append(round_int(h.body_parts[i].x * image_w))
                new_keypoints.append(round_int(h.body_parts[i].y * image_h))
                new_keypoints.append(2)

        json_out.append({"image_id": image_id,
                         "category_id" : 1,
                         "keypoints" : new_keypoints,
                         "score" : h.score})

def inference(base_model_name, path_to_npz, data_format, input_files, plot, writejson, outputdir, preExp13, is_resize = False):

    NumberOfIteration = path_to_npz[int(path_to_npz.rfind('pose'))+4:int(path_to_npz.rfind('.npz'))]

    model_func = get_model(base_model_name, is_resize)

    if is_resize:
        #Recommends : 368x432 or 368x656 or 736x1312
        width, height = (432,368)
    else:
        width, height = (None,None)

    e = measure(lambda: TfPoseEstimator(path_to_npz, model_func, target_size=(width, height), data_format=data_format),
                    'create TfPoseEstimator')
#    var_vgg19_stage1_kernel = []
#    var_vgg19_stage1_bias = []
#    var_stage2_kernel = []
#    var_stage2_bias = []
#
#    for variable in tf.trainable_variables():
#
#
#        # stage = 1 and vgg layer
#        if variable.name.find('vgg19') != -1 or variable.name.find('stage1') != -1:
#
#            if variable.name.find('kernel') != -1:
#                var_vgg19_stage1_kernel.append(variable)
#            elif variable.name.find('bias') != -1:
#                var_vgg19_stage1_bias.append(variable)

#        # stage > 1
#        elif variable.name.find('stage2') != -1 or variable.name.find('stage3') != -1 or variable.name.find('stage4') != -1 or variable.name.find('stage5') != -1 or variable.name.find('stage6') != -1:
#            if variable.name.find('kernel') != -1:
#                var_stage2_kernel.append(variable)
#            elif variable.name.find('bias') != -1:
#                var_stage2_bias.append(variable)

#    print('var_vgg19_stage1_kernel')
#    pprint(var_vgg19_stage1_kernel)
#    print('var_vgg19_stage1_bias')
#    pprint(var_vgg19_stage1_bias)
#    print('var_stage2_kernel')
#    pprint(var_stage2_kernel)
#    print('var_stage2_bias')
#    pprint(var_stage2_bias)

    #print([n.name for n in tf.get_default_graph().as_graph_def().node])
    t0 = time.time()

    json_out = []

    for idx, img_name in enumerate(input_files):

        image, orig_image = measure(lambda: read_imgfile(img_name, width, height, preExp13=preExp13, data_format=data_format), 'read_imgfile')

        if orig_image is None:
            sys.exit(-1)

        humans, heatMap, pafMap = measure(lambda: e.inference(image), 'e.inference')

        # for logging purposes
        tl.logging.info('got %d humans from %s' % (len(humans), img_name))
        if humans:
            for h in humans:
                tl.logging.debug(h)
        if plot:
            if data_format == 'channels_first':
                image = image.transpose([1, 2, 0])
                orig_image = orig_image.transpose([1, 2, 0])
            plot_humans(orig_image, heatMap, pafMap, humans, '%02d' % (idx + 1), os.path.basename(img_name), outputdir)
        if writejson:
            write_json(img_name, humans, json_out)

    if writejson:
        with open(outputdir+'/json_detection_'+NumberOfIteration+'.json', 'w') as f:
            json.dump(json_out, f, ensure_ascii=False)

    tot = time.time() - t0
    mean = tot / len(input_files)
    tl.logging.info('inference all took: %f, mean: %f, FPS: %f' % (tot, mean, 1.0 / mean))

    return outputdir+'/json_detection_'+NumberOfIteration+'.json'

def parse_args():
    parser = argparse.ArgumentParser(description='inference')
    parser.add_argument('--path-to-npz', type=str, default='', help='path to npz', required=True)
    parser.add_argument('--images', type=str, default='', help='comma separate list of image filenames', required=True)
    parser.add_argument('--base-model', type=str, default='vgg', help='vgg19 | vgg | vggtiny | mobilenet')
    parser.add_argument('--data-format', type=str, default='channels_last', help='channels_last | channels_first.')
    parser.add_argument('--plot', type=bool, default=False, help='draw the results')
    parser.add_argument('--repeat', type=int, default=1, help='repeat the images for n times for profiling.')
    parser.add_argument('--limit', type=int, default=100, help='max number of images.')
    parser.add_argument('--writejson', type=bool, default=False, help='write results to json')
    parser.add_argument('--outputdir', type=str, default='vis', help='write results to out dir')
    parser.add_argument('--gpu', type=str, default='0', help='write results to out dir')
    parser.add_argument('--preExp13', type=str, default='None', help='Type True or False')
    parser.add_argument('--iter', type=int, default=0, help='iteration to start inference')

    return parser.parse_args()

def main():
    args = parse_args()

    annFile  = args.images

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    # if os.path.isdir(args.outputdir):
    # 	for f in glob.glob(os.path.join(args.outputdir,'*')):
    # 		os.remove(f)
    # else:
    # 	os.mkdir(args.outputdir)

    if args.images.endswith(('.jpg', '.JPG', '.png', '.PNG')):
        image_files = ([f for f in args.images.split(',') if f] * args.repeat)[:args.limit]
        inference(args.base_model, args.path_to_npz, args.data_format, image_files, args.plot, args.writejson, args.outputdir.strip('/'), args.preExp13)

    elif os.path.isdir(args.images):
        image_files = []
        for image in glob.glob(os.path.join(args.images, '*.png')):
            image_files.append(image)
        for image in glob.glob(os.path.join(args.images, '*.PNG')):
            image_files.append(image)
        for image in glob.glob(os.path.join(args.images, '*.jpg')):
            image_files.append(image)
        for image in glob.glob(os.path.join(args.images, '*.JPG')):
            image_files.append(image)

        inference(args.base_model, args.path_to_npz, args.data_format, image_files, args.plot, args.writejson, args.outputdir.strip('/'), args.preExp13)

    elif os.path.isfile(args.images):
        image_files = []

        if args.images.endswith('.json'):
            print('Running Coco dataset.')
            dataDir  = os.path.dirname(os.path.dirname(args.images))#'/home/std/data/work/Dennis/openpose-plus/data/mscoco2017'
            dataType = os.path.basename(args.images).split('_')[-1].replace('.json', '')#'train2017'
            annType  = os.path.basename(args.images)[:os.path.basename(args.images).rfind('_')] #'person_keypoints'

            annFile  = '%s/annotations/%s_%s.json'%(dataDir, annType, dataType)

            print("{:10}[{}]".format('annFile:',annFile))

            cocoGt = COCO(annFile)
            catIds = cocoGt.getCatIds(catNms=['person'])
            keys = cocoGt.getImgIds(catIds=catIds)

            for annid in keys:

                img_meta = cocoGt.loadImgs(annid)[0]

                if dataType.endswith('2017'):
                    #filename = str(annid).zfill(12)+'.jpg'
                    #image_files.append('{}/{}/{}'.format(dataDir, dataType, filename))

                    image_files.append('{}/{}/{}'.format(dataDir, dataType, img_meta['file_name']))
                elif dataType.endswith('2014'):
                    if dataType.startswith('val'):
                        filename = img_meta['file_name']#'COCO_val2014_'+str(annid).zfill(12)+'.jpg'
                    elif dataType.startswith('train'):
                        filename = img_meta['file_name']#'COCO_train2014_'+str(annid).zfill(12)+'.jpg'

                    image_files.append('{}/{}/{}'.format(dataDir, dataType, filename))

            if os.path.isdir(args.path_to_npz):
                list_of_npz = sorted(glob.glob(args.path_to_npz+'*.npz'), key=os.path.getmtime)
                for npz_file in list_of_npz:

                    iteration = int(npz_file[int(npz_file.rfind('pose'))+4:int(npz_file.rfind('.npz'))])
                    if iteration >= args.iter and iteration%2000 == 0:

                        resFile = inference(args.base_model, npz_file, args.data_format, image_files, args.plot, args.writejson, args.outputdir.strip('/'),args.preExp13)

                        ## load ground truth annotations
                        coco_gt = COCO(annFile)
                        catIds = coco_gt.getCatIds(catNms=['person'])
                        keys = coco_gt.getImgIds(catIds=catIds)

                        ## initialize COCO detections api
                        coco_dt   = coco_gt.loadRes(resFile)

                        ## initialize COCO eval api
                        coco_analyze = COCOanalyze(coco_gt, coco_dt, 'keypoints')
                        coco_analyze.evaluate(verbose=True, makeplots=True, savedir=resFile[:int(resFile.rfind('json_detection'))], team_name=str(iteration))

            else:
                inference(args.base_model, args.path_to_npz, args.data_format, image_files, args.plot, args.writejson, args.outputdir.strip('/'), args.preExp13)

        elif args.images.endswith('.txt'):

            with open(images, 'r') as textfile:
                for line in textfile:
                    image_files.append(line.strip())
            inference(args.base_model, args.path_to_npz, args.data_format, image_files, args.plot, args.writejson, args.outputdir.strip('/'), args.preExp13)

        else:
            print('Wrong file format input. Only accepts .json or .txt file')

if __name__ == '__main__':
    measure(main)