from __future__ import division

import warnings

from math import sqrt

import networkx as nx
import parmed as pmd

__all__ = ['write_mcf']

def write_mcf(structure, filename, angle_style,
                      dihedral_style, lj14=None, coul14=None):
    """Output a Cassandra molecular connectivity file (MCF).

    Outputs a Cassandra MCF from a Parmed structure object.

    Parameters
    ----------
    structure : parmed.Structure
        ParmEd structure object
    filename : str
        Path of the output file
    angle_style : str
        Type of angles. 'fixed' and 'harmonic'
        are valid choices
    dihedral_style : str
        Type of dihedrals. 'harmonic', 'OPLS', 'CHARMM',
        and 'none' are valid choices
    lj14 : float
        Scaling factor for LJ interactions on 1-4 pairs
    coul14 : float
        Scaling factor for Coulombic interactions on 1-4 pairs

    Notes
    -----
    See https://cassandra.nd.edu/index.php/documentation for
    a complete description of the MCF format.

    """
    if not isinstance(structure, pmd.Structure):
        raise ValueError("MCF writer requires parmed structure.")
    if not all ([a.type for a in structure.atoms]):
        raise ValueError("MCF writing not supported without "
            "parameterized forcefield.")

    # Conversion factors
    IG_CONSTANT_KCAL = 0.00198720425864083 # kcal/mol*K
    KCAL_TO_KJ = 4.184

    # Check some things before we start writing the MCF
    # Only will write MCF for Cassandra-supported options
    if (angle_style.casefold() != 'fixed' and
            angle_style.casefold() != 'harmonic'):
        raise ValueError("Invalid selection for angle_style. "
                "Please choose 'fixed' or 'harmonic'")

    if len(structure.urey_bradleys) > 0 :
        raise ValueError("Urey bradley terms detected. Cassandra only "
                        "currently supports fixed or harmonic angles.")

    if (dihedral_style.casefold() != 'opls' and
            dihedral_style.casefold() != 'charmm' and
            dihedral_style.casefold() != 'none'):
        raise ValueError("Invalid selection for dihedral_style. "
                "Please choose 'OPLS', 'CHARMM', or 'none'")

    if dihedral_style.casefold() != 'none':
        if (len(structure.rb_torsions) > 0 and
                dihedral_style.casefold() != 'opls'):
            raise ValueError("Dihedral style declared as {} but "
                "RB torsions found.".format(dihedral_style))

        if (len(structure.dihedrals) > 0 and
                dihedral_style.casefold() != 'charmm'):
            raise ValueError("Dihedral style declared as {} but "
                "charmm-style dihedrals found.".format(dihedral_style))

        if (len(structure.rb_torsions) > 0 and
                len(structure.dihedrals) > 0):
            raise ValueError("Multiple dihedral styles detected, check your "
                             "Forcefield XML and structure")

    # Identify atoms in rings and Cassandra 'fragments'
    in_ring,frag_list,frag_conn = _id_rings_fragments(structure)

    # Infer 1-4 scaling if not specified
    if coul14 is None:
        if len(structure.adjusts) > 0:
            coul14 = structure.adjusts[0].type.chgscale
        else:
            coul14 = 0.0
            if (len(structure.dihedrals) > 0 or
                    len(structure.rb_torsions) > 0):
                warnings.warn('Unable to infer coulombic 1-4 '
                        'scaling factor. Setting to {:.1f}'.format(coul14))
    if lj14 is None:
        if len(structure.adjusts) > 0:
            type1_eps = structure.adjusts[0].atom1.epsilon
            type2_eps = structure.adjusts[0].atom2.epsilon
            scaled_eps = structure.adjusts[0].type.epsilon
            if (structure.combining_rule == 'geometric' or
                    structure.combining_rule == 'lorentz'):
                combined_eps = sqrt(type1_eps*type2_eps)
                lj14 = scaled_eps/combined_eps
            else:
                lj14 = 0.0
                warnings.warn('Unable to infer LJ 1-4 scaling'
                        'factor. Setting to {:.1f}'.format(lj14))
        else:
            lj14 = 0.0
            if (len(structure.dihedrals) > 0 or
                    len(structure.rb_torsions) > 0):
                warnings.warn('Unable to infer LJ 1-4 scaling'
                        'factor. Setting to {:.1f}'.format(lj14))

    if coul14 < 0.0 or coul14 > 1.0:
        raise ValueError("Unreasonable value {} for "
                "coul14 scaling.".format(coul14))
    if lj14 < 0.0 or lj14 > 1.0:
        raise ValueError("Unreasonable value {} for "
                "lj14 scaling.".format(lj14))

    # Now we write the MCF file
    with open(filename, 'w') as mcf_file:

        header = ( '!***************************************'
                   '****************************************\n'
                   '!Molecular connectivity file\n'
                   '!***************************************'
                   '****************************************\n'
                   '!'+filename+' - created by mBuild\n\n'
                 )

        mcf_file.write(header)
        _write_atom_information(mcf_file, structure, in_ring,
                IG_CONSTANT_KCAL)
        _write_bond_information(mcf_file, structure)
        _write_angle_information(mcf_file, structure, angle_style,
                IG_CONSTANT_KCAL)
        _write_dihedral_information(mcf_file, structure, dihedral_style,
                KCAL_TO_KJ)
        _write_improper_information(mcf_file, structure, KCAL_TO_KJ)
        _write_fragment_information(mcf_file, structure, frag_list, frag_conn)
        _write_intrascaling_information(mcf_file, lj14, coul14)

        # That's all, folks!
        mcf_file.write('\n\nEND\n')

def _id_rings_fragments(structure):
    """Identifies the rings and fragments of the molecule

    Parameters
    ----------
    structure : parmed.Structure
        Parmed structure object

    Returns
    ---------
    in_ring : list
        True for each atom in a ring
    frag_list : list
        Atom ids belonging to each fragment
    frag_conn : list
        Fragment ids of connected fragments

    """

    # Identify atoms in rings
    bond_graph = nx.Graph()
    bond_graph.add_edges_from([ [bond.atom1.idx, bond.atom2.idx]
                                 for bond in structure.bonds ])

    if len(structure.bonds) == 0:
        warnings.warn("No bonds found. Cassandra will interpet "
                "this as a rigid species")
        in_ring = [ False ] * len(structure.atoms)
        frag_list = []
        frag_conn = []
        return in_ring, frag_list, frag_conn

    # Check if entire molecule is connected. Warn if not.
    if nx.is_connected(bond_graph) == False:
        raise ValueError("Not all components of the molecule are connected. "
                         "MCF files are for a single molecule and thus "
                         "everything should be connected through bonds.")

    all_rings = nx.cycle_basis(bond_graph)
    in_ring = [False]*bond_graph.number_of_nodes()
    adj_to_ring = [False]*bond_graph.number_of_nodes()
    for ring in all_rings:
        for idx in ring:
            in_ring[idx] = True

    # Identify fragments
    # See Shah and Maginn, JCP, 135, 134121, 2011, doi:10.1063/1.3644939
    frag_list = []
    frag_conn = []

    # First create a neighbor list for each atom
    neigh_dict = {i:list(bond_graph.neighbors(i))
                  for i in range(bond_graph.number_of_nodes())}

    # Handle fused/adjoining rings
    rings_changed = True
    while rings_changed:
        rings_changed = False
        for ring1 in all_rings:
            if rings_changed:
                break
            for ring2 in all_rings:
                if ring1 == ring2:
                    continue
                if len(set(ring1) & set(ring2)) > 0:
                    all_rings.remove(ring1)
                    all_rings.remove(ring2)
                    all_rings.append(list(set(ring1+ring2)))
                    rings_changed=True
                    break

    # ID fragments which contain a ring
    for ring in all_rings:
        adjacentatoms = []
        for idx in ring:
            if len(neigh_dict[idx]) > 2:
                adjacentatoms.append(list(set(neigh_dict[idx])-set(ring)))
        tmp=filter(None, adjacentatoms)
        adjacentatoms = [x for sublist in tmp for x in sublist]
        frag_list.append(ring+adjacentatoms)
        for idx in adjacentatoms:
            adj_to_ring[idx] = True
    # Now ID the other fragments
    for idx in neigh_dict:
        if len(neigh_dict[idx]) > 1:
            if in_ring[idx] == True:
                continue
            else:
                frag_list.append([idx]+neigh_dict[idx])
    # Now find connectivity (shared bonds)
    for i in range(len(frag_list)):
        frag1 = frag_list[i]
        for j in range(i+1, len(frag_list)):
            frag2 = frag_list[j]
            shared_atoms = list(set(frag1) & set(frag2))
            if len(shared_atoms) == 2:
                frag_conn.append([i, j])
            elif len(shared_atoms) > 2:
                warnings.warn('Fragments share more than two atoms...'
                      'something may be going awry unless there are'
                      'fused rings in your system. See below for details.')
                print('Fragment 1 atoms:')
                print(frag1)
                print('Fragment 2 atoms:')
                print(frag2)

    return in_ring, frag_list, frag_conn

def _write_atom_information(mcf_file, structure, in_ring, IG_CONSTANT_KCAL):
    """Write the atoms in the system.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object
    in_ring : list
        Boolean for each atom idx True if atom belongs to a ring
    IG_CONSTANT_KCAL : float
        Ideal gas constant in kcal/mol K

    """

    elements = [atom.element_name for atom in structure.atoms]
    types = [atom.type for atom in structure.atoms]
    masses = [atom.mass for atom in structure.atoms]
    charges = [atom.charge for atom in structure.atoms]
    # Convert energy to units of K
    epsilons = [atom.epsilon/IG_CONSTANT_KCAL for atom in structure.atoms]
    sigmas = [atom.sigma for atom in structure.atoms]

    # Check constraints on atom type length and element name length
    max_element_length = 6
    max_atomtype_length = 20
    n_unique_elements = len(set(elements))
    for element in elements:
        if len(element) > max_element_length:
            warnings.warn("Warning, element name {} will be shortened "
                 "to {} characters. Please confirm your final "
                 "MCF.".format(element, max_element_length))

    elements = [ element[:max_element_length] for element in elements ]
    if len(set(elements)) < n_unique_elements:
        warnings.warn("Warning, the number of unique elements has been "
              "reduced due to shortening the element name to {} "
              "characters.".format(max_element_length))

    n_unique_types = len(set(types))
    for itype in types:
        if len(itype) > max_atomtype_length:
            warnings.warn("Warning, type name {} will be shortened to {} "
                      "characters as {}. Please confirm your final "
                      "MCF.".format(
                          itype,
                          max_atomtype_length,
                          itype[-max_atomtype_length:]))
        types = [ itype[-max_atomtype_length:] for itype in types ]
    if len(set(types)) < n_unique_types:
        warnings.warn("Warning, the number of unique atomtypes has been "
              "reduced due to shortening the atomtype name to {} "
              "characters.".format(max_atomtype_length))

    vdw_type = 'LJ'
    header = ('!Atom Format\n'
              '!index type element mass charge vdw_type parameters\n'
              '!vdw_type="LJ", parms=epsilon sigma\n'
              '!vdw_type="Mie", parms=epsilon sigma '
              'repulsion_exponent dispersion_exponent\n'
              '\n# Atom_Info\n'
              )

    mcf_file.write(header)
    mcf_file.write('{:d}\n'.format(len(structure.atoms)))
    for i in range(len(structure.atoms)):
        mcf_file.write('{:<4d}  {:<6s}  {:<2s}  {:7.3f}  {:12.8f}  '
                '{:3s}  {:8.3f}  {:8.3f}'.format(
            i+1, types[i], elements[i], masses[i], charges[i],
            vdw_type, epsilons[i], sigmas[i]))
        if in_ring[i] == True:
            mcf_file.write('  ring')
        mcf_file.write('\n')

def _write_bond_information(mcf_file, structure):
    """Write the bonds in the system.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object

    """

    bond_parms = [ str('{:8.3f}'.format(bond.type.req))
            for bond in structure.bonds ]

    mcf_file.write('\n!Bond Format\n')
    mcf_file.write('!index i j type parameters\n' +
              '!type="fixed", parms=bondLength\n')
    mcf_file.write('\n# Bond_Info\n')
    mcf_file.write('{:d}\n'.format(len(structure.bonds)))
    for i, bond in enumerate(structure.bonds):
        mcf_file.write('{:<4d}  {:<4d}  {:<4d}  {:s}  {:s}\n'.format(
            i+1, bond.atom1.idx+1, bond.atom2.idx+1, 'fixed', bond_parms[i]))

def _write_angle_information(mcf_file, structure, angle_style,
        IG_CONSTANT_KCAL):
    """Write the angles in the system.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object
    angle_style : string
        Angle style for Cassandra to use
    IG_CONSTANT_KCAL : float
        Ideal gas constant in kcal/mol K

    """

    if angle_style.casefold() == 'fixed':
        angle_parms = [ str('{:8.2f}'.format(angle.type.theteq))
                        for angle in structure.angles ]
    elif angle_style.casefold() == 'harmonic':
        # Convert energies to units of K
        angle_parms = [ str('{:8.1f}'.format(angle.type.k/IG_CONSTANT_KCAL)) +
                        '  ' + str('{:8.2f}'.format(angle.type.theteq))
                        for angle in structure.angles ]
    else:
        raise ValueError("Only 'fixed' and 'harmonic' angle styles "
                         "are supported by Cassandra")

    header = ( '\n!Angle Format\n'
               '!index i j k type parameters\n'
               '!type="fixed", parms=equilibrium_angle\n'
               '!type="harmonic", parms=force_constant equilibrium_angle\n'
               '\n# Angle_Info\n'
             )

    mcf_file.write(header)
    mcf_file.write('{:d}\n'.format(len(structure.angles)))

    for i, angle in enumerate(structure.angles):
        mcf_file.write('{:<4d}  {:<4d}  {:<4d}  {:<4d}  {:s}  {:s}\n'.format(
            i+1, angle.atom1.idx+1, angle.atom2.idx+1, angle.atom3.idx+1,
            angle_style.lower(), angle_parms[i]))

def _write_dihedral_information(mcf_file, structure, dihedral_style,
        KCAL_TO_KJ):
    """Write the dihedrals in the system.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object
    dihedral_style : string
        Dihedral style for Cassandra to use
    KCAL_TO_KJ : float
        Conversion factor from kcal to kJ
    """

    # Dihedral info
    header = ( '\n!Dihedral Format\n'
               '!index i j k l type parameters\n'
               '!type="none"\n'
               '!type="CHARMM", parms=a0 a1 delta\n'
               '!type="OPLS", parms=c0 c1 c2 c3\n'
               '!type="harmonic", parms=force_constant equilibrium_dihedral\n'
               '\n# Dihedral_Info\n'
             )

    mcf_file.write(header)

    if len(structure.dihedrals) > 0 or len(structure.rb_torsions) > 0:
        if dihedral_style.casefold() == 'opls':
            dihedral_style = dihedral_style.upper()
            dihedrals = structure.rb_torsions
            # Two things happen here:
            #  (1) convert from RB form to Cassandra OPLS
            #  (2) convert units from kcal/mol to kJ/mol
            dihedral_parms = []
            for dihedral in dihedrals:
                a0 = ( dihedral.type.c0 + dihedral.type.c1 +
                        dihedral.type.c2 + dihedral.type.c3 )
                a1 = -dihedral.type.c1 - (3./4.)*dihedral.type.c3
                a2 = (-1./2.)*dihedral.type.c2
                a3 = (-1./4.)*dihedral.type.c3
                if not dihedral.type.c4 == 0. and dihedral.type.c4 == 0.:
                    raise ValueError("Can only convert Ryckaert-Bellemans "
                                     "dihedrals to OPLS if c4==0 and c5==0")
                dihedral_parms.append(str('{:8.3f} '.format(a0*KCAL_TO_KJ)) +
                                      str('{:8.3f} '.format(a1*KCAL_TO_KJ)) +
                                      str('{:8.3f} '.format(a2*KCAL_TO_KJ)) +
                                      str('{:8.3f} '.format(a3*KCAL_TO_KJ)))

        elif dihedral_style.casefold() == 'charmm':
            dihedral_style = dihedral_style.upper()
            dihedrals = structure.dihedrals
            # type.per = periodicity (a1)
            # type.phase = phase offset (delta)
            dihedral_parms = [ str('{:8.3f} '.format(
                                        dihedral.type.phi_k*KCAL_TO_KJ)) +
                               str('{:8.3f} '.format(dihedral.type.per)) +
                               str('{:8.3f}'.format(dihedral.type.phase))
                               for dihedral in dihedrals ]

        elif dihedral_style.casefold() == 'none':
            warnings.warn("Dihedral style 'none' selected. "
                          "Ignoring dihedral parameters")
            dihedral_style = dihedral_style.lower()
            if structure.dihedrals:
                dihedrals = structure.dihedrals
                dihedral_parms = [ '' for dihedral in dihedrals ]
            elif structure.rb_torsions:
                dihedrals = structure.rb_torsions
                dihedral_parms = [ '' for dihedral in dihedrals ]
        else:
            raise ValueError("Only 'OPLS', 'CHARMM', and 'none' "
                             "dihedral styles are supported.")

        mcf_file.write('{:d}\n'.format(len(dihedrals)))
        for i,dihedral in enumerate(dihedrals):
            # Sorting here to match LEAP behavior
            # See https://github.com/choderalab/openmoltools/issues/24
            # If atom types are identical, too bad.
            if dihedral.improper:
                improper_atoms = [dihedral.atom1,
                                  dihedral.atom2,
                                  dihedral.atom4]
                improper_atoms.sort(key=lambda x: x.type)
                atom1 = improper_atoms[0]
                atom2 = improper_atoms[1]
                atom3 = dihedral.atom3
                atom4 = improper_atoms[2]
            else:
                atom1 = dihedral.atom1
                atom2 = dihedral.atom2
                atom3 = dihedral.atom3
                atom4 = dihedral.atom4
            mcf_file.write('{:<4d}  {:<4d}  {:<4d}  {:<4d}  {:<4d}'
                           '  {:s}  {:s}\n'.format(
                i+1, atom1.idx+1, atom2.idx+1, atom3.idx+1, atom4.idx+1,
                dihedral_style, dihedral_parms[i]))
    else:
        mcf_file.write('0\n')

def _write_improper_information(mcf_file, structure, KCAL_TO_KJ):
    """Write the impropers in the system.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object
    KCAL_TO_KJ : float
        Conversion factor from kcal to kJ
    """

    header = ( '\n!Improper Format\n'
               '!index i j k l type parameters\n'
               '!type="harmonic", parms=force_constant equilibrium_improper\n'
               '\n# Improper_Info\n'
             )

    mcf_file.write(header)
    mcf_file.write('{:d}\n'.format(len(structure.impropers)))

    improper_type = 'harmonic'
    for i, improper in enumerate(structure.impropers):
        mcf_file.write('{:<4d}  {:<4d}  {:<4d}  {:<4d}  {:<4d}'
                       '  {:s}  {:8.3f}  {:8.3f}\n'.format(
            i+1, improper.atom1.idx+1, improper.atom2.idx+1,
            improper.atom3.idx+1, improper.atom4.idx+1,
            improper_type, improper.type.psi_k*KCAL_TO_KJ,
            improper.type.psi_eq))

def _write_fragment_information(mcf_file, structure, frag_list, frag_conn):
    """Write the fragments in the molecule.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    structure : parmed.Structure
        Parmed structure object
    frag_list : list
        Atom ids belonging to each fragment
    frag_conn : list
        Fragment ids of connected fragments

    """

    header = ('\n!Fragment Format\n'
              '!index number_of_atoms_in_fragment branch_point other_atoms\n'
              '\n# Fragment_Info\n'
             )

    mcf_file.write(header)

    # Special cases first
    if len(frag_list) == 0:
        if len(structure.atoms) == 1:
            mcf_file.write('1\n')
            mcf_file.write('1 1 1\n')
        elif len(structure.atoms) == 2:
            mcf_file.write('1\n')
            mcf_file.write('1 2 1 2\n')
        else:
            warnings.warn('More than two atoms present but '
                          'no fragments identified.')
            mcf_file.write('0\n')
    else:
        mcf_file.write('{:d}\n'.format(len(frag_list)))
        for i, frag in enumerate(frag_list):
            mcf_file.write('{:d}    {:d}'.format(i+1, len(frag)))
            for idx in frag:
                mcf_file.write('    {:d}'.format(idx+1))
            mcf_file.write('\n')

    mcf_file.write('\n\n# Fragment_Connectivity\n')
    mcf_file.write('{:d}\n'.format(len(frag_conn)))
    for i, conn in enumerate(frag_conn):
        mcf_file.write('{:d}    {:d}    {:d}\n'.format(i+1, conn[0]+1, conn[1]+1))

def _write_intrascaling_information(mcf_file, lj14, coul14):
    """Write the intramolecular scaling in the molecule.

    Parameters
    ----------
    mcf_file : file object
        The file object of the Cassandra mcf being written
    lj14 : float
        The 1-4 scaling parameter for LJ interactions
    coul14 : float
        The 1-4 scaling parameter for Coulombic interactions

    """

    header = ( '\n!Intra Scaling\n'
               '!vdw_scaling    1-2 1-3 1-4 1-N\n'
               '!charge_scaling 1-2 1-3 1-4 1-N\n'
               '\n# Intra_Scaling\n'
             )

    mcf_file.write(header)
    mcf_file.write('0. 0. {:.4f} 1.\n'.format(lj14))
    mcf_file.write('0. 0. {:.4f} 1.\n'.format(coul14))


import torch
import torch.utils.data
from torch import nn, optim
from torch.autograd import Variable

from dataloader import RNA_Dataset
from dataloader import NucleiDatasetNew as NucleiDataset
from model import FC_Autoencoder, FC_Classifier, VAE, FC_VAE, Simple_Classifier

import os
import argparse
import numpy as np
import imageio

torch.manual_seed(1)

#============ PARSE ARGUMENTS =============

def setup_args():

    options = argparse.ArgumentParser()

    # save and directory options
    options.add_argument('-sd', '--save-dir', action="store", dest="save_dir", default="/Users/esthomas/Andor_Rotation/github_repo/cross-modal-autoencoders/save_dir")
    options.add_argument('--save-freq', action="store", dest="save_freq", default=20, type=int)
    options.add_argument('--pretrained-file', action="store", default="/Users/esthomas/Andor_Rotation/github_repo/cross-modal-autoencoders/save_dir/models/0.pth")

    # training parameters
    options.add_argument('-bs', '--batch-size', action="store", dest="batch_size", default=32, type=int)
    options.add_argument('-w', '--num-workers', action="store", dest="num_workers", default=10, type=int)
    options.add_argument('-lrAE', '--learning-rate-AE', action="store", dest="learning_rate_AE", default=1e-4, type=float)
    options.add_argument('-lrD', '--learning-rate-D', action="store", dest="learning_rate_D", default=1e-4, type=float)
    options.add_argument('-e', '--max-epochs', action="store", dest="max_epochs", default=2, type=int)
    options.add_argument('-wd', '--weight-decay', action="store", dest="weight_decay", default=0, type=float)
    options.add_argument('--train-imagenet', action="store_true")
    options.add_argument('--conditional', action="store_true")
    options.add_argument('--conditional-adv', action="store_true")

    # hyperparameters
    options.add_argument('--alpha', action="store", default=0.1, type=float)
    options.add_argument('--beta', action="store", default=1., type=float)
    options.add_argument('--lamb', action="store", default=0.00000001, type=float)
    options.add_argument('--latent-dims', action="store", default=128, type=int)

    # gpu options
    options.add_argument('-gpu', '--use-gpu', action="store_false", dest="use_gpu")

    return options.parse_args()

args = setup_args()
if not torch.cuda.is_available():
    args.use_gpu = False

os.makedirs(args.save_dir, exist_ok=True)

#============= TRAINING INITIALIZATION ==============

# initialize autoencoder
netRNA = FC_VAE(n_input=7633, nz=args.latent_dims)

netImage = VAE(latent_variable_size=args.latent_dims, batchnorm=True)
netImage.load_state_dict(torch.load(args.pretrained_file))
print("Pre-trained model loaded from %s" % args.pretrained_file)

if args.conditional_adv:
    netClf = FC_Classifier(nz=args.latent_dims+10)
    assert(not args.conditional)
else:
    netClf = FC_Classifier(nz=args.latent_dims)

if args.conditional:
    netCondClf = Simple_Classifier(nz=args.latent_dims)

if args.use_gpu:
    netRNA.cuda()
    netImage.cuda()
    netClf.cuda()
    if args.conditional:
        netCondClf.cuda()

# load data
genomics_dataset = RNA_Dataset(datadir="data_folder/data/nCD4_gene_exp_matrices/")
image_dataset = NucleiDataset(datadir="data_folder/data/nuclear_crops_all_experiments", mode='test')

image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=args.batch_size, drop_last=True, shuffle=True)
genomics_loader = torch.utils.data.DataLoader(genomics_dataset, batch_size=args.batch_size, drop_last=True, shuffle=True)

# setup optimizer
opt_netRNA = optim.Adam(list(netRNA.parameters()), lr=args.learning_rate_AE)
opt_netClf = optim.Adam(list(netClf.parameters()), lr=args.learning_rate_D, weight_decay=args.weight_decay)
opt_netImage = optim.Adam(list(netImage.parameters()), lr=args.learning_rate_AE)

if args.conditional:
    opt_netCondClf = optim.Adam(list(netCondClf.parameters()), lr=args.learning_rate_AE)

# loss criteria
criterion_reconstruct = nn.MSELoss()
criterion_classify = nn.CrossEntropyLoss()

# setup logger
with open(os.path.join(args.save_dir, 'log.txt'), 'w') as f:
    print(args, file=f)
    print(netRNA, file=f)
    print(netImage, file=f)
    print(netClf, file=f)
    if args.conditional:
        print(netCondClf, file=f)

# define helper train functions

def compute_KL_loss(mu, logvar):
    if args.lamb>0:
        KLloss = -0.5*torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return args.lamb * KLloss
    return 0

def train_autoencoders(rna_inputs, image_inputs, rna_class_labels=None, image_class_labels=None):

    netRNA.train()
    if args.train_imagenet:
        netImage.train()
    else:
        netImage.eval()
    netClf.eval()
    if args.conditional:
        netCondClf.train()

    # process input data
    rna_inputs, image_inputs = Variable(rna_inputs), Variable(image_inputs)

    if args.use_gpu:
        rna_inputs, image_inputs = rna_inputs.cuda(), image_inputs.cuda()

    # reset parameter gradients
    netRNA.zero_grad()

    # forward pass
    rna_recon, rna_latents, rna_mu, rna_logvar = netRNA(rna_inputs)
    image_recon, image_latents, image_mu, image_logvar = netImage(image_inputs)

    if args.conditional_adv:
        rna_class_labels, image_class_labels = rna_class_labels.cuda(), image_class_labels.cuda()
        rna_scores = netClf(torch.cat((rna_latents, rna_class_labels.float().view(-1,1).expand(-1,10)), dim=1))
        image_scores = netClf(torch.cat((image_latents, image_class_labels.float().view(-1,1).expand(-1,10)), dim=1))
    else:
        rna_scores = netClf(rna_latents)
        image_scores = netClf(image_latents)

    rna_labels = torch.zeros(rna_scores.size(0),).long()
    image_labels = torch.ones(image_scores.size(0),).long()

    if args.conditional:
        rna_class_scores = netCondClf(rna_latents)
        image_class_scores = netCondClf(image_latents)

    if args.use_gpu:
        rna_labels, image_labels = rna_labels.cuda(), image_labels.cuda()
        if args.conditional:
            rna_class_labels, image_class_labels = rna_class_labels.cuda(), image_class_labels.cuda()

    # compute losses
    rna_recon_loss = criterion_reconstruct(rna_inputs, rna_recon)
    image_recon_loss = criterion_reconstruct(image_inputs, image_recon)
    kl_loss = compute_KL_loss(rna_mu, rna_logvar) + compute_KL_loss(image_mu, image_logvar)
    clf_loss = 0.5*criterion_classify(rna_scores, image_labels) + 0.5*criterion_classify(image_scores, rna_labels)
    loss = args.alpha*(rna_recon_loss + image_recon_loss) + clf_loss + kl_loss

    if args.conditional:
        clf_class_loss = 0.5*criterion_classify(rna_class_scores, rna_class_labels) + 0.5*criterion_classify(image_class_scores, image_class_labels)
        loss += args.beta*clf_class_loss

    # backpropagate and update model
    loss.backward()
    opt_netRNA.step()
    if args.conditional:
        opt_netCondClf.step()

    if args.train_imagenet:
        opt_netImage.step()

    summary_stats = {'rna_recon_loss': rna_recon_loss.item()*rna_scores.size(0), 'image_recon_loss': image_recon_loss.item()*image_scores.size(0),
            'clf_loss': clf_loss.item()*(rna_scores.size(0)+image_scores.size(0))}

    if args.conditional:
        summary_stats['clf_class_loss'] = clf_class_loss.item()*(rna_scores.size(0)+image_scores.size(0))

    return summary_stats

def train_classifier(rna_inputs, image_inputs, rna_class_labels=None, image_class_labels=None):

    netRNA.eval()
    netImage.eval()
    netClf.train()

    # process input data
    rna_inputs, image_inputs = Variable(rna_inputs), Variable(image_inputs)

    if args.use_gpu:
        rna_inputs, image_inputs = rna_inputs.cuda(), image_inputs.cuda()

    # reset parameter gradients
    netClf.zero_grad()

    # forward pass
    _, rna_latents, _, _ = netRNA(rna_inputs)
    _, image_latents, _, _ = netImage(image_inputs)

    if args.conditional_adv:
        rna_class_labels, image_class_labels = rna_class_labels.cuda(), image_class_labels.cuda()
        rna_scores = netClf(torch.cat((rna_latents, rna_class_labels.float().view(-1,1).expand(-1,10)), dim=1))
        image_scores = netClf(torch.cat((image_latents, image_class_labels.float().view(-1,1).expand(-1,10)), dim=1))
    else:
        rna_scores = netClf(rna_latents)
        image_scores = netClf(image_latents)

    rna_labels = torch.zeros(rna_scores.size(0),).long()
    image_labels = torch.ones(image_scores.size(0),).long()

    if args.use_gpu:
        rna_labels, image_labels = rna_labels.cuda(), image_labels.cuda()

    # compute losses
    clf_loss = 0.5*criterion_classify(rna_scores, rna_labels) + 0.5*criterion_classify(image_scores, image_labels)

    loss = clf_loss

    # backpropagate and update model
    loss.backward()
    opt_netClf.step()

    summary_stats = {'clf_loss': clf_loss*(rna_scores.size(0)+image_scores.size(0)), 'rna_accuracy': accuracy(rna_scores, rna_labels), 'rna_n_samples': rna_scores.size(0),
            'image_accuracy': accuracy(image_scores, image_labels), 'image_n_samples': image_scores.size(0)}

    return summary_stats

def accuracy(output, target):
    pred = output.argmax(dim=1).view(-1)
    correct = pred.eq(target.view(-1)).float().sum().item()
    return correct

def generate_image(epoch):
    img_dir = os.path.join(args.save_dir, "images")
    os.makedirs(img_dir, exist_ok=True)
    netRNA.eval()
    netImage.eval()

    for i in range(5):
        samples = genomics_loader.dataset[np.random.randint(30)]
        rna_inputs = samples['tensor']
        rna_inputs = Variable(rna_inputs.unsqueeze(0))
        samples = image_loader.dataset[np.random.randint(30)]
        image_inputs = samples['image_tensor']
        image_inputs = Variable(image_inputs.unsqueeze(0))

        if torch.cuda.is_available():
            rna_inputs = rna_inputs.cuda()
            image_inputs = image_inputs.cuda()

        _, rna_latents, _, _ = netRNA(rna_inputs)
        recon_inputs = netImage.decode(rna_latents)
        imageio.imwrite(os.path.join(img_dir, "epoch_%s_trans_%s.jpg" % (epoch, i)), np.uint8(recon_inputs.cpu().data.view(64,64).numpy()*255))
        recon_images, _, _, _ = netImage(image_inputs)
        imageio.imwrite(os.path.join(img_dir, "epoch_%s_recon_%s.jpg" % (epoch, i)), np.uint8(recon_images.cpu().data.view(64,64).numpy()*255))

### main training loop
for epoch in range(args.max_epochs):
    print(epoch)

    if epoch % args.save_freq == 0:
        generate_image(epoch)

    recon_rna_loss = 0
    recon_image_loss = 0
    clf_loss = 0
    clf_class_loss = 0
    AE_clf_loss = 0

    n_rna_correct = 0
    n_rna_total = 0
    n_atac_correct = 0
    n_atac_total = 0

    for idx, (rna_samples, image_samples) in enumerate(zip(genomics_loader, image_loader)):
        rna_inputs = rna_samples['tensor']
        image_inputs = image_samples['image_tensor']

        if args.conditional or args.conditional_adv:
            rna_labels = rna_samples['binary_label']
            image_labels = image_samples['binary_label']
            out = train_autoencoders(rna_inputs, image_inputs, rna_labels, image_labels)
        else:
            out = train_autoencoders(rna_inputs, image_inputs)

        recon_rna_loss += out['rna_recon_loss']
        recon_image_loss += out['image_recon_loss']
        AE_clf_loss += out['clf_loss']

        if args.conditional:
            clf_class_loss += out['clf_class_loss']

        if args.conditional_adv:
            out = train_classifier(rna_inputs, image_inputs, rna_labels, image_labels)
        else:
            out = train_classifier(rna_inputs, image_inputs)

        clf_loss += out['clf_loss']
        n_rna_correct += out['rna_accuracy']
        n_rna_total += out['rna_n_samples']
        n_atac_correct += out['image_accuracy']
        n_atac_total += out['image_n_samples']

    recon_rna_loss /= n_rna_total
    clf_loss /= n_rna_total+n_atac_total
    AE_clf_loss /= n_rna_total+n_atac_total

    if args.conditional:
        clf_class_loss /= n_rna_total + n_atac_total

    with open(os.path.join(args.save_dir, 'log.txt'), 'a') as f:
        print('Epoch: ', epoch, ', rna recon loss: %.8f' % float(recon_rna_loss), ', image recon loss: %.8f' % float(recon_image_loss),
                ', AE clf loss: %.8f' % float(AE_clf_loss), ', clf loss: %.8f' % float(clf_loss), ', clf class loss: %.8f' % float(clf_class_loss),
                ', clf accuracy RNA: %.4f' % float(n_rna_correct / n_rna_total), ', clf accuracy ATAC: %.4f' % float(n_atac_correct / n_atac_total), file=f)

    # save model
    if epoch % args.save_freq == 0:
        torch.save(netRNA.cpu().state_dict(), os.path.join(args.save_dir,"netRNA_%s.pth" % epoch))
        torch.save(netImage.cpu().state_dict(), os.path.join(args.save_dir,"netImage_%s.pth" % epoch))
        torch.save(netClf.cpu().state_dict(), os.path.join(args.save_dir,"netClf_%s.pth" % epoch))
        if args.conditional:
            torch.save(netCondClf.cpu().state_dict(), os.path.join(args.save_dir,"netCondClf_%s.pth" % epoch))

    if args.use_gpu:
        netRNA.cuda()
        netClf.cuda()
        netImage.cuda()
        if args.conditional:
            netCondClf.cuda()


"""
Utility scripts to generate YOLO training data
"""
import os
from os.path import join
import sys
import json
import numpy as np
import random
from PIL import Image
from pprint import pprint
import matplotlib.pyplot as plt
plt.style.use('ggplot')
import itertools
import shutil
from pathlib import Path
from utils import fiox
from training.vcat.vcat_utils import VcatUtils

class YoloUtils:

  def __init__(self):
    pass

  def generate_image_list(self,anno_objs,dir_project_images):
    # for each object-class
    ims = []
    for class_idx, anno_obj in anno_objs.items():
      for regions in anno_obj['regions']:
        ims.append('{}/{}{}'.format(dir_project_images,regions['fn'],regions['ext']))
    return ims

  def generate_yolo_config(self,yolo_version,n_classes,cfg_type,n_subdiv=16,n_batch=64):

    # load existing cfg file

    # (Generally filters depends on the classes, num and coords
    # i.e. equal to (classes + coords + 1)*num, where num is number of anchors)
    # for yolov2 (num_classes + (num_coords +1)) * num_anchors

    # load the template from github.com/pjreddie/darknet
    cwd = os.path.dirname(os.path.realpath(__file__))
    fp_cfg_orig = os.path.join(cwd,'cfgs','yolov{}.cfg'.format(yolo_version))

    if yolo_version == '2':

      num_filters = (n_classes + 5) * 5
      s_replace_base = [
        ('filters=425', 'filters={}'.format(num_filters)),
        ('classes=80', 'classes={}'.format(n_classes)),
        ]

      if cfg_type == 'train':
        s_replace = s_replace_base.copy()
        s_replace.append(('batch=1','#batch=1'))
        s_replace.append(('subdivisions=1\n','#subdivisions=1\n'))
        s_replace.append(('# batch=64','batch={}'.format(n_batch)))
        s_replace.append(('# subdivisions=8','subdivisions={}'.format(n_subdiv)))
      elif cfg_type == 'test':
        # change nclasses, but keep the subdiv and batch size at default
        s_replace = s_replace_base.copy()

    elif yolo_version == '3':
      # load the template from github.com/pjreddie/darknet
      num_filters = (n_classes + 5) * 3
      s_replace_base = [
        ('filters=255', 'filters={}'.format(num_filters)),
        ('classes=80', 'classes={}'.format(n_classes)),
        ]

      if cfg_type == 'train':
        s_replace = s_replace_base.copy()
        s_replace.append(('batch=1','#batch=1'))
        s_replace.append(('subdivisions=1\n','#subdivisions=1\n'))
        s_replace.append(('# batch=64','batch={}'.format(n_batch)))
        s_replace.append(('# subdivisions=16','subdivisions={}'.format(n_subdiv)))
      elif cfg_type == 'test':
        # change nclasses, but keep the subdiv and batch size at default
        s_replace = s_replace_base.copy()

    # search and replace, return list of lines
    cfg_orig = open(fp_cfg_orig,'r')
    cfg_new_data = []
    for line in cfg_orig:
      for query,targ in s_replace:
        if query in line:
          line = line.replace(query,targ)
      cfg_new_data.append(line.replace('\n',''))

    return cfg_new_data

  def create_label_anno_index(self,image_anno_index):
    """Create index of labels with Darknet label format
      One text file per image with one line for each annotation
    """
    label_anno_index = {}

    for fn_key,image_obj in image_anno_index.items():
      bname = os.path.splitext(image_obj['basename'])[0]
      fname = os.path.splitext(bname)[0]
      im_anno_txt = []
      for region_obj in image_obj['regions']:
        w = float(region_obj['width'])
        h = float(region_obj['height'])
        cx = (w/2.) + float(region_obj['x'])
        cy = (h/2.) + float(region_obj['y'])
        class_label_index = region_obj['class_label_index']
        txt = '{} {} {} {} {}'.format(class_label_index,cx,cy,w,h)
        im_anno_txt.append(txt)
        if region_obj['is_attribute']:
          # get parent id-->class_label_index
          parent_class_label_index = region_obj['parent_class_label_index']
          txt = '{} {} {} {} {}'.format(parent_class_label_index,cx,cy,w,h)
          im_anno_txt.append(txt)

      label_anno_index[fn_key] = im_anno_txt

    return label_anno_index

  def create_image_symlinks(self,image_anno_index,dir_project_images):
    """Create symlinks to image files so project can be
    contained in single directory"""
    for fn_key,image_obj in image_anno_index.items():
        src = image_obj['filepath'] # path to local file
        dst = os.path.join(dir_project_images,'{}{}'.format(fn_key,image_obj['ext']))
        # remove symlink if exists
        if os.path.isfile(dst):
            os.remove(dst)
        # add symlink
        os.symlink(src,dst)

  def get_random_image(self,annos):
    """Return a random image for testing"""
    # get random class image
    annos_with_regions = [a for k,a in annos.items() if len(a['regions']) >1]
    v = random.choice(annos_with_regions)
    region_obj = random.choice(v['regions'])
    return region_obj['filepath']

  def create_project(self,kwargs):

    vcat_utils = VcatUtils()

    # load vcat JSON
    vcat_data = json.load(kwargs['input'])

    # hierarchy is
    hierarchy = vcat_data['hierarchy']

    # if exclusions, remove classes
    if kwargs['exclude'] is not None:
        exclusions = [int(i) for i in kwargs['exclude'].split(',')]
        print('[+] Excluding: {}'.format(exclusions))
        # remove from hierarchy
        hierarchy_tmp = hierarchy.copy()
        for obj_class_id, obj_class_meta in hierarchy_tmp.items():
          if int(obj_class_id) in exclusions:
            print('[+] Removing ID: {} ({})'.format(obj_class_id,obj_class_meta['slug']))
            del hierarchy[obj_class_id]

    # class label index lookup (YOLO  ID --> VCAT ID)
    class_label_index_lkup = vcat_utils.create_class_label_index(vcat_data,kwargs)

    print('----Classes----')
    for k,v in class_label_index_lkup.items():
      print('VCAT: {} --> YOLO: {}, Label: {}'.format(k,v,hierarchy[str(k)]['slug']))

    # create object class lookup
    anno_index = vcat_utils.create_annotation_index(
      vcat_data, class_label_index_lkup, kwargs)

    for a in anno_index:
      print('Annotation index: {}'.format(a))

    sys.exit()

    # inflate annotation index to include parent class (ghost) annotations
    # TODO

    # create image index lookup
    image_anno_index = vcat_utils.create_image_anno_index(anno_index)

    # Create randomize, class-based splits for train and valid
    anno_splits = vcat_utils.create_splits(anno_index,split_val=kwargs['split_val'])
    annos_train = anno_splits['train']
    annos_valid = anno_splits['valid']

    # Statistics
    n_regions_train = np.sum([len(v['regions']) for k,v in annos_train.items()])
    n_regions_test = np.sum([len(v['regions']) for k,v in annos_valid.items()])
    print('[+] Regions train: {}, Test: {}'.format(n_regions_train, n_regions_test))

    print('----Classes----')
    for k,v in class_label_index_lkup.items():
      print('VCAT: {} --> YOLO: {}, Label: {}'.format(k,v,hierarchy[str(k)]['slug']))

    print('----Train----')
    for k,v in annos_train.items():
      print('{}\t{}'.format(len(v['regions']),v['slug_vcat']))

    print('----Validate----')
    for k,v in annos_valid.items():
      print('{}\t{}'.format(len(v['regions']),v['slug_vcat']))

    print('----Labels----')
    for k,v in class_label_index_lkup.items():
      print('{}\t{}'.format(k,hierarchy[str(k)]['slug']))

    sys.exit()

    # ---------------------------------------------------
    # Filenames and paths
    # ---------------------------------------------------

    # convenience vars
    n_classes = len(anno_index)
    yolov = str(kwargs['yolo_version'])
    project_name = kwargs['name']
    dir_output = kwargs['dir_output']
    dir_project = os.path.join(dir_output,project_name)

    # create project_name directory
    fiox.ensure_dir(dir_project)

    # create images directory
    dir_project_images = os.path.join(dir_project,'images')
    fiox.ensure_dir(dir_project_images)

    # Create dirs if not exist
    dir_labels = os.path.join(dir_project,'labels')
    fiox.ensure_dir(dir_labels)

    # create dir for weights/backup
    dir_weights = os.path.join(dir_project,'weights')
    fiox.ensure_dir(dir_weights)

    # ---------------------------------------------------
    # Create "classes.txt"
    # ---------------------------------------------------

    fp_classes = os.path.join(dir_project,'classes.txt')
    class_labels = []
    for k,v in class_label_index_lkup.items():
      # maps vcat id to yolo id
      slug = hierarchy[str(k)]['slug']
      slug_sp = slug.split(':')
      if len(slug_sp) > 1:
        display_name = '{} ({})'.format(slug_sp[0],slug_sp[1])
      else:
        display_name = slug_sp[0]

      class_labels.append(display_name)

    with open(fp_classes,'w') as fp:
      fp.write('\n'.join(class_labels))

    # ---------------------------------------------------
    # Create "train.txt", "valid.txt", "classes.txt"
    # ---------------------------------------------------

    fp_train = os.path.join(dir_project,'train.txt')
    im_list = self.generate_image_list(annos_train,dir_project_images)
    with open(fp_train,'w') as fp:
      fp.write('\n'.join(im_list))

    fp_valid = os.path.join(dir_project,'valid.txt')
    im_list = self.generate_image_list(annos_valid,dir_project_images)
    with open(fp_valid,'w') as fp:
      fp.write('\n'.join(im_list))

    # ---------------------------------------------------
    # Crate sym-linked image files
    # ---------------------------------------------------

    # TODO remove folder if exists

    self.create_image_symlinks(image_anno_index,dir_project_images)

    # ---------------------------------------------------
    # Create "labels.txt"
    # ---------------------------------------------------

    # labels.txt contains one line per annotation
    # "1 0.434 0.605 0.2980 0.37222"
    # class_idx, cx, cy, w, h
    # for YoloV3 add parent class hierarchy to the label

    label_anno_index = self.create_label_anno_index(image_anno_index)

    for fname, anno_obj in label_anno_index.items():
      fp_txt = os.path.join(dir_labels,'{}.txt'.format(fname))
      with open(fp_txt,'w') as fp:
        fp.write('\n'.join(anno_obj))

    # ---------------------------------------------------
    # Create "meta.data"
    # ---------------------------------------------------

    txts = []
    txts.append('classes = {}'.format(n_classes))
    txts.append('train = {}'.format(fp_train))
    txts.append('valid = {}'.format(fp_valid))
    txts.append('names = {}'.format(fp_classes))
    txts.append('backup = {}'.format(dir_weights))

    fp_meta = os.path.join(dir_project,'meta.data')
    with open(fp_meta,'w') as fp:
      fp.write('\n'.join(txts))

    # ---------------------------------------------------
    # create train and test .cfg files
    # ---------------------------------------------------

    cfg_train_data = self.generate_yolo_config(yolov,n_classes,'train',
      n_subdiv=kwargs['subdivisions'],n_batch=kwargs['batch_size'])
    cfg_test_data = self.generate_yolo_config(yolov,n_classes,'test')

    # write data to new .cfg files
    fp_cfg_train = os.path.join(dir_project,'yolov{}.cfg'.format(yolov))
    fp_cfg_test = os.path.join(dir_project,'yolov{}_test.cfg'.format(yolov))

    with open(fp_cfg_train, 'w') as fp:
      fp.write('\n'.join(cfg_train_data))

    with open(fp_cfg_test, 'w') as fp:
      fp.write('\n'.join(cfg_test_data))

    # ---------------------------------------------------
    # Create shell scripts to start and resume training
    # ---------------------------------------------------

    txts_base = []
    txts_base.append('#!/bin/bash')
    txts_base.append('DARKNET=/opt/darknet_pjreddie/darknet')
    txts_base.append('DIR_PROJECT={}'.format(dir_project))
    txts_base.append('FP_CFG={}'.format(fp_cfg_train))
    txts_base.append('FP_META={}'.format(fp_meta))
    txts_base.append('CMD="detector train"')
    txts_base.append('LOGFILE="2>&1 | tee train_log.txt"')

    # run_train_init.sh
    fp_sh_train_init = os.path.join(dir_project,'run_train_init.sh')

    txts_init = txts_base.copy()
    txts_init.append('GPUS="-gpus {}"'.format(kwargs['gpu_init']))
    txts_init.append('FP_WEIGHTS={}'.format(kwargs['init_weights']))
    txts_init.append('$DARKNET $CMD $FP_META $FP_CFG $FP_WEIGHTS $GPUS')

    with open(fp_sh_train_init,'w') as fp:
      fp.write('\n'.join(txts_init))

    # run_train_resume.sh
    fp_sh_train_resume = os.path.join(dir_project,'run_train_resume.sh')

    txts_resume = txts_base.copy()
    txts_resume.append('GPUS="-gpus {}"'.format(kwargs['gpu_full']))
    txts_resume.append('FP_WEIGHTS={}/{}.backup'.format(dir_weights,project_name))
    txts_resume.append('$DARKNET $CMD $FP_META $FP_CFG $FP_WEIGHTS $GPUS')

    with open(fp_sh_train_resume, 'w') as fp:
      fp.write('\n'.join(txts_resume))

    # run_test.sh
    fp_sh_test_image = os.path.join(dir_project,'run_test_image.sh')

    txts = []
    txts.append('#!/bin/bash')
    txts.append('DIR_DARKNET=/opt/darknet_pjreddie/')
    txts.append('DIR_PROJECT={}'.format(dir_project))
    txts.append('FP_CFG={}'.format(fp_cfg_test))
    txts.append('FP_META={}'.format(fp_meta))
    rn_filepath = self.get_random_image(annos_valid)
    txts.append('FP_IMG={}'.format(rn_filepath))
    txts.append('FP_WEIGHTS={}/{}_900.weights'.format(dir_weights,project_name))
    txts.append('#FP_WEIGHTS={}/{}_10000.weights'.format(dir_weights,project_name))
    txts.append('#FP_WEIGHTS={}/{}_30000.weights'.format(dir_weights,project_name))
    txts.append('CMD="detector test"')
    txts.append('cd $DIR_DARKNET')
    txts.append('./darknet $CMD $FP_META $FP_CFG $FP_WEIGHTS $FP_IMG')

    with open(fp_sh_test_image, 'w') as fp:
      fp.write('\n'.join(txts))

    print('[+] Wrote Darknet files for {}'.format(project_name))
    print('[+] Project path: {}'.format(dir_project))

  # ------------------------------------------
  # Plot YOLO Loss Data
  # ------------------------------------------

  def create_plot(self,kwargs):
    """Create plot of Yolo Loss Avg"""
    # can stop training when it reaches ~0.067
    # TODO make dynamic plot

    f = kwargs['input']
    lines  = [line.rstrip("\n") for line in f.readlines()]

    numbers = {'1','2','3','4','5','6','7','8','9'}

    iters = []
    loss = []

    fig,ax = plt.subplots()

    for line in lines:
      args = line.split(' ')
      if args[0][-1:]==':' and args[0][0] in numbers :
        iters.append(int(args[0][:-1]))
        loss.append(float(args[2]))

    ax.plot(iters,loss)
    ax.set_xlim([kwargs['xmin'],kwargs['xmax']])
    ax.set_ylim([kwargs['ymin'],kwargs['ymax']])
    project_name = os.path.basename(os.path.dirname(kwargs['input'].name))
    plt.title('YoloV3: {}'.format(project_name))
    plt.xlabel('Iters')
    plt.ylabel('Loss')
    plt.grid()

    plt.show()


# -*- coding: utf-8 -*-

from __future__ import absolute_import

import numpy as np
import numpy.linalg as la
from PyQt5 import QtCore, QtWidgets
try:
    from PyQt5.QtWidgets import QOpenGLWidget
    has_qopenglwidget = True
except ImportError:
    from PyQt5 import QtOpenGL
    from PyQt5.QtOpenGL import QGLWidget
    has_qopenglwidget = False
from config.configuration import config
from OpenGL.GL import glReadPixels, GL_FLOAT, GL_DEPTH_COMPONENT

from util.gl_util import create_perspective_projection_matrix, create_look_at_matrix

class UpdateGLEvent(QtCore.QEvent):
    def __init__(self):
        t = QtCore.QEvent.registerEventType()
        QtCore.QEvent.__init__(self, QtCore.QEvent.Type(t))

class GLWidget(QOpenGLWidget if has_qopenglwidget else QGLWidget):
    """
    OpenGL widget to show the 3D-scene
    """

    def __init__(self, parent, main_window):
        if has_qopenglwidget:
            super(GLWidget, self).__init__(parent)
        else:
            super(GLWidget, self).__init__(QtOpenGL.QGLFormat(QtOpenGL.QGL.SampleBuffers), parent)
        self.setFocusPolicy(QtCore.Qt.StrongFocus)
        self.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
        self.update_needed = False
        self.dataset_loaded = False
        self.control = parent.control
        self.main_window = main_window
        self.setDisabled(True)

    @property
    def vis(self):
        return self.control.visualization

    def initializeGL(self):
        self.vis.assign_opengl_context(self)

    def minimumSizeHint(self):
        return QtCore.QSize(400, 400)

    def sizeHint(self):
        return QtCore.QSize(config.OpenGL.gl_window_size[0], config.OpenGL.gl_window_size[1])

    def mouseMoveEvent(self, e):
        dx = e.x() - self.x
        dy = e.y() - self.y
        if e.buttons() & QtCore.Qt.LeftButton:
            if e.modifiers() == QtCore.Qt.AltModifier:
                self.vis.translate_mouse(dx, dy)
            else:
                self.vis.rotate_mouse(dx, dy)
        self.x = e.x()
        self.y = e.y()
        self.updateGL()

    def wheelEvent(self, e):
        self.update_needed = True
	rot_v = 0.1
        if e.modifiers() == QtCore.Qt.ShiftModifier:
            if (e.angleDelta().x() != 0) or (e.angleDelta().y() != 0):
                self.vis.rotate_mouse(-e.angleDelta().x() * rot_v, -e.angleDelta().y() * rot_v)
        else:
            if e.angleDelta().y() != 0:
                self.vis.zoom(-e.angleDelta().y())

        QtWidgets.QApplication.postEvent(self, UpdateGLEvent())

    def mousePressEvent(self, e):
        if e.buttons() and QtCore.Qt.LeftButton:
            self.x = e.x()
            self.y = e.y()

    def mouseDoubleClickEvent(self, e):
        if e.buttons() and QtCore.Qt.LeftButton:
            x = e.x()
            y = self.height() - e.y()
            z = glReadPixels(x, y, 1, 1, GL_DEPTH_COMPONENT, GL_FLOAT)
            obj = self.vis.get_object_at_2dposition(x, y)
            if obj is None:
                x /= 1.0*self.width()
                y /= 1.0*self.height()
                x = x * 2 - 1
                y = y * 2 - 1
                z = z[0][0]
                z = 2 * z - 1
                z = self.vis.proj_mat[2, 3] / (-z - self.vis.proj_mat[2, 2])
                x = -x*z/self.vis.proj_mat[0, 0]
                y = -y*z/self.vis.proj_mat[1, 1]
                z += self.vis.d
                x, y, z, w = np.dot(self.vis.mat, np.array((x, y, z, 1)))
                obj = self.vis.get_object_at_3dposition(x, y, z)
            if obj is not None:
                object_type, object_index = obj
                if object_type == 'atom':
                    self.window().statistics_dock.statistics_tab.html_view.show_atom(object_index)
                elif object_type == 'domain':
                    self.window().statistics_dock.statistics_tab.html_view.show_domain(object_index)
                elif object_type == 'surface cavity':
                    self.window().statistics_dock.statistics_tab.html_view.show_surface_cavity(object_index)
                elif object_type == 'center cavity':
                    self.window().statistics_dock.statistics_tab.html_view.show_center_cavity(object_index)
                self.window().statistics_dock.raise_()
            self.updateGL()

    def customEvent(self, e):
        if self.update_needed:
            self.setDisabled(False)
            self.updateGL()
            self.update_needed = False

#     def create_scene(self, show_box, show_atoms, show_domains, show_cavities=True, center_based_cavities=False):
    def create_scene(self):
        self.vis.create_scene()
        self.updateGL()

    def keyPressEvent(self, e):
        """
        Catches and processes key presses
        """
        rot_v_key = 15
        if e.key() == QtCore.Qt.Key_Right:
            self.vis.rotate_mouse(rot_v_key, 0)
        elif e.key() == QtCore.Qt.Key_Left:
            self.vis.rotate_mouse(-rot_v_key, 0)
        elif e.key() == QtCore.Qt.Key_Up:
            self.vis.rotate_mouse(0, -rot_v_key)
        elif e.key() == QtCore.Qt.Key_Down:
            self.vis.rotate_mouse(0, rot_v_key)
        elif e.key() == QtCore.Qt.Key_D:            # Domains
            self.vis.settings.show_domains = True
            self.vis.settings.show_surface_cavities = False
            self.vis.settings.show_center_cavities = False
            self.main_window.view_dock.view_tab.domain_check.setChecked(True)
            # self.vis.create_scene()
        elif e.key() == QtCore.Qt.Key_S:            # Cavities
            self.vis.settings.show_domains = False
            self.vis.settings.show_surface_cavities = True
            self.vis.settings.show_center_cavities = False
            self.main_window.view_dock.view_tab.surface_cavity_check.setChecked(True)
            # self.vis.create_scene()
        elif e.key() == QtCore.Qt.Key_C:            # center based cavities
            self.vis.settings.show_domains = False
            self.vis.settings.show_surface_cavities = False
            self.vis.settings.show_center_cavities = True
            self.main_window.view_dock.view_tab.center_cavity_check.setChecked(True)
            # self.vis.create_scene()
        if e.key() == QtCore.Qt.Key_R:
            self.vis.reset_view()
        else:
            e.ignore()
        self.updateGL()

    def paintGL(self):
        """
        Refresh scene
        """
        self.vis.paint(self.width(), self.height(), has_qopenglwidget, self.devicePixelRatio())

    def updateGL(self):
        if has_qopenglwidget:
            self.update()
        else:
            super(GLWidget, self).updateGL()

    def activate(self):
        pass

    def updatestatus(self):
        self.update_needed = True


# coding:utf8
import cv2
import numpy as np
from model.facenet import get_model
from keras.models import Model

def get_feature_model():
    """Get face features extraction model.

    # Returns
        feat_model: Model, face features extraction model.
    """
    model = get_model((64, 64, 3))
    model.load_weights('model/weight.h5')

    feat_model = Model(inputs=model.get_layer('model_1').get_input_at(0),
                       outputs=model.get_layer('model_1').get_output_at(0))

    return feat_model

def process_image(img):
    """Resize, reduce and expand image.

    # Returns
        image: ndarray(64, 64, 3), processed image.
    """
    image = cv2.resize(img, (64, 64),
                       interpolation=cv2.INTER_CUBIC)
    image = np.array(image, dtype='float32')
    image /= 255.
    image = np.expand_dims(image, axis=0)

    return image


import os
from unittest import mock

import idelib
import numpy as np
import pandas as pd
import pytest
import hypothesis as hyp
import hypothesis.strategies as hyp_st
import hypothesis.extra.numpy as hyp_np

import endaq.batch.analyzer
from endaq.calc.stats import rms, L2_norm

np.random.seed(0)

@pytest.fixture()
def ide_SSX70065():
    with idelib.importFile(os.path.join("tests", "batch", "SSX70065.IDE")) as doc:
        yield doc

@pytest.fixture()
def analyzer_raw():
    analyzer_mock = mock.create_autospec(
        endaq.batch.analyzer.CalcCache, spec_set=False, instance=True
    )

    analyzer_mock.MPS2_TO_G = endaq.batch.analyzer.MPS2_TO_G
    analyzer_mock.MPS_TO_MMPS = endaq.batch.analyzer.MPS_TO_MMPS
    analyzer_mock.M_TO_MM = endaq.batch.analyzer.M_TO_MM
    analyzer_mock.PV_NATURAL_FREQS = endaq.batch.analyzer.CalcCache.PV_NATURAL_FREQS

    return analyzer_mock

@pytest.fixture()
def analyzer_bulk(analyzer_raw):
    analyzer_mock = analyzer_raw

    analyzer_mock._channels = {
        "acc": mock.Mock(axis_names=list("XYZ")),
        "gps": mock.Mock(axis_names=["Latitude", "Longitude"]),
        "spd": mock.Mock(axis_names=["Ground"]),
        "gyr": mock.Mock(axis_names=list("XYZ")),
        "mic": mock.Mock(axis_names=["Mic"]),
        "tmp": mock.Mock(axis_names=["Control"]),
        "pre": mock.Mock(axis_names=["Control"]),
    }

    analyzer_mock._accelerationFs = 3000
    analyzer_mock._accelerationData = pd.DataFrame(
        np.random.random((21, 3)),
        index=pd.Series(np.arange(21) / 3000, name="time"),
        columns=pd.Series(["X", "Y", "Z"], name="axis"),
    )
    analyzer_mock._accelerationResultantData = analyzer_mock._accelerationData.apply(
        L2_norm, axis="columns"
    ).to_frame()
    analyzer_mock._microphoneData = pd.DataFrame(
        np.random.random(21),
        index=pd.Series(np.arange(21) / 3000, name="time"),
        columns=pd.Series(["Mic"], name="axis"),
    )
    analyzer_mock._velocityData = pd.DataFrame(
        np.random.random((21, 3)),
        index=pd.Series(np.arange(21) / 3000, name="time"),
        columns=pd.Series(["X", "Y", "Z"], name="axis"),
    )
    analyzer_mock._displacementData = pd.DataFrame(
        np.random.random((21, 3)),
        index=pd.Series(np.arange(21) / 3000, name="time"),
        columns=pd.Series(["X", "Y", "Z"], name="axis"),
    )
    analyzer_mock._pressureData = pd.DataFrame(
        np.random.random(5),
        index=pd.Series(np.arange(5) / 5, name="time"),
        columns=pd.Series(["Control"], name="axis"),
    )
    analyzer_mock._temperatureData = pd.DataFrame(
        np.random.random(5),
        index=pd.Series(np.arange(5) / 5, name="time"),
        columns=pd.Series(["Control"], name="axis"),
    )
    analyzer_mock._gyroscopeData = pd.DataFrame(
        np.random.random(11),
        index=pd.Series(np.arange(11) / 5, name="time"),
        columns=pd.Series(["Gyro"], name="axis"),
    )

    return analyzer_mock

# ==============================================================================
# Analyzer class tests
# ==============================================================================

class TestAnalyzer:
    def test_from_ide_vs_from_literal(self, ide_SSX70065):
        dataset = ide_SSX70065
        calc_params = endaq.batch.analyzer.CalcParams(
            accel_start_time=None,
            accel_end_time=None,
            accel_start_margin=None,
            accel_end_margin=None,
            accel_highpass_cutoff=1,
            accel_integral_tukey_percent=0,
            accel_integral_zero="mean",
            psd_freq_bin_width=1,
            psd_window="hann",
            pvss_init_freq=1,
            pvss_bins_per_octave=12,
            vc_init_freq=1,
            vc_bins_per_octave=3,
        )

        dataset_cache = endaq.batch.analyzer.CalcCache.from_ide(dataset, calc_params)

        raw_cache = endaq.batch.analyzer.CalcCache.from_raw_data(
            [
                (
                    endaq.ide.to_pandas(dataset.channels[32], time_mode="timedelta"),
                    ("Acceleration", "g"),
                ),
                (
                    endaq.ide.to_pandas(
                        dataset.channels[36].subchannels[0], time_mode="timedelta"
                    ),
                    ("Pressure", "Pa"),
                ),
                (
                    endaq.ide.to_pandas(
                        dataset.channels[36].subchannels[1], time_mode="timedelta"
                    ),
                    ("Temperature", "°C"),
                ),
            ],
            calc_params,
        )

        assert set(dataset_cache._channels) == set(raw_cache._channels)

        for (ds_struct, raw_struct) in (
            (dataset_cache._channels[measure_key], raw_cache._channels[measure_key])
            for measure_key in dataset_cache._channels
        ):
            assert ds_struct.units == raw_struct.units
            pd.testing.assert_frame_equal(
                ds_struct.to_pandas(time_mode="timedelta"),
                raw_struct.to_pandas(time_mode="timedelta"),
            )

    @hyp.given(
        df=hyp_np.arrays(
            elements=hyp_st.floats(-1e7, 1e7),
            shape=(20, 2),
            dtype=np.float64,
        ).map(
            lambda array: pd.DataFrame(
                array, index=np.timedelta64(200, "ms") * np.arange(20)
            )
        ),
    )
    def test_accelerationData(self, df):
        calc_params = endaq.batch.analyzer.CalcParams(
            accel_start_time=None,
            accel_end_time=None,
            accel_start_margin=None,
            accel_end_margin=None,
            accel_highpass_cutoff=1,
            accel_integral_tukey_percent=0,
            accel_integral_zero="mean",
            psd_freq_bin_width=1,
            psd_window="hann",
            pvss_init_freq=1,
            pvss_bins_per_octave=12,
            vc_init_freq=1,
            vc_bins_per_octave=3,
        )
        data_cache = endaq.batch.analyzer.CalcCache.from_raw_data(
            [(df, ("Acceleration", "m/s\u00b2"))], calc_params
        )

        df_accel = endaq.calc.filters.butterworth(
            df, low_cutoff=calc_params.accel_highpass_cutoff
        )
        pd.testing.assert_frame_equal(data_cache._accelerationData, df_accel)

        (_df_accel, df_vel, df_displ) = endaq.calc.integrate.integrals(
            df_accel,
            n=2,
            zero=calc_params.accel_integral_zero,
            highpass_cutoff=calc_params.accel_highpass_cutoff,
            tukey_percent=calc_params.accel_integral_tukey_percent,
        )
        pd.testing.assert_frame_equal(data_cache._velocityData, df_vel)
        pd.testing.assert_frame_equal(data_cache._displacementData, df_displ)

    def test_accRMSFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.accRMSFull.func(analyzer_bulk)[
            "Resultant"
        ]
        expt_result = endaq.batch.analyzer.MPS2_TO_G * rms(
            analyzer_bulk._accelerationData.apply(L2_norm, axis="columns")
        )

        assert calc_result == pytest.approx(expt_result)

    def test_velRMSFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.velRMSFull.func(analyzer_bulk)[
            "Resultant"
        ]
        expt_result = endaq.batch.analyzer.MPS_TO_MMPS * rms(
            analyzer_bulk._velocityData.apply(L2_norm, axis="columns")
        )
        assert calc_result == pytest.approx(expt_result)

    def test_disRMSFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.disRMSFull.func(analyzer_bulk)[
            "Resultant"
        ]
        expt_result = endaq.batch.analyzer.M_TO_MM * rms(
            analyzer_bulk._displacementData.apply(L2_norm, axis="columns")
        )
        assert calc_result == pytest.approx(expt_result)

    def test_accPeakFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.accPeakFull.func(analyzer_bulk)[
            "Resultant"
        ]
        expt_result = endaq.batch.analyzer.MPS2_TO_G * rms(
            analyzer_bulk._accelerationData.apply(L2_norm, axis="columns").max()
        )

        assert calc_result == pytest.approx(expt_result)

    def test_pseudoVelPeakFull(self, analyzer_bulk):
        pass

    def test_gpsLocFull(self, analyzer_bulk):
        pass

    def test_gpsSpeedFull(self, analyzer_bulk):
        pass

    def test_gyroRMSFull(self, analyzer_bulk):
        pass

    def test_micRMSFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.micRMSFull.func(analyzer_bulk)[
            "Mic"
        ]
        expt_result = rms(analyzer_bulk._microphoneData)

        assert calc_result == pytest.approx(expt_result)

    def test_pressFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.pressFull.func(analyzer_bulk)[
            "Control"
        ]
        expt_result = analyzer_bulk._pressureData.mean()

        assert calc_result == pytest.approx(expt_result)

    def test_tempFull(self, analyzer_bulk):
        calc_result = endaq.batch.analyzer.CalcCache.tempFull.func(analyzer_bulk)[
            "Control"
        ]
        expt_result = analyzer_bulk._temperatureData.mean()

        assert calc_result == pytest.approx(expt_result)

    ###########################################################################
    # Live File Tests

    def testLiveFile(self, ide_SSX70065):
        analyzer = endaq.batch.analyzer.CalcCache.from_ide(
            ide_SSX70065,
            endaq.batch.analyzer.CalcParams(
                accel_start_time=None,
                accel_end_time=None,
                accel_start_margin=None,
                accel_end_margin=None,
                accel_highpass_cutoff=1,
                accel_integral_tukey_percent=0,
                accel_integral_zero="mean",
                psd_freq_bin_width=1,
                psd_window="hann",
                pvss_init_freq=1,
                pvss_bins_per_octave=12,
                vc_init_freq=1,
                vc_bins_per_octave=3,
            ),
        )

        raw_accel = ide_SSX70065.channels[32].getSession().arrayValues()
        np.testing.assert_allclose(
            analyzer.accRMSFull["Resultant"],
            rms(L2_norm(raw_accel - raw_accel.mean(axis=-1, keepdims=True), axis=0)),
            rtol=1e-3,
        )
        np.testing.assert_allclose(
            analyzer.pressFull,
            1e-3
            * ide_SSX70065.channels[36]
            .subchannels[0]
            .getSession()
            .arrayValues()
            .mean(),
            rtol=1e-5,
        )
        np.testing.assert_allclose(
            analyzer.tempFull,
            ide_SSX70065.channels[36].subchannels[1].getSession().arrayValues().mean(),
            rtol=1e-4,
        )

    @pytest.mark.parametrize(
        "filename",
        [
            os.path.join("tests", "batch", "test1.IDE"),
            os.path.join("tests", "batch", "test2.IDE"),
        ],
    )
    def testLiveFiles1(self, filename):
        ds = idelib.importFile(filename)
        analyzer = endaq.batch.analyzer.CalcCache.from_ide(
            ds,
            endaq.batch.analyzer.CalcParams(
                accel_start_time=None,
                accel_end_time=None,
                accel_start_margin=None,
                accel_end_margin=None,
                accel_highpass_cutoff=1,
                accel_integral_tukey_percent=0,
                accel_integral_zero="mean",
                psd_freq_bin_width=1,
                psd_window="hann",
                pvss_init_freq=1,
                pvss_bins_per_octave=12,
                vc_init_freq=1,
                vc_bins_per_octave=3,
            ),
        )

        raw_accel = ds.channels[8].getSession().arrayValues(subchannels=[0, 1, 2])
        np.testing.assert_allclose(
            analyzer.accRMSFull["Resultant"],
            rms(L2_norm(raw_accel - raw_accel.mean(axis=-1, keepdims=True), axis=0)),
            rtol=0.55,  # this is... probably a little too high...
        )
        audio_scale = 5.307530522779073
        np.testing.assert_allclose(
            analyzer.micRMSFull,
            audio_scale * rms(ds.channels[8].subchannels[3].getSession().arrayValues()),
            rtol=1e-3,
        )
        np.testing.assert_allclose(
            analyzer.gyroRMSFull["Resultant"],
            rms(L2_norm(ds.channels[84].getSession().arrayValues(), axis=0)),
            rtol=1e-2,
        )
        np.testing.assert_allclose(
            analyzer.pressFull,
            1e-3 * ds.channels[36].subchannels[0].getSession().arrayValues().mean(),
            rtol=1e-3,
        )
        np.testing.assert_allclose(
            analyzer.tempFull,
            ds.channels[36].subchannels[1].getSession().arrayValues().mean(),
            rtol=0.05,  # ...GEFGW?
        )
        np.testing.assert_allclose(
            analyzer.humidFull,
            ds.channels[59].subchannels[2].getSession().arrayValues().mean(),
        )

    @pytest.mark.parametrize(
        "filename",
        [
            os.path.join("tests", "batch", "DAQ12006_000005.IDE"),
        ],
    )
    def testLiveFiles2(self, filename):
        """
        Checks that audio in units of Pascals are properly handled.
        """
        ds = idelib.importFile(filename)
        analyzer = endaq.batch.analyzer.CalcCache.from_ide(
            ds,
            endaq.batch.analyzer.CalcParams(
                accel_start_time=None,
                accel_end_time=None,
                accel_start_margin=None,
                accel_end_margin=None,
                accel_highpass_cutoff=1,
                accel_integral_tukey_percent=0,
                accel_integral_zero="mean",
                psd_freq_bin_width=1,
                psd_window="hann",
                pvss_init_freq=1,
                pvss_bins_per_octave=12,
                vc_init_freq=1,
                vc_bins_per_octave=3,
            ),
        )

        np.testing.assert_allclose(
            analyzer.micRMSFull,
            rms(ds.channels[8].subchannels[3].getSession().arrayValues()),
            rtol=1e-3,
        )

    @pytest.mark.parametrize(
        "filename",
        [
            os.path.join("tests", "batch", "test3.IDE"),
        ],
    )
    def testLiveFile3(self, filename):
        ds = idelib.importFile(filename)
        analyzer = endaq.batch.analyzer.CalcCache.from_ide(
            ds,
            endaq.batch.analyzer.CalcParams(
                accel_start_time=None,
                accel_end_time=None,
                accel_start_margin=None,
                accel_end_margin=None,
                accel_highpass_cutoff=1,
                accel_integral_tukey_percent=0,
                accel_integral_zero="mean",
                psd_freq_bin_width=1,
                psd_window="hann",
                pvss_init_freq=1,
                pvss_bins_per_octave=12,
                vc_init_freq=1,
                vc_bins_per_octave=3,
            ),
        )

        ch_rot = ds.channels[84]
        np.testing.assert_allclose(
            analyzer.gyroRMSFull["Resultant"],
            rms(L2_norm(ch_rot.getSession().arrayValues(), axis=0)),
            rtol=0.005,
        )

    @pytest.mark.parametrize(
        "filename, sample_index",
        [
            (os.path.join("tests", "batch", "test_GPS_2.IDE"), -2),
            (os.path.join("tests", "batch", "test_GPS_3.IDE"), -4),
        ],
    )
    def testLiveFileGPS(self, filename, sample_index):
        ds = idelib.importFile(filename)
        analyzer = endaq.batch.analyzer.CalcCache.from_ide(
            ds,
            endaq.batch.analyzer.CalcParams(
                accel_start_time=None,
                accel_end_time=None,
                accel_start_margin=None,
                accel_end_margin=None,
                accel_highpass_cutoff=1,
                accel_integral_tukey_percent=0,
                accel_integral_zero="mean",
                psd_freq_bin_width=1,
                psd_window="hann",
                pvss_init_freq=1,
                pvss_bins_per_octave=12,
                vc_init_freq=1,
                vc_bins_per_octave=3,
            ),
        )

        ch_gps = ds.channels[88]
        # confirming channel format for gps file
        assert ch_gps.subchannels[0].name == "Latitude"
        assert ch_gps.subchannels[1].name == "Longitude"
        assert ch_gps.subchannels[3].name == "Ground Speed"

        assert tuple(analyzer.gpsLocFull) == tuple(
            ch_gps.getSession().arrayValues()[[0, 1], sample_index],
        )
        # Resampling throws off mean calculation
        # -> use resampled data for comparsion
        # gps_speed = analyzer.MPS_TO_KMPH * ch_gps.getSession().arrayValues(subchannels=[3])
        gps_speed = analyzer._gpsSpeedData
        np.testing.assert_allclose(
            analyzer.gpsSpeedFull,
            np.mean(gps_speed[gps_speed != 0]),
        )


""" Use the MDRNN to predict the future. """
import torch
from torch.distributions.categorical import Categorical
from torch.nn import functional as F
from torchvision import transforms
from torchvision.utils import save_image
import argparse
from os.path import join, exists
from os import mkdir
import numpy as np
from models import MDRNNCell, VAE, Controller
import gym
import gym.envs.box2d
from utils.misc import ACTION_SIZE, LATENT_SIZE, LATENT_RECURRENT_SIZE, IMAGE_RESIZE_DIM, SIZE
from utils.misc import sample_continuous_policy

## WARNING : THIS SHOULD BE REPLACE WITH PYTORCH 0.5
env = gym.make("CarRacing-v0")
env.reset()
seq_len=1000
time_limit = seq_len
dream_point = 30
device = 'cpu' #torch.device('cuda' if torch.cuda.is_available() else 'cpu')

a_rollout = sample_continuous_policy(env.action_space, seq_len, 1. / 50)

mdir = 'exp_dir'
# load in the VAE and MDRNN:
vae_file, rnn_file, ctrl_file = \
    [join(mdir, m, 'best.tar') for m in ['vae', 'mdrnn', 'ctrl']]

assert exists(vae_file) and exists(rnn_file),\
    "Either vae or mdrnn is untrained or the file is in the wrong place."

vae_state, rnn_state = [
    torch.load(fname, map_location={'cuda:0': str(device)})
    for fname in (vae_file, rnn_file)]

for m, s in (('VAE', vae_state), ('MDRNN', rnn_state)):
    print("Loading {} at epoch {} "
        "with test loss {}".format(
            m, s['epoch'], s['precision']))

vae = VAE(3, LATENT_SIZE).to(device)
vae.load_state_dict(vae_state['state_dict'])

mdrnn = MDRNNCell(LATENT_SIZE, ACTION_SIZE, LATENT_RECURRENT_SIZE, 5).to(device)
mdrnn.load_state_dict(
    {k.strip('_l0'): v for k, v in rnn_state['state_dict'].items()})

'''controller = Controller(LATENT_SIZE, LATENT_RECURRENT_SIZE, ACTION_SIZE).to(device)

# load controller if it was previously saved
if exists(ctrl_file):
    ctrl_state = torch.load(ctrl_file, map_location={'cuda:0': str(device)})
    print("Loading Controller with reward {}".format(
        ctrl_state['reward']))
    controller.load_state_dict(ctrl_state['state_dict'])'''

transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((IMAGE_RESIZE_DIM, IMAGE_RESIZE_DIM)),
    transforms.ToTensor()
])

def rollout(render=False):
    """ Execute a rollout and returns minus cumulative reward.

    Load :params: into the controller and execute a single rollout. This
    is the main API of this class.

    :args params: parameters as a single 1D np array

    :returns: minus cumulative reward
    # Why is this the minus cumulative reward?!?!!?
    """
    print('a rollout dims', len(a_rollout))

    #env.seed(int(rand_env_seed)) # ensuring that each rollout has a differnet random seed.
    obs = env.reset()

    # This first render is required !
    env.render()

    next_hidden = [
        torch.zeros(1, LATENT_RECURRENT_SIZE).to(device)
        for _ in range(2)]

    cumulative = 0
    i = 0
    rollout_dict = {k:[] for k in ['obs', 'rew', 'act', 'term']}

    obs = transform(obs).unsqueeze(0).to(device)
    mu, logsigma = vae.encoder(obs)
    next_z =  mu + logsigma.exp() * torch.randn_like(mu)

    while True:
        #print(i)

        action = torch.Tensor(a_rollout[i]).to(device).unsqueeze(0)

        #print('into mdrnn',action.shape, next_z.shape, next_hidden[0].shape)

        # commented out reward and done.
        mus, sigmas, logpi, _, _, next_hidden = mdrnn(action, next_z, next_hidden)

        # decode current z to see what it looks like.
        recon_obs = vae.decoder(next_z)

        if i>dream_point:
            if type(obs) != torch.Tensor:
                obs = transform(obs).unsqueeze(0)
            to_save = torch.cat([obs, recon_obs.cpu()], dim=0)
            #print(to_save.shape)
            # .view(args.batch_size*2, 3, IMAGE_RESIZE_DIM, IMAGE_RESIZE_DIM)
            save_image(to_save,
                        join(mdir, 'dream/sample_' + str(i) + '.png'))

        obs, reward, done, _ = env.step(a_rollout[i])

        if i < dream_point or np.random.random()>0.95:
            print('using real obs at point:', i)
            obs = transform(obs).unsqueeze(0).to(device)
            mu, logsigma = vae.encoder(obs)
            next_z =  mu + logsigma.exp() * torch.randn_like(mu)
        else:
            # sample the next z.
            g_probs = Categorical(probs=torch.exp(logpi).permute(0,2,1))
            which_g = g_probs.sample()
            #print(logpi.shape, mus.permute(0,2,1)[:,which_g].shape ,mus[:,:,which_g].shape,  which_g, mus.shape )
            #print(mus.squeeze().permute(1,0).shape, which_g.permute(1,0))
            mus_g, sigs_g = torch.gather(mus.squeeze(), 0, which_g), torch.gather(sigmas.squeeze(), 0, which_g)
            #print(mus_g.shape)
            next_z = mus_g + sigs_g * torch.randn_like(mus_g)
            #print(next_z.shape)

        #for key, var in zip(['obs', 'rew', 'act', 'term'], [obs,reward, action, done]):
        #    rollout_dict[key].append(var)

        if render:
            env.render()

        cumulative += reward
        if done or i >= time_limit:

            return - cumulative

        i += 1

rollout()

'''if __name__ == "__main__":'''


from avatarpy.core import Core
from avatarpy.transform import Transform
from avatarpy.animate import Animate
from avatarpy.describe import Describe
from avatarpy.annotation import Annotation

import numpy as np
import pandas as pd
from itertools import combinations
from scipy.stats import zscore

class Avatar(Core):
    _nodes={
        'nose'  :[0,1,2],
        'neck'  :[3,4,5],
        'anus'  :[6,7,8],
        'chest' :[9,10,11],
        'rfoot' :[12,13,14],
        'lfoot' :[15,16,17],
        'rhand' :[18,19,20],
        'lhand' :[21,22,23],
        'tip'   :[24,25,26],
    }

    _vectors={
        'head'      :{'head':'nose',    'tail':'neck'},
        'fbody'     :{'head':'neck',    'tail':'chest'},
        'hbody'     :{'head':'chest',   'tail':'anus'},
        'tail'      :{'head':'tip',     'tail':'anus'},
        'rarm'      :{'head':'rhand',   'tail':'chest'},
        'larm'      :{'head':'lhand',   'tail':'chest'},
        'rleg'      :{'head':'rfoot',   'tail':'anus'},
        'lleg'      :{'head':'lfoot',   'tail':'anus'},
        'nose2anus' :{'head':'anus',   'tail':'nose'},
        'lhand2nose':{'head':'nose',   'tail':'lhand'},
        'rhand2nose':{'head':'nose',   'tail':'rhand'},
    }

    _angles={
        'head':{'left':'head', 'right':'fbody'},
        'body':{'left':'fbody', 'right':'hbody'},
        'tail':{'left':'hbody', 'right':'tail'},
        'larm':{'left':'larm', 'right':'fbody'},
        'rarm':{'left':'rarm', 'right':'fbody'},
        'lleg':{'left':'lleg', 'right':'hbody'},
        'rleg':{'left':'rleg', 'right':'hbody'},
    }

    def __init__(self, csv_path, frame_rate=20, ID=None, tags={}, horizontal_correction=True):
        self._csv_path = csv_path
        self._data = pd.read_csv(csv_path, header=None)
        self._frame_rate = frame_rate
        if frame_rate: self.data.index/=frame_rate
        self._ID=ID if ID else self.csv_path
        self._tags = tags
        self.set_nodes()
        self.set_vectors()

        self._transform = Transform(parent=self)
        self._animate = Animate(parent=self)
        self._describe = Describe(parent=self)
        self._annotation = Annotation(parent=self)
        # self._human_annotation = HumanAnnotation(parent=self)
        # self._heuristic_annotation = HeuristicAnnotation(parent=self)

        if horizontal_correction:
            self.data = self.transform.level().data
            self.set_nodes()
            self.set_vectors()

    @property
    def csv_path(self):
        """Original coordinates file path"""
        return self._csv_path
    @csv_path.setter
    def csv_path(self, v):
        self._csv_path = v
    @property
    def data(self):
        """Raw data of repeated x, y, z coordinates of nodes"""
        return self._data
    @data.setter
    def data(self, v):
        self._data = v
    @property
    def frame_rate(self):
        """Number of frames recorded in 1 second. (a.k.a. data rate, sampling rate)"""
        return self._frame_rate
    @frame_rate.setter
    def frame_rate(self, v):
        self._frame_rate = v
    @property
    def ID(self):
        """User provided ID for avatar instance. (default: csv_path)"""
        return self._ID
    @ID.setter
    def ID(self, v):
        self._ID = v
    @property
    def tags(self):
        """User provided tags for avatar instance. add tags by dict (ex, {'genotype':'wt'})"""
        return self._tags
    @tags.setter
    def tags(self, v):
        self._tags = v

    def __repr__(self):
        return f'Avatar({self.ID})'

    @property
    def help(self):
        """Prints list of available functions and attributes of avatar"""
        functions, attributes  = [], []
        for item in dir(self.__class__):
            if item.startswith('_'):
                continue
            docstring = getattr(self.__class__, item).__doc__
            headline = docstring.split('\n')[0].strip() if docstring else ''
            if callable(getattr(self.__class__, item)):
                functions.append(f"\t{item}: \n\t\t{headline}")
            else:
                attributes.append(f"\t{item}: \n\t\t{headline}")
        funcs = '\n'.join(functions)
        attrs = '\n'.join(attributes)
        print(f"""Visit https://github.com/jeakwon/avatarpy\n
            Functions:\n{funcs}\nAttributes:\n{attrs}""")

    def set_nodes(self):
        """Set node attributes in avatar with predefined node info in `cls._nodes`"""
        for name, cols in self._nodes.items():
            data = self.get_node(cols)
            setattr(self, name, data)

    def set_vectors(self):
        """Set vector attributes in avatar with predefined vector info in `cls._nodes`"""
        for name, labels in self._vectors.items():
            data = self.get_vector(self[labels['head']], self[labels['tail']])
            setattr(self, name, data)

    @property
    def index(self):
        """Index of recording timecourse. If frame rate is provided, unit is second"""
        return self.data.index

    @property
    def nodes(self):
        """Dictionary of all nodes data"""
        return {key:self[key] for key in self._nodes.keys()}

    @property
    def vectors(self):
        """Dictionary of all vectors data"""
        return {key:self[key] for key in self._vectors.keys()}

    @property
    def x(self):
        """x coordindates of nodes and vectors"""
        return self.get_axis_data('x')

    @property
    def y(self):
        """y coordindates of nodes and vectors"""
        return self.get_axis_data('y')

    @property
    def z(self):
        """z coordindates of nodes and vectors"""
        return self.get_axis_data('z')

    @property
    def x_max(self):
        """Maximum x coord value among all nodes"""
        return self.x[self.nodes].values.max()

    @property
    def x_min(self):
        """Minimum x coord value among all nodes"""
        return self.x[self.nodes].values.min()

    @property
    def y_max(self):
        """Maximum y coord value among all nodes"""
        return self.y[self.nodes].values.max()

    @property
    def y_min(self):
        """Minimum y coord value among all nodes"""
        return self.y[self.nodes].values.min()

    @property
    def z_max(self):
        """Maximum z coord value among all nodes"""
        return self.z[self.nodes].values.max()

    @property
    def z_min(self):
        """Minimum z coord value among all nodes"""
        return self.z[self.nodes].values.min()

    @property
    def unit_vector_x(self):
        """Numpy 2d array (N x 3) of x unit vector [1,0,0]"""
        return self.get_unit_vector(axis='x')

    @property
    def unit_vector_y(self):
        """Numpy 2d array (N x 3) of y unit vector [0,1,0]"""
        return self.get_unit_vector(axis='y')

    @property
    def unit_vector_z(self):
        """Numpy 2d array (N x 3) of z unit vector. [0,0,1]"""
        return self.get_unit_vector(axis='z')

    @property
    def node_data(self):
        """Coords data with labels of all nodes. [x, y, z, node]"""
        return self.get_node_data(self.nodes.keys())

    def get_node_data(self, nodes):
        """Returns coords data with labels of provided nodes names"""
        labeled_data = [self[node].assign(node=node) for node in nodes]
        return pd.concat(labeled_data).sort_index()

    def get_unit_vector(self, axis):
        """Returns numpy 2d array (N x 3) of given axis unit vector"""
        if axis=='x': arr = np.array([1,0,0])
        elif axis=='y': arr = np.array([0,1,0])
        elif axis=='z': arr = np.array([0,1,0])
        return np.stack([arr]*len(self.index))

    def get_node(self, columns):
        """Returns T-series node x, y, z coords by assigned index of columns"""
        return self.data[columns].set_axis(['x', 'y', 'z'], axis=1, inplace=False)

    def get_vector(self, head, tail):
        """Returns T-series vector x, y, z coords by assigned index of columns"""
        return pd.DataFrame(head.values-tail.values, columns=['x', 'y', 'z']).set_index(self.data.index)

    def get_axis_data(self, axis):
        """Returns axis data of all nodes and vectors"""
        data_dict = {}
        for name, data in self.nodes.items():
            data_dict[name]=data[axis].values
        for name, data in self.vectors.items():
            data_dict[name]=data[axis].values
        return pd.DataFrame(data_dict).set_index(self.data.index)

    def get_triangular_area_by_nodes(self, node1, node2, node3):
        r"""Calcultes triangular area by node names"""
        coord1, coord2, coord3 = self[node1], self[node2], self[node3]
        name = '_'.join([node1, node2, node3])
        data = self.get_triangular_area_by_coords(coord1, coord2, coord3)
        return pd.Series(data=data, name=name, index=self.data.index)

    @property
    def area(self):
        """Returns T-series areas from all combination of nodes"""
        nodes_combinations = list(combinations(self.nodes.keys(), 3))
        return pd.DataFrame([self.get_triangular_area_by_nodes(*c) for c in nodes_combinations]).T

    @property
    def angle(self):
        """Returns T-series angles between predefined two vectors"""
        data_dict = {}
        for name, vectors in self._angles.items():
            data_dict[name]=self.get_angle(self[vectors['left']], self[vectors['right']])
        return pd.DataFrame(data_dict).set_index(self.data.index)

    @property
    def angle_velocity(self):
        """Returns T-series angles velocity between predefined two vectors"""
        return self.angle.diff()*self.frame_rate

    @property
    def angle_acceleration(self):
        """Returns T-series angles acceleration between predefined two vectors"""
        return self.angle_velocity.diff()*self.frame_rate

    @property
    def vector_length(self):
        """Returns length of all vectors"""
        data_dict = {}
        for name, vector in self.vectors.items():
            data_dict[name]=self.get_distance(vector)
        return pd.DataFrame(data_dict).set_index(self.data.index)

    @property
    def stretch_index(self):
        """Returns stretch_index which is equal to zscore of vector length"""
        return self.vector_length.apply(zscore)

    @property
    def distance(self):
        """Returns inter-frame distances of all coords"""
        data_dict = {}
        for name, node in self.nodes.items():
            data_dict[name]=self.get_distance(node.diff())
        return pd.DataFrame(data_dict).set_index(self.data.index)

    @property
    def velocity(self):
        """Returns moment velocity of all coords"""
        return self.distance*self.frame_rate

    @property
    def acceleration(self):
        """Returns moment acceleration of all coords"""
        return self.velocity.diff()*self.frame_rate

    @property
    def cummulative_distance(self):
        """Returns T-series cumulative distance of all coords"""
        return self.distance.cumsum()

    @property
    def total_distance(self):
        """Returns total explored distance of all coords """
        return self.cummulative_distance.max()

    def get_vector_coords(self, vector_name):
        """Retruns dataframe dicts of nodes given by vector name."""
        return dict(
            x = self.x[self._vectors[vector_name].values()],
            y = self.y[self._vectors[vector_name].values()],
            z = self.z[self._vectors[vector_name].values()],
        )

    def get_projection(self, vector, to):
        """Project 3D coordinates on to axis or plane"""
        zero_axis = list(set('xyz')-set(to))
        kwargs = {axis:0 for axis in zero_axis}
        return vector.assign(**kwargs)

    def x_projection(self, node):
        """Project 3D coordinates on to x axis"""
        return self.get_projection(vector=self[node], to='x')

    def y_projection(self, node):
        """Project 3D coordinates on to y axis"""
        return self.get_projection(vector=self[node], to='y')

    def z_projection(self, node):
        """Project 3D coordinates on to z axis"""
        return self.get_projection(vector=self[node], to='z')

    def xy_projection(self, node):
        """Project 3D coordinates on to xy plane"""
        return self.get_projection(vector=self[node], to='xy')

    def yz_projection(self, node):
        """Project 3D coordinates on to yz plane"""
        return self.get_projection(vector=self[node], to='yz')

    def xz_projection(self, node):
        """Project 3D coordinates on to xz plane"""
        return self.get_projection(vector=self[node], to='xz')

    @property
    def transform(self):
        """Transform module for coordinate change"""
        return self._transform

    @property
    def animate(self):
        """Animation module for coordinate change"""
        return self._animate

    @property
    def describe(self):
        """Feature Extract module for coordinate change"""
        return self._describe

    @property
    def annotation(self):
        """Annotation module for behavior screening"""
        return self._annotation

    def corr(self, data, window=None, center=True, **kwargs):
        """Returns rolling correlation with given property of data

        :param data: (str|pd.DataFrame)
        :param window: (int|None) if provided rolling corr is provided
        """
        if isinstance(data, str):
            data = self[data]
        if window==None:
            return self.flatten_pairwise_df(data.corr())
        return self.get_rolling_corr(data, window, center, **kwargs)

    def xcorr_lag(self, data, flatten=True):
        """Returns cross correlation time lag by parwise column calculation. See also xcorr_max
        """
        if isinstance(data, str):
            data = self[data]
        df = self.pairwise_apply_func(data, lambda in1, in2: self.xcorr(in1, in2)['lag']/self.frame_rate)
        if flatten:
            return self.flatten_pairwise_df(df)
        return df

    def xcorr_max(self, data, flatten=True):
        """Returns cross correlation max value by parwise column calculation. See also xcorr_lag
        """
        if isinstance(data, str):
            data = self[data]
        df = self.pairwise_apply_func(data, lambda in1, in2: self.xcorr(in1, in2)['max'])
        if flatten:
            return self.flatten_pairwise_df(df)
        return data

    def apply(self, func):
        return func(self)

    def gather(self, features=['x', 'y', 'z', 'aop_x', 'aop_y', 'aop_z', 'aoa_x', 'aoa_y', 'aoa_z', 'velocity', 'acceleration', 'angle', 'angle_velocity', 'angle_acceleration', 'stretch_index', 'area'], multi_index=False):
        df = pd.concat({feature:self[feature] for feature in features}, axis=1)
        if not multi_index:
            df.columns = df.columns.map('_'.join)
        return df

    def gather_corrs(self, window=20, features=['velocity', 'acceleration', 'angle', 'angle_velocity', 'angle_acceleration', 'stretch_index'], multi_index=False):
        df = pd.concat({f:self.corr(f, window=window) for f in features}, axis=1)
        if not multi_index:
            df.columns = df.columns.map('_'.join)
        return df

    def aop(self, *args, **kwargs):
        """Transforms avatar coords. Align on plane.
        """
        return self.transform.align_on_plane(*args, **kwargs)

    def aoa(self, *args, **kwargs):
        """Transforms avatar coords. Align on axis.
        """
        return self.transform.align_on_axis(*args, **kwargs)

    @property
    def aop_x(self):
        return self.aop()['x']

    @property
    def aop_y(self):
        return self.aop()['y']

    @property
    def aop_z(self):
        return self.aop()['z']

    @property
    def aoa_x(self):
        return self.aoa()['x']

    @property
    def aoa_y(self):
        return self.aoa()['y']

    @property
    def aoa_z(self):
        return self.aoa()['z']