"""
Molecule
========

#. :class:`.Molecule`

Molecule class for optimisation.

"""

import numpy as np
import networkx as nx

class Molecule:
    """
    Molecule to optimize.

    """

    def __init__(self, atoms, bonds, position_matrix):
        """
        Initialize a :class:`Molecule` instance.

        Parameters
        ----------
        atoms : :class:`iterable` of :class:`.Atom`
            Atoms that define the molecule.

        bonds : :class:`iterable` of :class:`.Bond`
            Bonds between atoms that define the molecule.

        position_matrix : :class:`numpy.ndarray`
            A ``(n, 3)`` matrix holding the position of every atom in
            the :class:`.Molecule`.

        """

        self._atoms = tuple(atoms)
        self._bonds = tuple(bonds)
        self._position_matrix = np.array(
            position_matrix.T,
            dtype=np.float64,
        )

    def get_position_matrix(self):
        """
        Return a matrix holding the atomic positions.

        Returns
        -------
        :class:`numpy.ndarray`
            The array has the shape ``(n, 3)``. Each row holds the
            x, y and z coordinates of an atom.

        """

        return np.array(self._position_matrix.T)

    def with_position_matrix(self, position_matrix):
        """
        Return clone Molecule with new position matrix.

        Parameters
        ----------
        position_matrix : :class:`numpy.ndarray`
            A position matrix of the clone. The shape of the matrix
            is ``(n, 3)``.

        """

        clone = self.__class__.__new__(self.__class__)
        Molecule.__init__(
            self=clone,
            atoms=self._atoms,
            bonds=self._bonds,
            position_matrix=np.array(position_matrix),
        )
        return clone

    def _write_xyz_content(self):
        """
        Write basic `.xyz` file content of Molecule.

        """
        coords = self.get_position_matrix()
        content = [0]
        for i, atom in enumerate(self.get_atoms(), 1):
            x, y, z = (i for i in coords[atom.get_id()])
            content.append(
                f'{atom.get_element_string()} {x:f} {y:f} {z:f}\n'
            )
        # Set first line to the atom_count.
        content[0] = f'{i}\n\n'

        return content

    def write_xyz_file(self, path):
        """
        Write basic `.xyz` file of Molecule to `path`.

        Connectivity is not maintained in this file type!

        """

        content = self._write_xyz_content()

        with open(path, 'w') as f:
            f.write(''.join(content))

    def _write_pdb_content(self):
        """
        Write basic `.pdb` file content of Molecule.

        """

        content = []
        atom_counts = {}
        hetatm = 'HETATM'
        alt_loc = ''
        res_name = 'UNL'
        chain_id = ''
        res_seq = '1'
        i_code = ''
        occupancy = '1.00'
        temp_factor = '0.00'

        coords = self.get_position_matrix()
        # This set will be used by bonds.
        atoms = set()
        for i, atom in enumerate(self.get_atoms(), 1):
            x, y, z = (i for i in coords[atom.get_id()])
            atom_id = atom.get_id()
            atoms.add(atom_id)
            serial = atom_id+1
            element = atom.get_element_string()
            charge = 0
            atom_counts[element] = atom_counts.get(element, 0) + 1
            name = f'{element}{atom_counts[element]}'

            content.append(
                f'{hetatm:<6}{serial:>5} {name:<4}'
                f'{alt_loc:<1}{res_name:<3} {chain_id:<1}'
                f'{res_seq:>4}{i_code:<1}   '
                f' {x:>7.3f} {y:>7.3f} {z:>7.3f}'
                f'{occupancy:>6}{temp_factor:>6}          '
                f'{element:>2}{charge:>2}\n'
            )

        conect = 'CONECT'
        for bond in self.get_bonds():
            a1 = bond.get_atom1_id()
            a2 = bond.get_atom2_id()
            if a1 in atoms and a2 in atoms:
                content.append(
                    f'{conect:<6}{a1+1:>5}{a2+1:>5}               \n'
                )

        content.append('END\n')

        return content

    def write_pdb_file(self, path):
        """
        Write basic `.pdb` file of Molecule to `path`.

        """

        content = self._write_pdb_content()

        with open(path, 'w') as f:
            f.write(''.join(content))

    def get_atoms(self):
        """
        Yield the atoms in the molecule, ordered as input.

        Yields
        ------
        :class:`.Atom`
            An atom in the molecule.

        """

        for atom in self._atoms:
            yield atom

    def get_bonds(self):
        """
        Yield the bonds in the molecule, ordered as input.

        Yields
        ------
        :class:`.Bond`
            A bond in the molecule.

        """

        for bond in self._bonds:
            yield bond

    def get_centroid(self, atom_ids=None):
        """
        Return the centroid.

        Parameters
        ----------
        atom_ids : :class:`iterable` of :class:`int`, optional
            The ids of atoms which are used to calculate the
            centroid. Can be a single :class:`int`, if a single
            atom is to be used, or ``None`` if all atoms are to be
            used.

        Returns
        -------
        :class:`numpy.ndarray`
            The centroid of atoms specified by `atom_ids`.

        Raises
        ------
        :class:`ValueError`
            If `atom_ids` has a length of ``0``.

        """

        if atom_ids is None:
            atom_ids = range(len(self._atoms))
        elif isinstance(atom_ids, int):
            atom_ids = (atom_ids, )
        elif not isinstance(atom_ids, (list, tuple)):
            atom_ids = list(atom_ids)

        if len(atom_ids) == 0:
            raise ValueError('atom_ids was of length 0.')

        return np.divide(
            self._position_matrix[:, atom_ids].sum(axis=1),
            len(atom_ids)
        )

    def get_subunits(self, bond_pair_ids):
        """
        Get connected graphs based on Molecule separated by bonds.

        Parameters
        ----------
        bond_pair_ids :
            :class:`iterable` of :class:`tuple` of :class:`ints`
            Iterable of pairs of atom ids with bond between them to
            optimize.

        Returns
        -------
        subunits : :class:`.dict`
            The subunits of `mol` split by bonds defined by
            `bond_pair_ids`. Key is subunit identifier, Value is
            :class:`iterable` of atom ids in subunit.

        """

        # Produce a graph from the molecule that does not include edges
        # where the bonds to be optimized are.
        mol_graph = nx.Graph()
        for atom in self.get_atoms():
            mol_graph.add_node(atom.get_id())

        # Add edges.
        for bond in self.get_bonds():
            pair_ids = (bond.get_atom1_id(), bond.get_atom2_id())
            if pair_ids not in bond_pair_ids:
                mol_graph.add_edge(*pair_ids)

        # Get atom ids in disconnected subgraphs.
        subunits = {
            i: sg
            for i, sg in enumerate(nx.connected_components(mol_graph))
        }

        return subunits

    def __str__(self):
        return repr(self)

    def __repr__(self):
        return f'<{self.__class__.__name__} at {id(self)}>'


import sys
import argparse
import os
import numpy as np
import similarity_score
import TAD_calling
import TAD_split
import shutil

def printHelp():
    print('\nTADsplimer version 1.0.3')
    print('For help information for each function, try:\npython3 TADsplimer.py <function> -h')
    print('\nFunctions:')
    print('\tsplit_TADs:\n\t\tsplit TAD detection(input contact maps in two conditions)\n')
    print('\tTAD_calculator:\n\t\tidentify the topological domain\n')
    print('\tsplit_TADs_alternate:\n\t\tsplit TAD detection (input contact maps in two conditions with pre-called TADs)\n')
    print('\tTAD_similarity:\n\t\tfour algorithms of TAD similarity\n')
    print('')

def TAD_calculator(command='TAD_calculator'):
    '''
    Description:
        This function provide an entrance to identify TAD.
    parameters:
        none
    '''

    if (len(sys.argv) < 3) and ('-h' not in sys.argv) and ('--help' not in sys.argv):
        print("\nusage:\npython3 TADsplimer.py TAD_calculator <contact_map_file_paths> [optional arguments]\n\n"
              "for more help, please try: python3 TADsplimer.py TAD_calculator -h\n")
        return 0

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     usage="\n\npython3 TADsplimer.py TAD_calculator "
                                           "[optional arguments]\n\n", description='')
    parser.add_argument('command', default=None,
                        help="set as 'TAD_calculator' to identify TADs relative to Hi-C contact map")

    parser.add_argument('-c', '--contact_map', dest="contact_map", default=None,
                        help="path to Hi-C contact map")

    parser.add_argument('-u', '--up', dest="up", default=0, type=float,
                        help="up cutoff for Hi-C contact map")

    parser.add_argument('-d', '--down', dest="down", default=0, type=float,
                        help="down cutoff for Hi-C contact map")

    parser.add_argument('-o', '--TAD_output', dest="output", default=None,
                        help="path for the output file of TADs")

    parser.add_argument('-p', '--TAD_plot', dest="plot", default=1, type=int,
                        help="Set to 1 to plot the contact map and TAD, else set to 0 to cancel this analysis.")

    parser.add_argument('--sub_map', dest="sub_map", default=1, type=int,
                        help="Set to 1 to output sub contact maps and TADs, else set to 0 to cancel this analysis.")

    args = parser.parse_args()

    if not os.path.exists(args.output):
        print('path does not exist\n')
    else:
        up, down, _ = TAD_calling.file_split(args.contact_map, args.output)
        dirs = os.listdir(args.output)
        if args.up != 0 or args.down != 0:
            up = args.up
            down = args.down
        tad = np.empty(shape=[0, 3])
        for file in dirs:
            if file[-7:] == '.subchr':
                start = file.split('.')[-3]
                tad_temp = TAD_calling.TAD_calling(os.path.join(args.output, file), up, down)
                tad_location = tad_temp
                if args.plot == 1:
                    try:
                        TAD_calling.TAD_plot(os.path.join(args.output, file), os.path.join(args.output, file + '.tiff'),
                                             tad_temp, down, up)
                    except:
                        pass

                try:
                    tad_location[:, 1:3] = tad_location[:, 1:3] + int(start)
                    if not np.array_equal(tad_location, np.array([0, int(start), int(start)])):
                        tad = np.vstack((tad, tad_location))
                    if args.sub_map == 1:
                        np.savetxt(os.path.join(args.output, file + '.TAD.txt'), tad_location, delimiter="\t")
                except:
                    tad_location = np.array([[0, 0, 0]])
                    if args.sub_map == 1:
                        np.savetxt(os.path.join(args.output, file + '.TAD.txt'), tad_location, delimiter="\t")
        sorted_idx = np.lexsort(tad[:, 1:3].T)
        sorted_tad = tad[sorted_idx,:]
        row_mask = np.append([True], np.any(np.diff(sorted_tad[:, 1:3], axis=0), 1))
        unique_tad = sorted_tad[row_mask]
        np.savetxt(os.path.join(args.output, file + '.TAD.merge.txt'), unique_tad, delimiter="\t")

def split_TADs_alternate(command='split_TADs_alternate'):
    '''
    corner split algorithm for identifying split TAD
    '''

    if (len(sys.argv) < 3) and ('-h' not in sys.argv) and ('--help' not in sys.argv):
        # at least two parameter need to be specified, will print help message if no parameter is specified
        print("\nusage:\npython3 TADsplimer.py TAD_calculator <contact_map_file_paths> [optional arguments]\n\n"
              "for more help, please try: python3 TADsplimer.py TAD_calculator -h\n")
        return 0

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     usage="\n\npython3 TADsplimer.py corner_split <contact_map_file_paths> "
                                           "[optional arguments]\n\n", description='')
    parser.add_argument('command', default=None,
                        help="set as 'split_TADs_alternate' to identify split TAD")

    parser.add_argument('-c', '--contact_maps', dest="contact_map", default=None,
                        help="paths to Hi-C contact maps in two conditions. paths must be separated by the comma ','.")

    parser.add_argument('--contact_maps_aliases', dest="aliases", default=None,
                        help="A set of short aliases for the contact map. Paths must be separated by the comma ','.")

    parser.add_argument('-t', '--TAD', dest="TAD", default=None,
                        help="input files of TADs for two compared Hi-C contact maps. Paths must be separated by the"
                             " comma ','.")

    parser.add_argument('-u', '--up_cutoff', dest="up", default="0,0",
                        help="up cutoff for two compared Hi-C contact maps, paths must be separated by the comma ','.")

    parser.add_argument('-j', '--adjust_quality', dest="adjust_quality", default=0, type=int,
                        help="set as 1 to normalize sequence quality for two Hi-C contact maps, set as 0 not to "
                             "normalize sequence quality for two Hi-C contact maps")

    parser.add_argument('-o', '--output', dest="output", default=None,
                        help="path to output files")

    parser.add_argument('-d', '--split_direction', dest="direction", default=0, type=int,
                        help="set as 0: output TADs split in both two contact maps, set as 1: output TADs split in "
                             "contact map 1, set as 2: output TADs split in contact map 2")

    args = parser.parse_args()

    file = args.contact_map.split(',')
    TAD = args.TAD.split(',')
    up = args.up.split(',')
    aliases = args.aliases.split(',')

    TAD1 = np.loadtxt(TAD[0])
    TAD2 = np.loadtxt(TAD[1])

    TAD_s, TAD1_only, TAD2_only= TAD_split.TAD_matching(TAD2, TAD1)

    D1, D2, D3 = TAD_split.corner_split_score(TAD1, TAD2, TAD_s)

    map1 = np.loadtxt(file[0])
    map2 = np.loadtxt(file[1])

    up1, _ , _ = TAD_calling.file_split(TAD[0], outfile="", print_subcontact=1)
    up2, _ , _ = TAD_calling.file_split(TAD[1], outfile="", print_subcontact=1)

    if up[0] != 0 or up[1] != 0:
        up1 = float(up[0])
        up2 = float(up[1])

    if args.adjust_quality == 0:
        split1_1, split2_1, loc_d, loc_u = TAD_split.split_region(file[1], file[0], D3, D1, up2, up1, 1)
    elif args.adjust_quality == 1:
        map1 = np.loadtxt(file[0])
        map2 = np.loadtxt(file[1])
        ratio1 = TAD_split.get_ratio(map1, TAD1, float(up[0]))
        ratio2 = TAD_split.get_ratio(map2, TAD2, float(up[1]))
        fold = np.mean(ratio1)/np.mean(ratio2)
        split1_1, split2_1, loc_d, loc_u = TAD_split.split_region(file[1], file[0], D3, D1, up2, up1, fold)
    else:
        print("\nusage:\npython3 TADsplimer.py TAD_calculator <contact_map_file_paths> [optional arguments]\n\n"
              "for more help, please try: python3 TADsplimer.py TAD_calculator -h\n")
    try:
        if split1_1 == 0:
            print('No split TAD')
    except:
        if (not os.path.exists(args.output)):
            os.makedirs(args.output)
        else:
            scc = similarity_score.similarity_scc(map1, map2, loc_u[:, 0:3])
            Laplacian = similarity_score.similarity_Laplacian(map1, map2, loc_u[:, 0:3])
            hash = similarity_score.hash_similarity(map1, map2, loc_u[:, 0:3])
            loc_u2 = np.c_[loc_u, scc[:, 3], Laplacian[:, 3], hash[:, 3]]
            np.savetxt(os.path.join(args.output, aliases[0]+'->'+aliases[1]+'.split.txt'), loc_d)
            np.savetxt(os.path.join(args.output, aliases[0]+'->'+aliases[1]+'.merge.txt'), loc_u2)

def split_TADs(command='split_TADs'):
    '''
    corner split algorithm for identifying split TAD
    '''

    if (len(sys.argv) < 3) and ('-h' not in sys.argv) and ('--help' not in sys.argv):
        # at least two parameter need to be specified, will print help message if no parameter is specified
        print("\nusage:\npython3 TADsplimer.py TAD_calculator <contact_map_file_paths> [optional arguments]\n\n"
              "for more help, please try: python3 TADsplimer.py TAD_calculator -h\n")
        return 0

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     usage="\n\npython3 TADsplimer.py corner_split <contact_map_file_paths> "
                                           "[optional arguments]\n\n", description='')
    parser.add_argument('command', default=None,
                        help="set as 'split_TADs' to identify split TADs")

    parser.add_argument('-c', '--contact_maps', dest="contact_map", default=None,
                        help="paths to two contact maps. paths must be separated by the comma ','.")

    parser.add_argument('--contact_maps_aliases', dest="aliases", default=None,
                        help="A set of short aliases for two contact maps. Paths must be separated by the comma ','.")

    parser.add_argument('--up_cutoff', dest="up_cutoff", default="None",
                        help="paths for up cutoff of two compared Hi-C contact maps, paths must be separated by the comma ','.")

    parser.add_argument('--down_cutoff', dest="down_cutoff", default="None",
                        help="paths for down cutoff of two compared Hi-C contact maps, paths must be separated by the comma ','.")

    parser.add_argument('-j', '--adjust_quality', dest="adjust_quality", default=0, type=int,
                        help="set as 1 to normalize sequence quality for two Hi-C contact maps, set as 0 not to "
                             "normalize sequence quality for two Hi-C contact maps")

    parser.add_argument('-o', '--output', dest="output", default=None,
                        help="path to output files")

    parser.add_argument('-d', '--split_direction', dest="direction", default=0, type=int,
                        help="set as 0: output TADs split in both two contact maps, set as 1: output TADs split in "
                             "contact map1, set as 2: output TADs split in contact map2")

    ### add
    parser.add_argument('-t', '--tmp', dest="tmp", default=0, type=int,
                        help="set as 0: save temporary file in the output folder, set as 1: not save temporary file in"
                             "the output folder")

    args = parser.parse_args()

    file = args.contact_map.split(',')
    aliases = args.aliases.split(',')

    tmp_dir1 = os.path.join(args.output, aliases[0])
    if not os.path.isdir(tmp_dir1):
        os.makedirs(tmp_dir1)

    up1, down1, files1 = TAD_calling.file_split(file[0], outfile=tmp_dir1, print_subcontact=0, aliases=aliases[0])

    print('auto:', up1, down1)

    if args.up_cutoff != 'None' and args.down_cutoff != 'None':
        up_cutoff = args.up_cutoff.split(',')
        up1 = float(up_cutoff[0])
        down_cutoff = args.down_cutoff.split(',')
        down1 = float(down_cutoff[0])
        print('manual:', up1, down1)

    TAD_calling.subTAD_calling(tmp_dir1, up1, down1, aliase=aliases[0])

    tmp_dir2 = os.path.join(args.output, aliases[1])
    if not os.path.isdir(tmp_dir2):
        os.makedirs(tmp_dir2)

    up2, down2, files2 = TAD_calling.file_split(file[1], outfile=tmp_dir2, print_subcontact=0, aliases=aliases[1])

    print('auto:', up2, down2)

    if args.up_cutoff != 'None' and args.down_cutoff != 'None':
        up_cutoff = args.up_cutoff.split(',')
        up2 = float(up_cutoff[1])
        down_cutoff = args.down_cutoff.split(',')
        down2 = float(down_cutoff[1])
        print('manual:', up2, down2)

    TAD_calling.subTAD_calling(tmp_dir2, up2, down2, aliase=aliases[1])

    for i in range(len(files1)):
        contact_map_file1 = files1[i]
        contact_map_file2 = files2[i]
        TAD_file1 = contact_map_file1+'.TAD.txt'
        TAD_file2 = contact_map_file2 + '.TAD.txt'
        base1 = os.path.basename(contact_map_file1)
        base2 = os.path.basename(contact_map_file2)
        start1 = int(base1.split('.')[1])
        start2 = int(base2.split('.')[1])
        end = base2.split('.')[2]
        if start1 != start2:
            print('input file wrong')
            break

        TAD1 = np.loadtxt(TAD_file1)
        TAD2 = np.loadtxt(TAD_file2)

        try:
            TAD1[:, 1:3] = TAD1[:, 1:3]-start1
            TAD2[:, 1:3] = TAD2[:, 1:3]-start2
        except:
            pass

        map1 = np.loadtxt(contact_map_file1)
        map2 = np.loadtxt(contact_map_file2)

        try:
            TAD_split.write_split_tad(TAD1, TAD2, contact_map_file1, contact_map_file2, map1, map2, up1, up2,
                                      args.adjust_quality, args.output, start1, end, aliases[0], aliases[1])

            TAD_split.write_split_tad(TAD2, TAD1, contact_map_file2, contact_map_file1, map2, map1, up2, up1,
                                      args.adjust_quality, args.output, start1, end, aliases[1], aliases[0])

        except:
            pass

    a = TAD_split.merge_outputfile(aliases[0], aliases[1], args.output)
    if a[0].shape != 0:
        np.savetxt(os.path.join(args.output, aliases[0] + '->' + aliases[1] + '.all.merge.txt'), a[0])
    if a[1].shape != 0:
        np.savetxt(os.path.join(args.output, aliases[0] + '->' + aliases[1] + '.all.split.txt'), a[1])

    b = TAD_split.merge_outputfile(aliases[1], aliases[0], args.output)
    if b[0].shape != 0:
        np.savetxt(os.path.join(args.output, aliases[1] + '->' + aliases[0] + '.all.merge.txt'), b[0])
    if b[1].shape != 0:
        np.savetxt(os.path.join(args.output, aliases[1] + '->' + aliases[0] + '.all.split.txt'), b[1])

    shutil.move(os.path.join(args.output, aliases[0], aliases[0]+'.TAD.merge.txt'),
                os.path.join(args.output, aliases[0]+'.TAD.txt'))
    shutil.move(os.path.join(args.output, aliases[1], aliases[1] + '.TAD.merge.txt'),
                os.path.join(args.output, aliases[1] + '.TAD.txt'))

    if args.tmp == 1:
        shutil.rmtree(os.path.join(args.output, aliases[0]))
        shutil.rmtree(os.path.join(args.output, aliases[1]))

def TAD_similarity(command='TAD_similarity'):
    if (len(sys.argv) < 3) and ('-h' not in sys.argv) and ('--help' not in sys.argv):
        # at least two parameter need to be specified, will print help message if no parameter is specified
        print("\nusage:\npython3 TADsplimer.py TAD_similarity <contact_map_file_paths> [optional arguments]\n\n"
              "for more help, please try: python3 TADsplimer.py TAD_similarity -h\n")
        return 0

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     usage="\n\npython3 TADsplimer.py TAD_similarity <contact_map_file_paths> "
                                           "[optional arguments]\n\n", description='')
    parser.add_argument('command', default=None,
                        help="set as 'TAD_similarity' to calculate similarity of two TADs")

    parser.add_argument('-c', '--contact_maps', dest="contact_map", default=None,
                        help="paths to Hi-C contact maps in two conditions. paths must be separated by the comma ','.")

    parser.add_argument('-t', '--TAD', dest="TAD", default=None,
                        help="input files of TADs for two compared Hi-C contact maps. Paths must be separated by the"
                             " comma ','.")

    parser.add_argument('-o', '--output', dest="output", default=None,
                        help="path to output files")

    args = parser.parse_args()

    file1 = args.contact_map.split(',')
    file2 = args.TAD.split(',')
    map1 = np.loadtxt(file1[0])
    map2 = np.loadtxt(file1[1])
    TAD1 = np.loadtxt(file2[0])
    TAD2 = np.loadtxt(file2[1])
    head1, tail1 = os.path.split(file1[0])
    base1 = os.path.splitext(tail1)
    head2, tail2 = os.path.split(file1[1])
    base2 = os.path.splitext(tail2)

    TAD = np.concatenate((TAD1, TAD2), axis=0)
    scc = similarity_score.similarity_scc(map1, map2, TAD)
    Laplacian = similarity_score.similarity_Laplacian(map1, map2, TAD)
    hash = similarity_score.hash_similarity(map1, map2, TAD)
    np.savetxt(os.path.join(args.output, base1[0]+'_'+base2[0]+'_scc.txt'), scc, delimiter='\t')
    np.savetxt(os.path.join(args.output, base1[0]+'_'+base2[0]+'_laplacian.txt'), Laplacian, delimiter='\t')
    np.savetxt(os.path.join(args.output, base1[0]+'_'+base2[0]+'_hash.txt'), hash, delimiter='\t')

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == 'TAD_calculator':
            TAD_calculator(command='TAD_calculator')
        elif sys.argv[1] == 'split_TADs':
            split_TADs(command='split_TADs')
        elif sys.argv[1] == 'split_TADs_alternate':
            split_TADs_alternate(command='split_TADs_alternate')
        elif sys.argv[1] == 'TAD_similarity':
            TAD_similarity(command='TAD_similarity')
        else:
            printHelp()
    else:
        print('\nTADsplimer version 1.0.3')
        print('For a list of functions in TADsplimer, please try:\npython3 TADsplimer.py -h')
        print('')


#!/usr/bin/env python
from ..utils.protocol import frame_path_at, track_box_at_frame, bbox_hash, \
    tubelets_overlap, tubelets_proto_from_tracks_proto, tubelet_box_at_frame, det_score, tubelet_box_proto_at_frame
from ..utils.common import imread, svm_from_rcnn_model, iou
from ..utils.log import logging
from ..vdet.image_det import googlenet_det, googlenet_features, svm_scores
from ..vdet.dataset import imagenet_vdet_classes, index_vdet_to_det
import numpy as np
import sys
import os
import copy
from scipy.interpolate import interp1d
from collections import defaultdict

def score_conv_cls(score_proto, net):
    new_score_proto = copy.copy(score_proto)
    print "{}: {} tubelet(s).".format(score_proto['video'], len(new_score_proto['tubelets']))
    for tubelet in new_score_proto['tubelets']:
        track = {}
        track['length'] = len(tubelet['boxes'])
        track['gt'] = tubelet['gt']
        track['mean_iou'] = np.mean([map(lambda x:x['gt_overlap'],
                                  tubelet['boxes'])])
        track['det_scores'] = map(lambda x:x['det_score'],
                                  tubelet['boxes'])
        track['track_scores'] = map(lambda x:x['track_score'],
                                  tubelet['boxes'])
        track['anchors'] = map(lambda x:x['anchor'] * 1. / track['length'],
                                  tubelet['boxes'])
        track['abs_anchors'] = map(abs, track['anchors'])
        track['gt_overlaps'] = map(lambda x:x['gt_overlap'],
                                  tubelet['boxes'])
        track['labels'] = [1 if iou >= 0.5 else 0 for iou in track['gt_overlaps']]

        # skip memory heavy features if possible
        if 'all_scores' in net.blobs.keys():
            track['all_scores'] = map(lambda x:x['all_score'],
                                  tubelet['boxes'])
        if 'feats' in net.blobs.keys():
            track['feats'] = map(lambda x:x['feat'],
                                  tubelet['boxes'])

        for blob_name in set(net.blobs.keys()).intersection(set(track.keys())):
            num_channels = net.blobs[blob_name].shape[1]
            net.blobs[blob_name].reshape(1, num_channels, 1, track['length'])
            net.blobs[blob_name].data[...] = np.asarray(track[blob_name], dtype='float32')
        blobs_out = net.forward()
        probs = blobs_out['probs'][:, 1,:]
        for box, prob in zip(tubelet['boxes'], probs.ravel()):
            box['conv_score'] = float(prob)
    return new_score_proto

def fast_rcnn_cls(video_proto, track_proto, net, class_idx):
    sys.path.insert(1, os.path.join(os.path.dirname(__file__),
        '../../External/fast-rcnn/lib/'))
    sys.path.insert(1, os.path.join(os.path.dirname(__file__),
        '../../External/fast-rcnn/caffe-fast-rcnn/python'))
    from fast_rcnn.test import im_detect
    new_tracks = [[] for _ in track_proto['tracks']]
    for frame in video_proto['frames']:
        frame_id = frame['frame']
        img = imread(frame_path_at(video_proto, frame_id))
        boxes = [track_box_at_frame(tracklet, frame_id) \
                 for tracklet in track_proto['tracks']]
        valid_boxes = np.asarray([box for box in boxes if box is not None])
        valid_index = [i for i in xrange(len(boxes)) if boxes[i] is not None]
        scores, pred_boxes = im_detect(net, img, valid_boxes)
        for score, box, track_id in zip(scores, pred_boxes, valid_index):
            new_tracks[track_id].append(
                {
                    "frame": frame_id,
                    "bbox": list(box),
                    "score": score[class_idx],
                    "hash": bbox_hash(video_proto['video'], frame_id, box)
                })
    return new_tracks

def googlenet_cls(video_proto, track_proto, net, class_idx):
    new_tracks = [[] for _ in track_proto['tracks']]
    logging.info("Classifying {}...".format(video_proto['video']))
    for frame in video_proto['frames']:
        frame_id = frame['frame']
        img = imread(frame_path_at(video_proto, frame_id))
        boxes = [track_box_at_frame(tracklet, frame_id) \
                 for tracklet in track_proto['tracks']]
        valid_boxes = np.asarray([box for box in boxes if box is not None])
        valid_index = [i for i, box in enumerate(boxes) if box is not None]
        logging.info("frame {}: {} boxes".format(frame_id, len(valid_index)))
        for box, track_id in zip(valid_boxes, valid_index):
            scores = googlenet_det(img, box, net)
            new_tracks[track_id].append(
                {
                    "frame": frame_id,
                    "bbox": list(box),
                    "score": scores[class_idx],
                    "hash": bbox_hash(video_proto['video'], frame_id, box)
                })
    return new_tracks

def rcnn_scoring(vid_proto, track_proto, net, class_idx, rcnn_model,
        save_feat=False, save_all_sc=False):
    svm_model = svm_from_rcnn_model(rcnn_model)
    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Scoring {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))
    for frame in vid_proto['frames']:
        frame_id = frame['frame']
        img = imread(frame_path_at(vid_proto, frame_id))
        boxes = [tubelet_box_at_frame(tubelet, frame_id) \
                 for tubelet in tubelets_proto]
        valid_boxes = np.asarray([box for box in boxes if box is not None])
        valid_index = [i for i, box in enumerate(boxes) if box is not None]
        logging.info("frame {}: {} boxes".format(frame_id, len(valid_index)))
        if len(valid_index) == 0:
            continue
        features = googlenet_features(img, valid_boxes, net, 'pool5')
        scores = svm_scores(features, svm_model)
        if scores.shape[1] == 200:
            cls_scores = scores[:, index_vdet_to_det[class_idx] - 1]
        else:
            raise
        for score, tubelet_id, feat, all_score in \
                zip(cls_scores, valid_index, features, scores):
            cur_box = [box for box in tubelets_proto[tubelet_id]['boxes'] \
                if box['frame'] == frame_id]
            assert len(cur_box) == 1
            cur_box[0]['det_score'] = score
            if save_feat:
                cur_box[0]['feat'] = feat.ravel().tolist()
            if save_all_sc:
                cur_box[0]['all_score'] = all_score.ravel().tolist()
    return tubelets_proto

def sampling_boxes(orig_box, num, ratio = 0.05, return_orig=True):
    h, w = orig_box[3] - orig_box[1], orig_box[2] - orig_box[0]
    offsets = np.random.uniform(-ratio, ratio, [num, 4]) * [w, h, w, h]
    if not return_orig:
        return orig_box + offsets
    else:
        return np.vstack((orig_box, orig_box+offsets))

def rcnn_sampling_scoring(vid_proto, track_proto, net, class_idx, rcnn_model,
        samples_per_box = 32, ratio = 0.05,
        save_feat=False, save_all_sc=False):
    svm_model = svm_from_rcnn_model(rcnn_model)
    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Scoring {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))
    for frame in vid_proto['frames']:
        frame_id = frame['frame']
        img = imread(frame_path_at(vid_proto, frame_id))
        boxes = [tubelet_box_at_frame(tubelet, frame_id) \
                 for tubelet in tubelets_proto]
        valid_boxes = np.asarray([box for box in boxes if box is not None])
        valid_index = [i for i, box in enumerate(boxes) if box is not None]
        logging.info("frame {}: {} boxes".format(frame_id, len(valid_index)))
        if len(valid_index) == 0:
            continue

        # sample nearby boxes to increase spatial robustness
        sampled_boxes = np.vstack([sampling_boxes(box, samples_per_box, ratio) \
                for box in valid_boxes])
        features = googlenet_features(img, sampled_boxes, net, 'pool5')
        scores = svm_scores(features, svm_model)
        if scores.shape[1] == 200:
            cls_scores = scores[:, index_vdet_to_det[class_idx] - 1]
            cls_scores = cls_scores.reshape((len(valid_index), -1))
            max_scores = cls_scores.max(axis=1)
            # extract cooresponding features and all class scores
            max_idx = np.argmax(cls_scores, axis=1)
        else:
            raise

        # extract features and scores of box with maximum score
        features = features.reshape((len(valid_index), samples_per_box+1, -1))
        features = features[xrange(len(valid_index)), max_idx,:]
        scores = scores.reshape((len(valid_index), samples_per_box+1, -1))
        scores = scores[xrange(len(valid_index)), max_idx,:]
        sampled_boxes = sampled_boxes.reshape((len(valid_index), samples_per_box+1, -1))
        boxes = sampled_boxes[xrange(len(valid_index)), max_idx,:]
        for score, tubelet_id, feat, all_score, max_box in \
                zip(max_scores, valid_index, features, scores, boxes):
            cur_box = [box for box in tubelets_proto[tubelet_id]['boxes'] \
                if box['frame'] == frame_id]
            assert len(cur_box) == 1
            cur_box[0]['det_score'] = score
            cur_box[0]['bbox'] = max_box.tolist()
            if save_feat:
                cur_box[0]['feat'] = feat.ravel().tolist()
            if save_all_sc:
                cur_box[0]['all_score'] = all_score.ravel().tolist()
    return tubelets_proto

def rcnn_sampling_dets_scoring(vid_proto, track_proto, det_proto,
        net, class_idx, rcnn_model, overlap_thres=0.7,
        save_feat=False, save_all_sc=False):
    svm_model = svm_from_rcnn_model(rcnn_model)
    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Scoring {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))
    for frame in vid_proto['frames']:
        frame_id = frame['frame']
        img = imread(frame_path_at(vid_proto, frame_id))
        boxes = [tubelet_box_at_frame(tubelet, frame_id) \
                 for tubelet in tubelets_proto]
        valid_boxes = np.asarray([box for box in boxes if box is not None])
        valid_index = [i for i, box in enumerate(boxes) if box is not None]
        logging.info("frame {}: {} boxes".format(frame_id, len(valid_index)))
        if len(valid_index) == 0:
            continue
        # compute rcnn scores
        features = googlenet_features(img, valid_boxes, net, 'pool5')
        scores = svm_scores(features, svm_model)
        if scores.shape[1] == 200:
            cls_scores = scores[:, index_vdet_to_det[class_idx] - 1]
        else:
            raise

        # find all detection proposal boxes in current frame
        dets = [det for det in det_proto['detections'] if det['frame']==frame_id]
        det_boxes = np.asarray(map(lambda x:x['bbox'], dets))
        det_scores = np.asarray(map(lambda x:det_score(x, class_idx), dets))

        for score, tubelet_id, feat, all_score in \
                zip(cls_scores, valid_index, features, scores):
            cur_box = [box for box in tubelets_proto[tubelet_id]['boxes'] \
                if box['frame'] == frame_id]
            assert len(cur_box) == 1
            # calculate overlaps with all det_boxes in current frame
            if len(det_boxes) > 0:
                overlaps = iou([cur_box[0]['bbox']], det_boxes)
                conf_idx = (overlaps > overlap_thres).ravel()
            else:
                conf_idx = [False]
            if np.any(conf_idx):
                conf_boxes = det_boxes[conf_idx]
                conf_scores = det_scores[conf_idx]
                max_idx = np.argmax(conf_scores)
                max_score = conf_scores[max_idx]
                max_box = conf_boxes[max_idx].tolist()
            else:
                max_score = -np.inf
            if max_score > score:
                cur_box[0]['det_score'] = max_score
                cur_box[0]['bbox'] = max_box
                max_feat = googlenet_features(img, [max_box], net, 'pool5')
                max_all_scores = svm_scores(max_feat, svm_model)
                if save_feat:
                    cur_box[0]['feat'] = max_feat.ravel().tolist()
                if save_all_sc:
                    cur_box[0]['all_score'] = max_all_scores.ravel().tolist()
            else:
                cur_box[0]['det_score'] = score
                if save_feat:
                    cur_box[0]['feat'] = feat.ravel().tolist()
                if save_all_sc:
                    cur_box[0]['all_score'] = all_score.ravel().tolist()
    return tubelets_proto

def scoring_tracks(vid_proto, track_proto, annot_proto,
        sc_method, net, class_idx):
    assert vid_proto['video'] == track_proto['video']
    score_proto = {}
    score_proto['video'] = vid_proto['video']
    score_proto['method'] = sc_method.__name__
    tubelets_proto = sc_method(vid_proto, track_proto, net, class_idx)
    if annot_proto is not None:
        tubelets_proto = tubelets_overlap(tubelets_proto, annot_proto, class_idx)
    score_proto['tubelets'] = tubelets_proto
    return score_proto

def classify_tracks(video_proto, track_proto, cls_method, net, class_idx):
    assert video_proto['video'] == track_proto['video']
    cls_track = {}
    cls_track['video'] = video_proto['video']
    cls_track['method'] = cls_method.__name__
    cls_track['tracks'] = cls_method(video_proto, track_proto, net, class_idx)
    return cls_track

def do_score_completion(score_proto):
    for tubelet in score_proto['tubelets']:
        # deal with -100000
        boxes = tubelet['boxes']
        for i, box in enumerate(boxes):
            if box['det_score'] > -10: continue
            j = i
            while j < len(boxes) and boxes[j]['det_score'] <= -10:
                j += 1
            if i == 0:
                for k in xrange(i, j):
                    boxes[k]['det_score'] = boxes[j]['det_score']
            elif j == len(boxes):
                for k in xrange(i, j):
                    boxes[k]['det_score'] = boxes[i-1]['det_score']
            else:
                l = boxes[i-1]['det_score']
                r = boxes[j]['det_score']
                for k in xrange(i, j):
                    boxes[k]['det_score'] = l + (r - l) * (k - i + 1) / (j - i + 1)

def dets_spatial_max_pooling(vid_proto, track_proto, det_proto, class_idx, overlap_thres=0.7):
    assert vid_proto['video'] == track_proto['video']
    score_proto = {}
    score_proto['video'] = vid_proto['video']
    score_proto['method'] = "spatial_max_pooling_IOU_{}".format(overlap_thres)

    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Sampling dets in {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))

    # build dict to fast indexing
    frame_to_tubelets_idx = defaultdict(list)
    for i, tubelet in enumerate(tubelets_proto):
        for j, box in enumerate(tubelet['boxes']):
            frame_to_tubelets_idx[box['frame']].append((i, j))
    frame_to_det_idx = defaultdict(list)
    for i, det in enumerate(det_proto['detections']):
        frame_to_det_idx[det['frame']].append(i)

    for frame in vid_proto['frames']:
        frame_id = frame['frame']
        det_idx = frame_to_det_idx[frame_id]
        if len(det_idx) == 0: continue
        det_boxes = np.asarray([det_proto['detections'][i]['bbox'] for i in det_idx])
        det_scores = np.asarray([det_proto['detections'][i]['scores'][class_idx - 1]['score'] for i in det_idx])
        for i, j in frame_to_tubelets_idx[frame_id]:
            cur_box = tubelets_proto[i]['boxes'][j]
            if cur_box is None: continue
            overlaps = iou([cur_box['bbox']], det_boxes)
            overlap_idx = (overlaps > overlap_thres).ravel()
            if np.any(overlap_idx):
                conf_boxes = det_boxes[overlap_idx]
                conf_scores = det_scores[overlap_idx]
                max_idx = np.argmax(conf_scores)
                max_score = conf_scores[max_idx]
                max_box = conf_boxes[max_idx].tolist()
            else:
                print "Warning: Tubelet {} has no overlapping dets (IOU > {}).".format(
                    i, overlap_thres)
                max_score = -1e5
                max_box = cur_box['bbox']
            cur_box['det_score'] = float(max_score)
            cur_box['bbox'] = max_box
    score_proto['tubelets'] = tubelets_proto
    do_score_completion(score_proto)
    return score_proto

def anchor_propagate(vid_proto, track_proto, det_proto, class_idx):
    assert vid_proto['video'] == track_proto['video']
    score_proto = {}
    score_proto['video'] = vid_proto['video']
    score_proto['method'] = "anchor_propagate"
    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Propagating anchor scores in {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))
    frame_to_det_idx = defaultdict(list)
    for i, det in enumerate(det_proto['detections']):
        frame_to_det_idx[det['frame']].append(i)
    for tubelet in tubelets_proto:
        # find anchor box
        anchor_box = [box for box in tubelet['boxes'] if box['anchor'] == 0]
        assert len(anchor_box) == 1
        anchor_box = anchor_box[0]

        # get anchor detection score
        anchor_frame = anchor_box['frame']
        det_idx = frame_to_det_idx[anchor_frame]
        det_boxes = np.asarray([det_proto['detections'][i]['bbox'] for i in det_idx])
        det_scores = np.asarray([det_proto['detections'][i]['scores'][class_idx - 1]['score'] for i in det_idx])
        overlaps = iou([anchor_box['bbox']], det_boxes)[0]
        anchor_index = np.argmax(overlaps)
        anchor_score = det_scores[anchor_index]

        # propagate anchor scores
        for box in tubelet['boxes']:
            box['det_score'] = anchor_score
    score_proto['tubelets'] = tubelets_proto
    return score_proto

def score_proto_temporal_maxpool(score_proto, window_size):
    if window_size == 1:
        return score_proto
    if window_size % 2 != 1:
        raise ValueError('Window size must be odd!')
    half_window_size = window_size / 2

    new_score_proto = copy.copy(score_proto)
    new_score_proto['method'] += '_temporal_maxpool_{}'.format(window_size)
    for tubelet in new_score_proto['tubelets']:
        if tubelet['gt'] == 1:
            raise ValueError('Dangerous: Score file contains gt tracks!')

        boxes = tubelet['boxes']
        scores = [boxes[i]['det_score'] for i in range(len(boxes))]
        scores = np.array(scores)
        scores = np.pad(scores, (0, window_size-1), 'constant', constant_values=(-1e+5,-1e+5))
        scores = np.tile(scores.T, (window_size, 1))

        for roll in range(window_size):
            scores[roll] = np.roll(scores[roll], roll)

        max_score = scores.max(0)
        max_score = max_score[half_window_size:-half_window_size]

        for i in range(len(boxes)):
            tubelet['boxes'][i]['det_score'] = float(max_score[i])

    return new_score_proto

def extrap1d(interpolator):
    xs = interpolator.x
    ys = interpolator.y

    def pointwise(x):
        if x < xs[0]:
            return ys[0]+(x-xs[0])*(ys[1]-ys[0])/(xs[1]-xs[0])
        elif x > xs[-1]:
            return ys[-1]+(x-xs[-1])*(ys[-1]-ys[-2])/(xs[-1]-xs[-2])
        else:
            return interpolator(x)

    return pointwise

def score_proto_interpolation(score_proto, vid_proto):
    '''Perform interpolation on score protocols is only part of
       the tracks are available'''
    new_score_proto = {}
    new_score_proto['video'] = score_proto['video']
    new_score_proto['method'] = score_proto['method'] + '_interpolation'
    tubelets_proto = []
    idx_fun = lambda x:x['frame']
    funcs = {
        'x1': lambda x:x['bbox'][0],
        'y1': lambda x:x['bbox'][1],
        'x2': lambda x:x['bbox'][2],
        'y2': lambda x:x['bbox'][3],
        'det_score': lambda x:x['det_score'],
        'anchor': lambda x:x['anchor']
    }

    max_frames = len(vid_proto['frames'])

    for tubelet in score_proto['tubelets']:
        if tubelet['gt'] == 1:
            raise ValueError('Dangerous: Score file contains gt tracks!')
        if len(tubelet['boxes']) < 2:
            tubelets_proto.append(copy.copy(tubelet))
            continue

        # Generate interpolation function for each field
        truth_idx = map(idx_fun, tubelet['boxes'])
        interp_funcs = {}
        for field in funcs.keys():
            field_fun = funcs[field]
            truth_values = map(field_fun, tubelet['boxes'])
            interpolator = interp1d(truth_idx, truth_values)
            extrapolator = extrap1d(interpolator)
            interp_funcs[field] = extrapolator
        new_tubelet = {}
        for key in ['gt', 'class', 'class_index']:
            new_tubelet[key] = tubelet[key]
        new_tubelet['boxes'] = []
        min_idx = min(truth_idx)
        max_idx = max(truth_idx)
        # extrapolate first and last frames
        if min_idx == 2:
            min_idx = 1
        if max_idx == max_frames - 1:
            max_idx = max_frames
        for dense_idx in xrange(min_idx, max_idx+1):
            box = {}
            box['frame'] = dense_idx
            box['det_score'] = float(interp_funcs['det_score'](dense_idx))
            box['anchor'] = float(interp_funcs['anchor'](dense_idx))
            x1 = interp_funcs['x1'](dense_idx)
            y1 = interp_funcs['y1'](dense_idx)
            x2 = interp_funcs['x2'](dense_idx)
            y2 = interp_funcs['y2'](dense_idx)
            box['bbox'] = map(float, [x1, y1, x2, y2])
            new_tubelet['boxes'].append(box)
        tubelets_proto.append(new_tubelet)

    new_score_proto['tubelets'] = tubelets_proto
    return new_score_proto

def raw_dets_spatial_max_pooling(vid_proto, track_proto, frame_to_det, class_idx, overlap_thres=0.7):
    assert vid_proto['video'] == track_proto['video']
    score_proto = {}
    score_proto['video'] = vid_proto['video']
    score_proto['method'] = "spatial_max_pooling_IOU_{}".format(overlap_thres)

    tubelets_proto = tubelets_proto_from_tracks_proto(track_proto['tracks'], class_idx)
    logging.info("Sampling dets in {} for {}...".format(vid_proto['video'],
                 imagenet_vdet_classes[class_idx]))

    # build dict to fast indexing
    frame_to_tubelets_idx = defaultdict(list)
    for i, tubelet in enumerate(tubelets_proto):
        for j, box in enumerate(tubelet['boxes']):
            frame_to_tubelets_idx[box['frame']].append((i, j))

    for frame in vid_proto['frames']:
        frame_id = frame['frame']
        if frame_id not in frame_to_det: continue
        det_boxes, det_scores = frame_to_det[frame_id]
        if det_boxes.size == 0: continue
        det_scores = det_scores[:, class_idx - 1].ravel()
        for i, j in frame_to_tubelets_idx[frame_id]:
            cur_box = tubelets_proto[i]['boxes'][j]
            if cur_box is None: continue
            overlaps = iou([cur_box['bbox']], det_boxes)
            overlap_idx = (overlaps > overlap_thres).ravel()
            if np.any(overlap_idx):
                conf_boxes = det_boxes[overlap_idx]
                conf_scores = det_scores[overlap_idx]
                max_idx = np.argmax(conf_scores)
                max_score = conf_scores[max_idx]
                max_box = conf_boxes[max_idx].tolist()
            else:
                print "Warning: Tubelet {} has no overlapping dets (IOU > {}).".format(
                    i, overlap_thres)
                max_score = -1e5
                max_box = cur_box['bbox']
            cur_box['det_score'] = float(max_score)
            cur_box['bbox'] = max_box
    score_proto['tubelets'] = tubelets_proto
    do_score_completion(score_proto)
    return score_proto


def db_seresnext54_LSEP_VD():

    import torch
    from torch import nn

    model_name = 'db_seresnext54_LSEP_VD'

    device = 'cuda:0'

    torch.backends.cudnn.benchmark = True

    import numpy as np
    import pandas as pd
    from sklearn.model_selection import KFold, StratifiedKFold
    from sklearn.utils import shuffle

    from skimage.transform import resize
    from skimage.color import rgb2gray, gray2rgb

    from sklearn.metrics import f1_score, precision_score, recall_score, roc_auc_score

    from tqdm import tqdm_notebook

    import gc
    import math
    import sys

    from fastai import *
    from fastai.vision import *

    np.random.seed(42)

    data_dir = '../../../input/'
    submit_l1_dir = "../../../submits/"
    weights_dir = "../../weights/"
    results_dir = '../../../results/'

    name_label_dict = {
    0:  'Nucleoplasm',
    1:  'Nuclear membrane',
    2:  'Nucleoli',
    3:  'Nucleoli fibrillar center',
    4:  'Nuclear speckles',
    5:  'Nuclear bodies',
    6:  'Endoplasmic reticulum',
    7:  'Golgi apparatus',
    8:  'Peroxisomes',
    9:  'Endosomes',
    10:  'Lysosomes',
    11:  'Intermediate filaments',
    12:  'Actin filaments',
    13:  'Focal adhesion sites',
    14:  'Microtubules',
    15:  'Microtubule ends',
    16:  'Cytokinetic bridge',
    17:  'Mitotic spindle',
    18:  'Microtubule organizing center',
    19:  'Centrosome',
    20:  'Lipid droplets',
    21:  'Plasma membrane',
    22:  'Cell junctions',
    23:  'Mitochondria',
    24:  'Aggresome',
    25:  'Cytosol',
    26:  'Cytoplasmic bodies',
    27:  'Rods & rings' }

    def kfold_threshold(y_true, y_pred):
        n_classes = len(name_label_dict)
        classes_thresholds = []
        classes_scores = []
        kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=239)
        for i in range(n_classes):
            kf_class_thresholds = []
            if (sum(y_true[:,i]) > 20):
                for _, tst_inx in kf.split(y_true,y_true[:,i]):
                    t_min = np.min(y_pred[tst_inx,i])
                    t_max = np.max(y_pred[tst_inx,i])
                    thresholds = np.linspace(t_min, t_max, 50)
                    scores = np.array([
                        f1_score(y_true[tst_inx,i], np.int32(y_pred[tst_inx,i] >= threshold)) for threshold in thresholds
                    ])
                    threshold_best_index = np.argmax(scores)
                    kf_class_thresholds.append(thresholds[threshold_best_index])
                threshold = np.mean(kf_class_thresholds)
                classes_thresholds.append(threshold)
                f1 = f1_score(y_true[:,i], np.int32(y_pred[:,i] >= threshold))
                classes_scores.append(f1)
            else:
                t_min = np.min(y_pred[:,i])
                t_max = np.max(y_pred[:,i])
                thresholds = np.linspace(t_min, t_max, 50)
                scores = np.array([
                    f1_score(y_true[:,i], np.int32(y_pred[:,i] >= threshold)) for threshold in thresholds
                ])
                threshold_best_index = np.argmax(scores)
                classes_thresholds.append(thresholds[threshold_best_index])
                f1 = f1_score(y_true[:,i], np.int32(y_pred[:,i] >= thresholds[threshold_best_index]))
                classes_scores.append(f1)
        return classes_thresholds, classes_scores

    import pretrainedmodels
    import pretrainedmodels.utils as pqutils

    _model_name = 'se_resnext50_32x4d'
    model = pretrainedmodels.__dict__[_model_name](num_classes=1000, pretrained='imagenet')
    tf_img = pqutils.TransformImage(model)
    model_stats = (tf_img.__dict__['mean'], tf_img.__dict__['std'])
    data_dir = '../../../input/'
    valid_df = pd.read_csv('../../../input/' + 'val_id.csv', header=None, names=['idx','Id'])
    train_df = pd.read_csv(data_dir + 'train.csv')
    test_matches = pd.read_csv('../../../input/test_matches.csv')
    test_matches.Extra = test_matches['Extra'].apply(lambda x : "_".join(x.split("_")[2:]))
    xtra_data = pd.read_csv(data_dir + 'xtra_train.csv')
    xtra_data['Extra'] = xtra_data.Id.apply(lambda x : x[:x.find("_classes")])

    xtra_matches_ids = test_matches.Extra.values.tolist()
    xtra_data_train = xtra_data.loc[~xtra_data.Extra.isin(xtra_matches_ids),['Id','Target']].reset_index(drop=True)
    xtra_data_valid = xtra_data.loc[xtra_data.Extra.isin(xtra_matches_ids),['Id','Target']].reset_index(drop=True)
    data = xtra_data_train
    labels = np.zeros((data.shape[0], 28), dtype=np.int32)
    if "Target" in data:
        for i, lbls in data['Target'].str.split().iteritems():
            for j in map(int, lbls):
                labels[i, j] = 1

    xtra_matches_ids = ['1054_E4_1_classes_25_16_0','1762_G4_5_classes_27','1335_C6_2_classes_3',
                        '935_D5_2_classes_22_0','27_H9_2_classes_10','669_D8_1_classes_16_2',
                        '1178_D4_2_classes_19_16_14','791_A9_1_classes_10_9','759_F9_9_classes_25_21_19_16',
                        '1283_F10_2_classes_16_0','688_E7_10_classes_23','1772_F9_7_classes_25_17',
                        '454_E5_1_classes_14_0','1020_C5_3_classes_23','1386_G4_2_classes_8',
                        '681_G8_5_classes_13','1609_C4_2_classes_16_0','690_D3_5_classes_22_21_1_0',
                        '1245_B2_2_classes_21_0','1335_C10_4_classes_16_0','693_A11_3_classes_23',
                        '1139_A12_4_classes_23','916_F8_1_classes_25_2_0','694_C1_2_classes_18_1',
                        '929_B8_1_classes_25_19','340_F5_3_classes_13','138_B12_1_classes_8',
                        '932_G11_2_classes_25_16','28_H9_1_classes_10','924_F12_1_classes_27',
                        '682_F12_2_classes_25_4','1147_D3_13_classes_16_0','346_A5_1_classes_12',
                        '616_F1_4_classes_8','73_A10_1_classes_27_25','663_A9_2_classes_16_14',
                        '859_C8_4_classes_16_14','933_C10_4_classes_22_21','1207_B10_7_classes_12',
                        '694_F10_1_classes_25_21','908_E3_1_classes_4','1758_C9_4_classes_17_2',
                        '1335_D2_2_classes_2_0','929_H2_2_classes_23','1717_G8_34_classes_25_17',
                        '1150_H4_7_classes_13','1054_E4_2_classes_25_16_0','504_B1_3_classes_25_16_0',
                        '747_B5_4_classes_10_9','1020_B1_7_classes_23_5','918_H10_2_classes_25_15',
                        '532_H3_1_classes_25_16_0','757_C6_3_classes_16_2','1346_H6_3_classes_16_5_0',
                        '496_D1_1_classes_16_0','1042_C3_3_classes_27','929_B12_1_classes_3',
                        '684_C4_2_classes_23_0','696_C9_5_classes_25_21_0','1144_A10_4_classes_2','846_A8_2_classes_16_14','903_F12_2_classes_23_5','1264_G1_1_classes_27','925_H8_2_classes_1_0','121_C6_2_classes_10_9','1657_E10_3_classes_25_17','932_G11_1_classes_25_16','704_G4_1_classes_25_12','1039_C3_2_classes_19_16','906_H7_2_classes_25_6','19_H7_2_classes_8','725_G10_2_classes_16_14','681_B2_4_classes_4','697_A6_4_classes_19_0','1581_B12_2_classes_16_14','926_F7_2_classes_5_0','1770_D2_4_classes_21_17_4','1037_F4_3_classes_19','1413_F11_6_classes_21_16','694_A2_1_classes_2','1049_D11_2_classes_25_16_0','1276_C3_2_classes_21_0','346_B12_3_classes_14_0','1773_G12_3_classes_16_12','1183_F4_2_classes_15','1158_H11_8_classes_16_5','380_C6_1_classes_16_0','792_B6_7_classes_13_0','682_C9_6_classes_25_12_2','906_A9_4_classes_20_0','400_D3_2_classes_25_7','1237_G1_4_classes_21_6','793_B1_1_classes_25_22_0','1308_A5_4_classes_5','800_E1_1_classes_16_14','1421_G5_7_classes_17','906_A9_6_classes_20_0','1245_B2_3_classes_21_0','626_D7_6_classes_25_21_12','344_G2_4_classes_11','901_E12_1_classes_25_6_2','1050_F6_6_classes_16_0','240_G8_1_classes_8','933_C2_1_classes_23_2_0','556_B9_1_classes_25_18_0','1335_C10_2_classes_16_0','1125_F6_3_classes_4','1495_F7_3_classes_7_0','694_C1_1_classes_18_1','918_B3_4_classes_14','1762_E6_5_classes_7','915_C6_5_classes_4','820_G4_3_classes_10_9','927_F12_12_classes_18_0','901_D10_2_classes_12_0','1642_G7_34_classes_25_16','928_G1_2_classes_14_7','682_G9_1_classes_7_0','903_F2_1_classes_2_0','1645_E1_32_classes_16_14','685_G10_5_classes_12_0','927_A9_10_classes_25_5','957_G6_4_classes_16','757_C6_2_classes_16_2','1213_C4_2_classes_4','909_A6_1_classes_2','694_D6_2_classes_1_0','480_D6_3_classes_25_16','1050_F1_3_classes_25_16_0','692_A1_5_classes_25_14_0','1772_H1_5_classes_18_17_16_0','991_G6_7_classes_10_9','782_F8_2_classes_25_16','693_H4_1_classes_7','1259_A11_4_classes_19_16','1414_D12_2_classes_21_0','1139_D5_5_classes_5','930_H3_2_classes_1','901_G9_5_classes_25_19_0','1754_G2_34_classes_5','353_A9_1_classes_21_13','1179_H7_1_classes_25_16_0','1423_A4_2_classes_16_14','686_F4_2_classes_22_21','1693_E1_2_classes_23_16','400_H8_2_classes_23','1680_G4_4_classes_16','935_G3_1_classes_5','838_E8_1_classes_3','1030_D8_2_classes_7_0','684_D12_4_classes_18','812_C10_2_classes_13_0','1416_D10_6_classes_21_16_0','1293_E3_2_classes_1_0','480_D6_2_classes_25_16','700_H6_2_classes_25_2','1773_E10_4_classes_16_0','611_E10_1_classes_25_13','346_B12_4_classes_14_0','523_A9_4_classes_5','1581_B12_3_classes_16_14','684_D8_6_classes_25_12_0','927_F12_11_classes_18_0','353_E4_2_classes_5','556_C1_5_classes_25_22_16','1179_H7_2_classes_25_16_0','1711_B12_3_classes_26_21_4','449_G8_2_classes_4_2','544_A8_5_classes_22_21_7','1772_H1_3_classes_18_17_16_0','1772_G2_6_classes_25_19_16_0','909_C11_2_classes_2_0','930_C12_1_classes_18_14_6','690_C10_2_classes_13','1009_B6_2_classes_10_9','757_E10_5_classes_12','88_D7_2_classes_8','383_E8_7_classes_25_17','1432_F2_2_classes_6','505_C10_1_classes_25_15','1104_E7_2_classes_16_14','699_E8_1_classes_1','1213_C4_3_classes_4','690_H5_1_classes_4','1169_D3_6_classes_16_0','686_F4_1_classes_22_21','532_D1_1_classes_16_0','896_G8_3_classes_5_0','934_G4_3_classes_21','344_G2_1_classes_11','369_C9_1_classes_18_14_0','682_F12_1_classes_25_4','683_E1_2_classes_25_1_0','697_G3_6_classes_13_7','1772_A6_7_classes_5','933_C4_6_classes_5','1231_F9_5_classes_7','802_D5_9_classes_16_0','682_G10_1_classes_7','850_C1_9_classes_21_0','929_B12_2_classes_3','1339_D3_3_classes_2_1','858_D4_2_classes_4','334_B12_2_classes_4','622_F1_7_classes_8','908_G5_2_classes_2_0','778_G6_2_classes_25_16_14','1027_C4_1_classes_7','886_C10_5_classes_23_0','807_C2_3_classes_4','1314_D2_2_classes_25_16_0','1770_B5_1_classes_21_16_11','1105_F10_2_classes_16_0','1283_B2_10_classes_16_0','583_E11_1_classes_25_16','820_G4_7_classes_10_9','928_H3_2_classes_14_0','970_H1_4_classes_25_18','1751_A7_32_classes_27','701_H10_2_classes_25_14','1773_B6_11_classes_23_17_16','1736_G7_31_classes_25_16','928_H3_1_classes_14_0','1645_E5_34_classes_17','539_B3_1_classes_25_21_0','683_E1_1_classes_25_1_0','484_G6_3_classes_22','928_A1_1_classes_4','1773_B6_7_classes_23_17_16','1255_A3_4_classes_16_0','698_C6_2_classes_25_21_4','1773_D5_6_classes_17','681_G8_4_classes_13','935_H11_2_classes_22_0','1125_B9_4_classes_25_7','698_F11_1_classes_13_0','344_F7_1_classes_25_21','906_C11_1_classes_4','1656_F5_2_classes_19_17','1761_A10_3_classes_23_17_14','1772_H5_7_classes_17_7','910_B8_1_classes_12_0','1283_F10_4_classes_16_0','508_C10_1_classes_25_15','681_B2_3_classes_4','868_E8_2_classes_17_16_0','1339_B9_2_classes_16_0','856_A2_4_classes_2_0','700_C3_6_classes_21','869_B3_1_classes_16_0','701_B9_2_classes_21_13_0','1178_F9_6_classes_16_0','542_G1_1_classes_11_2_0']

    xtra_data_train = xtra_data.loc[~xtra_data.Id.isin(xtra_matches_ids),['Id','Target']].reset_index(drop=True)
    xtra_data_valid = xtra_data.loc[xtra_data.Id.isin(xtra_matches_ids),['Id','Target']].reset_index(drop=True)

    exclude_valid = \
    ['5ae3db3a-bbc4-11e8-b2bc-ac1f6b6435d0',
     'e6d0b648-bbbc-11e8-b2ba-ac1f6b6435d0',
     '3202385a-bbca-11e8-b2bc-ac1f6b6435d0',
     '0cf36c82-bbca-11e8-b2bc-ac1f6b6435d0',
     '7cb0006e-bbaf-11e8-b2ba-ac1f6b6435d0',
     '87b77dd2-bba2-11e8-b2b9-ac1f6b6435d0',
     '62c88efa-bbc8-11e8-b2bc-ac1f6b6435d0',
     '44d819c2-bbbb-11e8-b2ba-ac1f6b6435d0',
     'b1ca2b40-bbbd-11e8-b2ba-ac1f6b6435d0',
     '8cd67266-bbbe-11e8-b2ba-ac1f6b6435d0',
     'cead83ec-bb9a-11e8-b2b9-ac1f6b6435d0',
     'a166d11a-bbca-11e8-b2bc-ac1f6b6435d0',
     '91a0a67e-bb9e-11e8-b2b9-ac1f6b6435d0',
     '2be24582-bbb1-11e8-b2ba-ac1f6b6435d0']
    exclude_train = \
    ['7138c4aa-bb9b-11e8-b2b9-ac1f6b6435d0',
     '8a10533e-bba6-11e8-b2ba-ac1f6b6435d0',
     'be92e108-bbb5-11e8-b2ba-ac1f6b6435d0',
     'abfa727e-bba4-11e8-b2ba-ac1f6b6435d0',
     '2384acac-bbae-11e8-b2ba-ac1f6b6435d0',
     'c7a7a462-bbb1-11e8-b2ba-ac1f6b6435d0',
     '559f7ce0-bbb2-11e8-b2ba-ac1f6b6435d0']

    valid_df = pd.read_csv('../../../input/' + 'val_id.csv', header=None, names=['idx','Id'])
    valid_df = valid_df.loc[~valid_df.Id.isin(exclude_valid),:]
    train_df = pd.read_csv(data_dir + 'train.csv')
    train_df = train_df.loc[~train_df.Id.isin(exclude_train),:]

    test_df = pd.read_csv('../../../input/' + "sample_submission.csv")
    train = train_df.loc[~train_df.Id.isin(valid_df.Id.values.tolist()),:].reset_index(drop=True)
    train = pd.concat([train,xtra_data_train], axis=0, sort=False)
    valid = train_df.loc[train_df.Id.isin(valid_df.Id.values.tolist()),:].reset_index(drop=True)
    valid = pd.concat([valid,xtra_data_valid], axis=0, sort=False)
    test  = test_df
    del train_df,valid_df,test_df,xtra_data_valid,xtra_data_train
    gc.collect()

    train_files = train.Id.apply(lambda s: '../../../input/' + 'train_png/'+s+'.png')
    train_labels = train.Target.astype(str).apply(lambda s : [name_label_dict[int(q)] for q in s.split(' ')])
    train_ds = ImageMultiDataset(fns = train_files, labels = train_labels, classes = list(name_label_dict.values()))
    del train_files, train_labels

    valid_files = valid.Id.apply(lambda s: '../../../input/' + 'train_png/'+s+'.png')
    valid_labels = valid.Target.astype(str).apply(lambda s : [name_label_dict[int(q)] for q in s.split(' ')])
    valid_ds = ImageMultiDataset(fns = valid_files, labels = valid_labels, classes = list(name_label_dict.values()))
    del valid_files, valid_labels

    test_files = test.Id.apply(lambda s: '../../../input/' + 'test_png/'+s+'.png')
    test_labels = test.Predicted.astype(str).apply(lambda s : [name_label_dict[int(q)] for q in s.split(' ')])
    test_ds = ImageMultiDataset(fns = test_files, labels = test_labels, classes = list(name_label_dict.values()))
    del test_files, test_labels

    magnitudes = torch.tensor(np.zeros(8))
    magnitudes[6]=0.5
    xtra = [RandTransform(squish, {}),jitter(magnitude=0.03,p=0.1),perspective_warp(magnitude=0.03,p=0.1)]
    tfms = get_transforms(do_flip=True, flip_vert=True, max_rotate=180.0,
                          max_zoom=1.2, max_lighting=0.2, max_warp=None,
                          p_affine=0.9, p_lighting=0.65, xtra_tfms=xtra)
    data = ImageDataBunch.create(train_ds, valid_ds, test_ds, path=data_dir, device=device,
                                 size=512, bs=12, ds_tfms=tfms, padding_mode='zeros')
    data.normalize(model_stats)

    from torch.autograd import Variable, Function

    def _to_one_hot(y, n_dims, dtype=torch.cuda.FloatTensor):
        scatter_dim = len(y.size())
        y_tensor = y.type(torch.cuda.LongTensor).view(*y.size(), -1)
        zeros = torch.zeros(*y.size(), n_dims).type(dtype)

        return zeros.scatter(scatter_dim, y_tensor, 1)

    class LSEP2(Function):
        """
        Autograd function of LSEP loss. Appropirate for multi-label
        - Reference: Li+2017
          https://arxiv.org/pdf/1704.03135.pdf
          Code-ref: https://github.com/Mipanox/Bird_cocktail/blob/196e9404a4f7022d1e56433112f581b34a334e53/utils.py#L332
        """

        @staticmethod
        def forward(ctx, input, target):
            batch_size = target.size()[0]
            label_size = target.size()[1]

            ##
            positive_indices = target.gt(0).float()
            negative_indices = target.eq(0).float()

            ## summing over all negatives and positives
            loss = 0.
            for i in range(input.size()[0]):
                pos = positive_indices[i].nonzero()
                neg = negative_indices[i].nonzero()
                pos_examples = input[i, pos]
                neg_examples = torch.transpose(input[i, neg], 0, 1)
                loss += torch.sum(torch.exp(neg_examples - pos_examples))

            loss = torch.log(1 + loss)

            ctx.save_for_backward(input, target)
            ctx.loss = loss
            ctx.positive_indices = positive_indices
            ctx.negative_indices = negative_indices

            return loss

        # This function has only a single output, so it gets only one gradient
        @staticmethod
        def backward(ctx, grad_output):
            dtype = torch.cuda.FloatTensor
            input, target = ctx.saved_variables
            N = input.size()[1]
            loss = Variable(ctx.loss, requires_grad = False)
            positive_indices = ctx.positive_indices
            negative_indices = ctx.negative_indices

            fac  = -1 / loss
            grad_input = torch.zeros(input.size()).type(dtype)

            scale = grad_input.size(0), -1
            phot = _to_one_hot(positive_indices.nonzero()[:, 1].view(*scale), N)
            nhot = _to_one_hot(negative_indices.nonzero()[:, 1].view(*scale), N)

            scale = (len(phot), *nhot.size())
            diffs = torch.sum(phot - nhot.expand(scale), dim=2)
            grads_input = (Variable(diffs * torch.exp(-input * diffs)) * (grad_output * fac))

            return grad_input, None, None

        class FocalLoss(nn.Module):
            def __init__(self, gamma=2):
                super().__init__()
                self.gamma = gamma

            def forward(self, input, target):
                if not (target.size() == input.size()):
                    raise ValueError("Target size ({}) must be the same as input size ({})"
                                     .format(target.size(), input.size()))

                max_val = (-input).clamp(min=0)
                loss = input - input * target + max_val + \
                    ((-max_val).exp() + (-input - max_val).exp()).log()

                invprobs = F.logsigmoid(-input * (target * 2.0 - 1.0))
                loss = (invprobs * self.gamma).exp() * loss

                return loss.sum(dim=1).mean()
    from fastai.vision import models
    from sklearn import metrics
    import torchvision

    q = None

    class F1Metric(Callback):
        def on_epoch_begin(self, **kwargs):
            self.tp, self.sum_true, self.sum_pred = torch.zeros([28]).cuda(device),\
    torch.zeros([28]).cuda(device), torch.zeros([28]).cuda(device)
            self.threshold = 0.5
            self.metric = 0

        def on_batch_end(self, last_output, last_target, **kwargs):
            y_pred = torch.ge(last_output, self.threshold).float()
            y_true = last_target.float()

            self.tp = self.tp + (y_pred * y_true).sum(dim=0)
            self.sum_true = self.sum_true + (y_true).sum(dim=0)
            self.sum_pred = self.sum_pred + (y_pred).sum(dim=0)

        def on_epoch_end(self, **kwargs):
            global q
            eps = 1e-9
            q = self.tp
            true_positive = self.tp
            precision = true_positive.div(self.sum_true.add(eps))
            recall = true_positive.div(self.sum_pred.add(eps))

            self.metric = ((precision*recall).div(precision + recall + eps).mul(2)).mean()

    f1metric = F1Metric()

    def body(pretrained = True):
        return pretrainedmodels.__dict__[_model_name](num_classes=1000, pretrained='imagenet').to(device)

    learner = create_cnn(data, arch=body, cut=-2,
                         custom_head = create_head(4096, len(data.classes), ps=0.3), metrics=[f1metric]) #
    learner.loss_fn = FocalLoss()

    import collections

    layers = []

    for i, l1 in enumerate(learner.model):
        for j, l2 in enumerate(learner.model[i]):
            if isinstance(learner.model[i][j], collections.Iterable):
                for k, l3 in enumerate(learner.model[i][j]):
                    print(l3)
                    layers.append(learner.model[i][j][k])
            else:
                    layers.append(learner.model[i][j])

    learner.split(layers)

    learner.load('pd_seresnext_v3_0.7062.pnt')

    num_iter=0
    def get_lrs2(lr):
        global num_iter
        num_iter=num_iter+17
        np.random.seed(num_iter)
        return np.logspace(np.log(lr/16), np.log(lr), len(learner.layer_groups), base=np.e)+\
                    np.random.normal(0,lr*1e-2,size=len(learner.layer_groups))

    learner.loss_fn = LSEP2()

    lr = 3e-5
    learner.fit_one_cycle(4, max_lr=get_lrs2(lr))

    y_pred_solo, avg_preds1, y = learner.TTA(beta=None)
    y = y.cpu().numpy().copy()
    _, avg_preds2, _ = learner.TTA(beta=None)
    _, avg_preds3, _ = learner.TTA(beta=None)
    _, avg_preds4, _ = learner.TTA(beta=None)

    avg_preds = y_pred_solo.cpu().numpy().copy()*0.4+\
    torch.stack([avg_preds1,avg_preds2,avg_preds3,avg_preds4]).mean(0).cpu().numpy().copy()*0.6

    results_dir = '../../../results/'
    np.save(results_dir+model_name+'_y.npy', y)
    np.save(results_dir+model_name+'_ids.npy', valid.Id.values)
    np.save(results_dir+model_name+'_holdout.npy', avg_preds)


import argparse
import os
import os.path as osp
import torch
import mmcv
from mmaction.apis import init_recognizer
from mmcv.parallel import collate, scatter
from mmaction.datasets.pipelines import Compose
from mmaction.datasets import build_dataloader, build_dataset
from mmcv.parallel import MMDataParallel
import numpy as np
from tqdm import tqdm
from sklearn.metrics import f1_score, roc_auc_score, accuracy_score
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

def parse_args():
    """
    experiments/baseline_rpl.py --config configs/recognition/tsm/inference_tsm_rpl.py \
    --checkpoint work_dirs/tsm/finetune_ucf101_tsm_rpl/latest.pth \
    --train_data data/ucf101/ucf101_train_split_1_videos.txt \
    --ind_data
    --result_prefix experiments/tsm/results_baselines/rpl/RPL
    """
    parser = argparse.ArgumentParser(description='MMAction2 test')
    # model and data config
    parser.add_argument('--config', help='test config file path')
    parser.add_argument('--checkpoint', help='checkpoint file/url')
    parser.add_argument('--train_data', help='the split file of in-distribution training data')
    parser.add_argument('--batch_size', type=int, default=8, help='the testing batch size')
    # test data config
    parser.add_argument('--ind_data', help='the split file of in-distribution testing data')
    parser.add_argument('--ood_data', help='the split file of out-of-distribution testing data')
    parser.add_argument('--ood_ncls', type=int, help='the number of classes in unknwon dataset')
    parser.add_argument('--ood_dataname', choices=['HMDB', 'MiT'], help='the name of out-of-distribution testing data')
    # env config
    parser.add_argument('--device', type=str, default='cuda:0', help='CPU/CUDA device option')
    parser.add_argument('--result_prefix', help='result file prefix')
    args = parser.parse_args()
    return args

def set_deterministic(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    np.random.seed(seed)  # Numpy module.
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True

def run_inference(config, checkpoint, data_split, batch_size, device):
    # initialize recognition model
    model = init_recognizer(config, checkpoint, device=device, use_frames=False)
    torch.backends.cudnn.benchmark = True
    model.cfg.data.test.test_mode = True
    model.cfg.test_cfg.average_clips = 'prob'  # we need the probability socore from softmax layer
    model.cfg.data.videos_per_gpu = batch_size  # batch size
    model.cfg.data.test.ann_file = data_split
    model.cfg.data.test.data_prefix = os.path.join(os.path.dirname(data_split), 'videos')

    # build the dataloader
    dataset = build_dataset(model.cfg.data.test, dict(test_mode=True))
    dataloader_setting = dict(
        videos_per_gpu=model.cfg.data.get('videos_per_gpu', 1),
        workers_per_gpu=model.cfg.data.get('workers_per_gpu', 1),
        dist=False,
        shuffle=False,
        pin_memory=False)
    dataloader_setting = dict(dataloader_setting, **model.cfg.data.get('test_dataloader', {}))
    data_loader = build_dataloader(dataset, **dataloader_setting)

    # running the inference
    model = MMDataParallel(model, device_ids=[0])
    all_scores, all_labels = [], []
    prog_bar = mmcv.ProgressBar(len(data_loader.dataset))
    for i, data in enumerate(data_loader):
        with torch.no_grad():
            scores = model(return_loss=False, **data)  # (B, C)
            all_scores.append(scores)
        # gather labels
        labels = data['label'].numpy()
        all_labels.append(labels)
        # use the first key as main key to calculate the batch size
        bs = len(next(iter(data.values())))
        for _ in range(bs):
            prog_bar.update()
    all_scores = np.concatenate(all_scores, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)
    return all_scores, all_labels

def evaluate_softmax(ind_softmax, ood_softmax, ind_labels, ood_labels, ood_ncls, thresh, num_rand=10):

    ind_ncls = ind_softmax.shape[1]
    ind_results = np.argmax(ind_softmax, axis=1)
    ood_results = np.argmax(ood_softmax, axis=1)
    # close-set accuracy (multi-class)
    acc = accuracy_score(ind_labels, ind_results)

    # open-set auc-roc (binary class)
    ind_conf = np.max(ind_softmax, axis=1)
    ood_conf = np.max(ood_softmax, axis=1)
    preds = np.concatenate((ind_results, ood_results), axis=0)
    confs = np.concatenate((ind_conf, ood_conf), axis=0)
    preds[confs < threshold] = 1  # unknown class
    preds[confs >= threshold] = 0  # known class
    labels = np.concatenate((np.zeros_like(ind_labels), np.ones_like(ood_labels)))
    auc = roc_auc_score(labels, preds)
    print('\nClosedSet Accuracy (multi-class): %.3lf, OpenSet AUC (bin-class): %.3lf'%(acc * 100, auc * 100))

    ind_results[ind_conf < thresh] = ind_ncls  # incorrect rejection
    # open set F1 score (multi-class)
    macro_F1_list = [f1_score(ind_labels, ind_results, average='macro')]
    std_list = [0]
    openness_list = [0]
    for n in range(ood_ncls):
        ncls_novel = n + 1
        openness = (1 - np.sqrt((2 * ind_ncls) / (2 * ind_ncls + ncls_novel))) * 100
        openness_list.append(openness)
        # randoml select the subset of ood samples
        macro_F1_multi = np.zeros((num_rand), dtype=np.float32)
        for m in range(num_rand):
            cls_select = np.random.choice(ood_ncls, ncls_novel, replace=False)
            ood_sub_results = np.concatenate([ood_results[ood_labels == clsid] for clsid in cls_select])
            ood_sub_labels = np.ones_like(ood_sub_results) * ind_ncls
            ood_sub_confs = np.concatenate([ood_conf[ood_labels == clsid] for clsid in cls_select])
            ood_sub_results[ood_sub_confs < thresh] = ind_ncls  # correct rejection
            # construct preds and labels
            preds = np.concatenate((ind_results, ood_sub_results), axis=0)
            labels = np.concatenate((ind_labels, ood_sub_labels), axis=0)
            macro_F1_multi[m] = f1_score(labels, preds, average='macro')
        macro_F1 = np.mean(macro_F1_multi)
        std = np.std(macro_F1_multi)
        macro_F1_list.append(macro_F1)
        std_list.append(std)

    # draw comparison curves
    macro_F1_list = np.array(macro_F1_list)
    std_list = np.array(std_list)

    w_openness = np.array(openness_list) / 100.
    open_maF1_mean = np.sum(w_openness * macro_F1_list) / np.sum(w_openness)
    open_maF1_std = np.sum(w_openness * std_list) / np.sum(w_openness)
    print('Open macro-F1 score: %.3f, std=%.3lf'%(open_maF1_mean * 100, open_maF1_std * 100))

    return openness_list, macro_F1_list, std_list

if __name__ == '__main__':

    args = parse_args()
    # assign the desired device.
    device = torch.device(args.device)
    set_deterministic(0)
    modelname = os.path.dirname(args.config).split('/')[-1].upper()

    ######## Compute threshold with training data ########
    result_file = os.path.join(os.path.dirname(args.result_prefix), modelname + '_RPL_trainset_softmax.npz')
    if not os.path.exists(result_file):
        # prepare result path
        result_dir = os.path.dirname(result_file)
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        # run the inference on training data
        trainset_scores, _ = run_inference(args.config, args.checkpoint, args.train_data, args.batch_size, device)
        # save
        np.savez(result_file[:-4], trainset_scores=trainset_scores)
    else:
        result = np.load(result_file)
        trainset_scores = result['trainset_scores']

    max_scores = np.max(trainset_scores, axis=1)
    scores_sort = np.sort(max_scores)[::-1]  # sort the uncertainties with descending order
    N = max_scores.shape[0]
    threshold = scores_sort[int(N * 0.95)-1]  # 95% percentile

    print('\nThe RPL softmax threshold on UCF-101 train set: %lf'%(threshold))

    ######## OOD and IND detection ########
    testset_result = os.path.join(os.path.dirname(args.result_prefix), modelname +'_RPL_'+ args.ood_dataname +'_result.npz')
    if not os.path.exists(testset_result):
        # prepare result path
        result_dir = os.path.dirname(testset_result)
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        # run the inference on OOD data
        ood_softmax, ood_labels = run_inference(args.config, args.checkpoint, args.ood_data, args.batch_size, device)
        # run the inference on IND data
        ind_softmax, ind_labels = run_inference(args.config, args.checkpoint, args.ind_data, args.batch_size, device)
        # save
        np.savez(testset_result[:-4], ind_softmax=ind_softmax, ood_softmax=ood_softmax,
                                   ind_label=ind_labels, ood_label=ood_labels)
    else:
        results = np.load(testset_result, allow_pickle=True)
        ind_softmax = results['ind_softmax']  # (N1, C)
        ood_softmax = results['ood_softmax']  # (N2, C)
        ind_labels = results['ind_label']  # (N1,)
        ood_labels = results['ood_label']  # (N2,)

    openness_list, macro_F1_list, std_list = evaluate_softmax(ind_softmax, ood_softmax, ind_labels, ood_labels, args.ood_ncls, threshold)


from config import COCO_PATH, IM_SCALE, BOX_SCALE
import os
from torch.utils.data import Dataset
from pycocotools.coco import COCO
from PIL import Image
from lib.fpn.anchor_targets import anchor_target_layer
from torchvision.transforms import Resize, Compose, ToTensor, Normalize
from dataloaders.image_transforms import SquarePad, Grayscale, Brightness, Sharpness, Contrast, \
    RandomOrder, Hue, random_crop
import numpy as np
from dataloaders.blob import Blob
import torch

class CocoDetection(Dataset):
    """
    Adapted from the torchvision code
    """

    def __init__(self, mode):
        """
        :param mode: train2014 or val2014
        """
        self.mode = mode
        self.root = os.path.join(COCO_PATH, mode)
        self.ann_file = os.path.join(COCO_PATH, 'annotations', 'instances_{}.json'.format(mode))
        self.coco = COCO(self.ann_file)
        self.ids = [k for k in self.coco.imgs.keys() if len(self.coco.imgToAnns[k]) > 0]

        tform = []
        if self.is_train:
            tform.append(RandomOrder([
                Grayscale(),
                Brightness(),
                Contrast(),
                Sharpness(),
                Hue(),
            ]))

        tform += [
            SquarePad(),
            Resize(IM_SCALE),
            ToTensor(),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]

        self.transform_pipeline = Compose(tform)
        self.ind_to_classes = ['__background__'] + [v['name'] for k, v in self.coco.cats.items()]
        # COCO inds are weird (84 inds in total but a bunch of numbers are skipped)
        self.id_to_ind = {coco_id: (ind + 1) for ind, coco_id in enumerate(self.coco.cats.keys())}
        self.id_to_ind[0] = 0

        self.ind_to_id = {x: y for y, x in self.id_to_ind.items()}

    @property
    def is_train(self):
        return self.mode.startswith('train')

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns: entry dict
        """
        img_id = self.ids[index]
        path = self.coco.loadImgs(img_id)[0]['file_name']
        image_unpadded = Image.open(os.path.join(self.root, path)).convert('RGB')
        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        anns = self.coco.loadAnns(ann_ids)
        gt_classes = np.array([self.id_to_ind[x['category_id']] for x in anns], dtype=np.int64)

        if np.any(gt_classes >= len(self.ind_to_classes)):
            raise ValueError("OH NO {}".format(index))

        if len(anns) == 0:
            raise ValueError("Annotations should not be empty")
        #     gt_boxes = np.array((0, 4), dtype=np.float32)
        # else:
        gt_boxes = np.array([x['bbox'] for x in anns], dtype=np.float32)

        if np.any(gt_boxes[:, [0, 1]] < 0):
            raise ValueError("GT boxes empty columns")
        if np.any(gt_boxes[:, [2, 3]] < 0):
            raise ValueError("GT boxes empty h/w")
        gt_boxes[:, [2, 3]] += gt_boxes[:, [0, 1]]

        # Rescale so that the boxes are at BOX_SCALE
        if self.is_train:
            image_unpadded, gt_boxes = random_crop(image_unpadded,
                                                   gt_boxes * BOX_SCALE / max(image_unpadded.size),
                                                   BOX_SCALE,
                                                   round_boxes=False,
                                                   )
        else:
            # Seems a bit silly because we won't be using GT boxes then but whatever
            gt_boxes = gt_boxes * BOX_SCALE / max(image_unpadded.size)
        w, h = image_unpadded.size
        box_scale_factor = BOX_SCALE / max(w, h)

        # Optionally flip the image if we're doing training
        flipped = self.is_train and np.random.random() > 0.5
        if flipped:
            scaled_w = int(box_scale_factor * float(w))
            image_unpadded = image_unpadded.transpose(Image.FLIP_LEFT_RIGHT)
            gt_boxes[:, [0, 2]] = scaled_w - gt_boxes[:, [2, 0]]

        img_scale_factor = IM_SCALE / max(w, h)
        if h > w:
            im_size = (IM_SCALE, int(w * img_scale_factor), img_scale_factor)
        elif h < w:
            im_size = (int(h * img_scale_factor), IM_SCALE, img_scale_factor)
        else:
            im_size = (IM_SCALE, IM_SCALE, img_scale_factor)

        entry = {
            'img': self.transform_pipeline(image_unpadded),
            'img_size': im_size,
            'gt_boxes': gt_boxes,
            'gt_classes': gt_classes,
            'scale': IM_SCALE / BOX_SCALE,
            'index': index,
            'image_id': img_id,
            'flipped': flipped,
            'fn': path,
        }

        return entry

    @classmethod
    def splits(cls, *args, **kwargs):
        """ Helper method to generate splits of the dataset"""
        train = cls('train2014', *args, **kwargs)
        val = cls('val2014', *args, **kwargs)
        return train, val

    def __len__(self):
        return len(self.ids)

def coco_collate(data, num_gpus=3, is_train=False):
    blob = Blob(mode='det', is_train=is_train, num_gpus=num_gpus,
                batch_size_per_gpu=len(data) // num_gpus)
    for d in data:
        blob.append(d)
    blob.reduce()
    return blob

class CocoDataLoader(torch.utils.data.DataLoader):
    """
    Iterates through the data, filtering out None,
     but also loads everything as a (cuda) variable
    """

    # def __iter__(self):
    #     for x in super(CocoDataLoader, self).__iter__():
    #         if isinstance(x, tuple) or isinstance(x, list):
    #             yield tuple(y.cuda(async=True) if hasattr(y, 'cuda') else y for y in x)
    #         else:
    #             yield x.cuda(async=True)

    @classmethod
    def splits(cls, train_data, val_data, batch_size=3, num_workers=1, num_gpus=3, **kwargs):
        train_load = cls(
            dataset=train_data,
            batch_size=batch_size * num_gpus,
            shuffle=True,
            num_workers=num_workers,
            collate_fn=lambda x: coco_collate(x, num_gpus=num_gpus, is_train=True),
            drop_last=True,
            # pin_memory=True,
            **kwargs,
        )
        val_load = cls(
            dataset=val_data,
            batch_size=batch_size * num_gpus,
            shuffle=False,
            num_workers=num_workers,
            collate_fn=lambda x: coco_collate(x, num_gpus=num_gpus, is_train=False),
            drop_last=True,
            # pin_memory=True,
            **kwargs,
        )
        return train_load, val_load

if __name__ == '__main__':
    train, val = CocoDetection.splits()
    gtbox = train[0]['gt_boxes']
    img_size = train[0]['img_size']
    anchor_strides, labels, bbox_targets = anchor_target_layer(gtbox, img_size)