from __future__ import division
from build.cioverlap import *
from qm.dftbplus.dftbplus import DFTBplus
from qm.dftbplus.dftbpar import spin_w, spin_w_lc, onsite_uu, onsite_ud, max_l
from misc import data, eps, eV_to_au, call_name
import os, shutil, re, textwrap
import numpy as np

class DFTB(DFTBplus):
    """ Class for (TD)DFTB method of DFTB+

        :param object molecule: Molecule object
        :param boolean l_scc: Include self-consistent charge (SCC) scheme
        :param double scc_tol: Stopping criteria for the SCC iterations
        :param integer scc_max_iter: Maximum number of SCC iterations
        :param boolean l_onsite: Include onsite correction to SCC term
        :param boolean l_range_sep: Include long-range corrected functional
        :param string lc_method: Algorithms for LC-DFTB
        :param boolean l_spin_pol: Include spin-polarisation scheme
        :param double unpaired_elec: Number of unpaired electrons
        :param string guess: Initial guess method for SCC scheme
        :param string guess_file: Initial guess file for charges
        :param double elec_temp: Electronic temperature in Fermi-Dirac scheme
        :param string mixer: Charge mixing method used in DFTB
        :param string ex_symmetry: Symmetry of excited state in TDDFTB
        :param double e_window: Energy window for TDDFTB. Increases efficiency of NACME calculation.
        :param integer,list k_point: Number of k-point samplings
        :param boolean l_periodic: Use periodicity in the calculations
        :param double,list cell_length: The lattice vectors of periodic unit cell
        :param string sk_path: Path for Slater-Koster files
        :param string install_path: Path for DFTB+ install directory
        :param boolean mpi: Use MPI parallelization
        :param string mpi_path: Path for MPI binary
        :param integer nthreads: Number of threads in the calculations
        :param string version: Version of DFTB+
    """
    def __init__(self, molecule, l_scc=True, scc_tol=1E-6, scc_max_iter=100, l_onsite=False, \
        l_range_sep=False, lc_method="MatrixBased", l_spin_pol=False, unpaired_elec=0., guess="h0", \
        guess_file="./charges.bin", elec_temp=0., mixer="Broyden", ex_symmetry="singlet", e_window=0., \
        k_point=[1, 1, 1], l_periodic=False, cell_length=[0., 0., 0., 0., 0., 0., 0., 0., 0.,], \
        sk_path="./", install_path="./", mpi=False, mpi_path="./", nthreads=1, version="20.1"):
        # Initialize DFTB+ common variables
        super(DFTB, self).__init__(molecule, sk_path, install_path, nthreads, version)

        # Initialize DFTB+ DFTB variables
        self.l_scc = l_scc
        self.scc_tol = scc_tol
        self.scc_max_iter = scc_max_iter

        self.l_onsite = l_onsite

        self.l_range_sep = l_range_sep
        self.lc_method = lc_method.lower()

        self.l_spin_pol = l_spin_pol
        self.unpaired_elec = unpaired_elec

        # Set initial guess for SCC term
        self.guess = guess.lower()
        self.guess_file = guess_file
        if not (self.guess in ["h0", "read"]):
            error_message = "Invalid initial guess for DFTB!"
            error_vars = f"guess = {self.guess}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.elec_temp = elec_temp
        self.mixer = mixer.lower()

        self.ex_symmetry = ex_symmetry.lower()
        self.e_window = e_window

        self.k_point = k_point
        self.l_periodic = l_periodic
        self.a_axis = np.copy(cell_length[0:3])
        self.b_axis = np.copy(cell_length[3:6])
        self.c_axis = np.copy(cell_length[6:9])

        # Check excitation symmetry in TDDFTB
        # TODO : Currently, allows only singlet excited states with TDDFTB
#        if not (self.ex_symmetry in ["singlet", "triplet"]):
        if (not self.ex_symmetry == "singlet"):
            error_message = "Invalid symmetry of excited states for TDDFTB given!"
            error_vars = f"ex_symmetry = {self.ex_symmetry}"
            raise ValueError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        self.mpi = mpi
        self.mpi_path = mpi_path

        # Set 'l_nacme' and 're_calc' with respect to the computational method
        # TDDFTB do not produce NACs, so we should get NACME from CIoverlap
        # TDDFTB cannot compute the gradient of several states simultaneously.
        molecule.l_nacme = True
        self.re_calc = True

        # Calculate number of basis for current system
        # Set new variable to decide the position of basis functions in terms of atoms
        # DFTB method considers only valence electrons, so core electrons should be removed
        core_elec = 0.
        self.nbasis = 0
        self.check_atom = [0]
        for iat in range(molecule.nat_qm):
            # Check number of basis functions with respect to maximum angular momentum
            max_ang = max_l[molecule.symbols[iat]]
            if (max_ang == 's'):
                self.nbasis += 1
            elif (max_ang == 'p'):
                self.nbasis += 4
            elif (max_ang == 'd'):
                self.nbasis += 9
            else:
                error_message = "Number of basis for f orbital not implemented, see '$PYUNIXMDHOME/src/qm/dftbplus/dftb.py'!"
                error_vars = f"current atom = {molecule.symbols[iat]}, max_ang = {max_ang}"
                raise NotImplementedError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")
            self.check_atom.append(self.nbasis)
            # Check number of core electrons with respect to atomic number
            sym_index = list(data.keys()).index(molecule.symbols[iat])
            if (sym_index > 0 and sym_index <= 2):
                core_elec += 0.
            elif (sym_index > 2 and sym_index <= 10):
                core_elec += 2.
            elif (sym_index > 10 and sym_index <= 18):
                core_elec += 10.
            elif (sym_index > 18 and sym_index <= 36):
                core_elec += 18.
            elif (sym_index > 36 and sym_index <= 54):
                core_elec += 36.
            else:
                error_message = "Core electrons for current element not implemented, see '$PYUNIXMDHOME/src/qm/dftbplus/dftb.py'!"
                error_vars = f"current atom = {molecule.symbols[iat]}, sym_index = {sym_index}"
                raise NotImplementedError (f"( {self.qm_method}.{call_name()} ) {error_message} ( {error_vars} )")

        # Set new variable to decide the position of atoms in terms of basis functions
        self.check_basis = []
        for ibasis in range(self.nbasis):
            for iat in range(molecule.nat_qm):
                ind_a = self.check_atom[iat] + 1
                ind_b = self.check_atom[iat + 1]
                if (ibasis + 1 >= ind_a and ibasis + 1 <= ind_b):
                    self.check_basis.append(iat + 1)

        # Initialize NACME variables
        # There is no core orbitals in TDDFTB (fixed occupations)
        # nocc is number of occupied orbitals and nvirt is number of virtual orbitals
        self.norb = self.nbasis
        self.nocc = int(int(molecule.nelec - core_elec) / 2)
        self.nvirt = self.norb - self.nocc

        # Replace norb by arrays containing the limits of the for loops.
        # For energy window calculations loops will not go from (0 to nocc/nvirt) or (0 to norb)
        # but from (nocc_min to nocc/0 to nvirt_max) or (nocc_min to norb).
        self.orb_ini = np.zeros(1, dtype=np.int32)
        self.orb_final = np.zeros(1, dtype=np.int32)
        self.orb_final[0] = self.norb

        if (self.e_window > eps):
            # Swap minimal/maximal values to replace them in reading of SPX.DAT by the minimal/maximal values.
            self.orb_ini[0] = self.norb
            self.orb_final[0] = 0

        self.ao_overlap = np.zeros((self.nbasis, self.nbasis))
        self.mo_coef_old = np.zeros((self.norb, self.nbasis))
        self.mo_coef_new = np.zeros((self.norb, self.nbasis))
        self.ci_coef_old = np.zeros((molecule.nst, self.nocc, self.nvirt))
        self.ci_coef_new = np.zeros((molecule.nst, self.nocc, self.nvirt))

    def get_data(self, molecule, base_dir, bo_list, dt, istep, calc_force_only):
        """ Extract energy, gradient and nonadiabatic couplings from (TD)DFTB method

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        self.copy_files(molecule, istep, calc_force_only)
        super().get_data(base_dir, calc_force_only)
        self.write_xyz(molecule)
        self.get_input(molecule, istep, bo_list, calc_force_only)
        self.run_QM(molecule, base_dir, istep, bo_list, calc_force_only)
        self.extract_QM(molecule, base_dir, istep, bo_list, dt, calc_force_only)
        self.move_dir(base_dir)

    def copy_files(self, molecule, istep, calc_force_only):
        """ Copy necessary scratch files in previous step

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Copy required files for NACME
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "geometry.xyz"), \
                os.path.join(self.scr_qm_dir, "../geometry.xyz.pre"))
            if (istep == 0):
                shutil.copy(os.path.join(self.scr_qm_dir, "eigenvec.bin"), \
                    os.path.join(self.scr_qm_dir, "../eigenvec.bin.pre"))
                shutil.copy(os.path.join(self.scr_qm_dir, "SPX.DAT"), \
                    os.path.join(self.scr_qm_dir, "../SPX.DAT.pre"))
                shutil.copy(os.path.join(self.scr_qm_dir, "XplusY.DAT"), \
                    os.path.join(self.scr_qm_dir, "../XplusY.DAT.pre"))

        # Copy required files to read initial guess
        if (self.guess == "read" and istep >= 0):
            # After T = 0.0 s
            shutil.copy(os.path.join(self.scr_qm_dir, "charges.bin"), \
                os.path.join(self.scr_qm_dir, "../charges.bin.pre"))

    def get_input(self, molecule, istep, bo_list, calc_force_only):
        """ Generate DFTB+ input files: geometry.gen, dftb_in.hsd

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Make 'geometry.gen' file
        os.system("xyz2gen geometry.xyz")
        if (self.l_periodic):
            # Substitute C to S in first line
            file_be = open('geometry.gen', 'r')
            file_af = open('tmp.gen', 'w')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm} S\n'
                    first_row = False
                file_af.write(row)
            # Add gamma-point and cell lattice information
            geom_periodic = textwrap.dedent(f"""\
            {0.0:15.8f} {0.0:15.8f} {0.0:15.8f}
            {self.a_axis[0]:15.8f} {self.a_axis[1]:15.8f} {self.a_axis[2]:15.8f}
            {self.b_axis[0]:15.8f} {self.b_axis[1]:15.8f} {self.b_axis[2]:15.8f}
            {self.c_axis[0]:15.8f} {self.c_axis[1]:15.8f} {self.c_axis[2]:15.8f}
            """)
            file_af.write(geom_periodic)
            file_be.close()
            file_af.close()
            os.rename('tmp.gen', 'geometry.gen')

        # Make 'double.gen' file for CIoverlap in TDDFTB
        # In this case, we do not need to consider periodicity
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # Move previous files to currect directory
            os.rename('../geometry.xyz.pre', './geometry.xyz.pre')
            if (istep == 0):
                os.rename('../eigenvec.bin.pre', './eigenvec.bin.pre')
                os.rename('../SPX.DAT.pre', './SPX.DAT.pre')
                os.rename('../XplusY.DAT.pre', './XplusY.DAT.pre')
            # Open 'geometry.xyz.pre'
            file_af = open('double.xyz', 'w')
            file_be = open('geometry.xyz.pre', 'r')
            first_row = True
            for row in file_be:
                if (first_row):
                    row = f'{molecule.nat_qm * 2}\n'
                    first_row = False
                file_af.write(row)
            file_be.close()
            # Open 'geometry.xyz'
            file_be = open('geometry.xyz', 'r')
            iline = 1
            for row in file_be:
                if (iline > 2):
                    file_af.write(row)
                iline += 1
            file_be.close()
            file_af.close()
            os.system("xyz2gen double.xyz")

        # Make 'dftb_in.hsd' file
        input_dftb = ""

        # Geometry Block
        input_geom = textwrap.dedent(f"""\
        Geometry = GenFormat{{
          <<< 'geometry.gen'
        }}
        """)
        input_dftb += input_geom

        # Hamiltonian Block
        input_ham_init = textwrap.dedent(f"""\
        Hamiltonian = DFTB{{
        """)
        input_dftb += input_ham_init

        # SCC-DFTB option
        if (self.l_scc):
            input_ham_scc = textwrap.indent(textwrap.dedent(f"""\
              SCC = Yes
              SCCTolerance = {self.scc_tol}
              MaxSCCIterations = {self.scc_max_iter}
              Mixer = {self.mixer}{{}}
            """), "  ")
            input_dftb += input_ham_scc

            # Onsite-corrected DFTB (OC-DFTB) option
            if (self.l_onsite):
                onsite_const_uu = ("\n" + " " * 18).join([f"  {itype}uu = {{ {onsite_uu[f'{itype}']} }}" for itype in self.atom_type])
                onsite_const_ud = ("\n" + " " * 18).join([f"  {itype}ud = {{ {onsite_ud[f'{itype}']} }}" for itype in self.atom_type])
                input_ham_oc = textwrap.indent(textwrap.dedent(f"""\
                  OnsiteCorrection = {{
                  {onsite_const_uu}
                  {onsite_const_ud}
                  }}
                """), "  ")
                input_dftb += input_ham_oc

            # Long-range corrected DFTB (LC-DFTB) option
            if (self.l_range_sep):
                input_ham_lc = textwrap.indent(textwrap.dedent(f"""\
                  RangeSeparated = LC{{
                    Screening = {self.lc_method}{{}}
                  }}
                """), "  ")
                input_dftb += input_ham_lc

            # Spin-polarized DFTB option
            if (self.l_spin_pol and molecule.nst == 1):
                input_ham_spin = textwrap.dedent(f"""\
                SpinPolarisation = Colinear{{
                  UnpairedElectrons = {self.unpaired_elec}
                }}
                """)
                input_dftb += input_ham_spin

            # Read atomic spin constants used in spin-polarized DFTB or TDDFTB
            # TODO : Currently, allows only singlet excited states with TDDFTB
#            if (self.l_spin_pol or self.ex_symmetry == "triplet"):
            if (self.l_spin_pol and molecule.nst == 1):
                if (self.l_range_sep):
                    spin_constant = ("\n" + " " * 18).join([f"  {itype} = {{ {spin_w_lc[f'{itype}']} }}" for itype in self.atom_type])
                else:
                    spin_constant = ("\n" + " " * 18).join([f"  {itype} = {{ {spin_w[f'{itype}']} }}" for itype in self.atom_type])
                input_ham_spin_w = textwrap.indent(textwrap.dedent(f"""\
                  SpinConstants = {{
                    ShellResolvedSpin = Yes
                  {spin_constant}
                  }}
                """), "  ")
                input_dftb += input_ham_spin_w

            # Read 'charges.bin' from previous step
            if (self.guess == "read"):
                if (istep == -1):
                    if (os.path.isfile(self.guess_file)):
                        # Copy guess file to currect directory
                        shutil.copy(self.guess_file, os.path.join(self.scr_qm_dir, "charges.bin"))
                        restart = "Yes"
                    else:
                        restart = "No"
                elif (istep >= 0):
                    # Move previous file to currect directory
                    os.rename("../charges.bin.pre", "./charges.bin")
                    restart = "Yes"
            elif (self.guess == "h0"):
                restart = "No"

            # Read 'charges.bin' for surface hopping dynamics when hop occurs
            if (calc_force_only):
                restart = "Yes"

            input_ham_restart = textwrap.indent(textwrap.dedent(f"""\
              ReadInitialCharges = {restart}
            """), "  ")
            input_dftb += input_ham_restart

        # TODO: for QM/MM, point_charge??

        if (self.l_periodic):
            num_k_point = np.sum(self.k_point)
            if (num_k_point == 3):
                # gamma-point sampling
                input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
                  KPointsAndWeights = {{
                    0.0 0.0 0.0 1.0
                  }}
                """), "  ")
            else:
                # K-point sampling
                shift_vector = [0.5 if (ik % 2 == 0) else 0 for ik in self.k_point]
                input_ham_periodic = textwrap.indent(textwrap.dedent(f"""\
                  KPointsAndWeights = SupercellFolding{{
                    {self.k_point[0]} 0.0 0.0
                    0.0 {self.k_point[1]} 0.0
                    0.0 0.0 {self.k_point[2]}
                    {shift_vector[0]} {shift_vector[1]} {shift_vector[2]}
                  }}
                """), "  ")
            input_dftb += input_ham_periodic

        angular_momentum = ("\n" + " " * 10).join([f"  {itype} = '{max_l[f'{itype}']}'" for itype in self.atom_type])
        input_ham_basic = textwrap.dedent(f"""\
          Charge = {molecule.charge}
          Filling = Fermi{{
            Temperature[K] = {self.elec_temp}
          }}
          MaxAngularMomentum = {{
          {angular_momentum}
          }}
          SlaterKosterFiles = Type2FileNames{{
            Prefix = '{self.sk_path}'
            Separator = '-'
            Suffix = '.skf'
            LowerCaseTypeName = No
          }}
        }}
        """)
        input_dftb += input_ham_basic

        # Analysis Block
        input_analysis = textwrap.dedent(f"""\
        Analysis = {{
          CalculateForces = Yes
          WriteBandOut = Yes
          WriteEigenvectors = Yes
          MullikenAnalysis = Yes
        }}
        """)
        input_dftb += input_analysis

        # Options Block
        input_options = textwrap.dedent(f"""\
        Options = {{
          WriteDetailedXml = Yes
          WriteDetailedOut = Yes
          TimingVerbosity = -1
        }}
        """)
        input_dftb += input_options

        # ExcitedState Block
        if (molecule.nst > 1):
            # Calculate excited state force for target state
            if (bo_list[0] > 0):
                ex_force = "Yes"
                rst = bo_list[0]
            else:
                ex_force = "No"
                rst = bo_list[0] + 1

            # Set number of excitations in TDDFTB
            # This part can be modified by users
            if (molecule.nat_qm <= 5):
                num_ex = molecule.nst + 2
            elif (molecule.nat_qm > 5 and molecule.nat_qm <= 15):
                num_ex = 2 * molecule.nst + 2
            else:
                num_ex = 3 * molecule.nst + 2

            # Write XplusY data?
            if (self.calc_coupling):
                xpy = "Yes"
            else:
                xpy = "No"

            input_excited = textwrap.dedent(f"""\
            ExcitedState = Casida{{
              NrOfExcitations = {num_ex}
              StateOfInterest = {rst}
              Symmetry = {self.ex_symmetry}
              WriteTransitions = Yes
              WriteSPTransitions = {xpy}
              WriteMulliken = Yes
              WriteXplusY = {xpy}
              EnergyWindow [eV] = {self.e_window}
              ExcitedStateForces = {ex_force}
            }}
            """)

            input_dftb += input_excited

        # ParserOptions Block
        if (self.version == "19.1"):
            parser_version = 7
        elif (self.version == "20.1"):
            parser_version = 8

        input_parseroptions = textwrap.dedent(f"""\
        ParserOptions = {{
          ParserVersion = {parser_version}
        }}
        """)
        input_dftb += input_parseroptions

        # Parallel Block
        if (self.mpi):
            if (self.l_spin_pol and self.nthreads > 1):
                groups = 2
            else:
                groups = 1
            input_parallel = textwrap.dedent(f"""\
            Parallel = {{
              Groups = {groups}
              UseOmpThreads = No
              Blacs = BlockSize {{ 32 }}
            }}
            """)
            input_dftb += input_parallel

        # Write 'dftb_in.hsd.geom' file
        file_name = f"dftb_in.hsd.geom.{bo_list[0]}"
        with open(file_name, "w") as f:
            f.write(input_dftb)

        # Write 'dftb_in.hsd.double' file
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            # New input for dftb
            input_dftb = ""

            # Geometry Block
            input_geom = textwrap.dedent(f"""\
            Geometry = GenFormat{{
              <<< 'double.gen'
            }}
            """)
            input_dftb += input_geom
            input_dftb += input_ham_init
            input_dftb += input_ham_basic

            # Options Block
            input_options = textwrap.dedent(f"""\
            Options = {{
              WriteDetailedXml = Yes
              WriteDetailedOut = Yes
              WriteHS = Yes
              TimingVerbosity = -1
            }}
            """)
            input_dftb += input_options

            file_name = "dftb_in.hsd.double"
            with open(file_name, "w") as f:
                f.write(input_dftb)

    def run_QM(self, molecule, base_dir, istep, bo_list, calc_force_only):
        """ Run (TD)DFTB calculation and save the output files to qm_log directory

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Set run command
        qm_command = os.path.join(self.qm_path, "dftb+")
        if (self.mpi):
            # MPI setting
            os.environ["OMP_NUM_THREADS"] = "1"
            mpi_command = os.path.join(self.mpi_path, "mpirun")
            command = f"{mpi_command} -np {self.nthreads} {qm_command} > log"
        else:
            # OpenMP setting
            os.environ["OMP_NUM_THREADS"] = f"{self.nthreads}"
            command = f"{qm_command} > log"

        # Run DFTB+ for calculation of overlap matrix
        if (self.calc_coupling and not calc_force_only and istep >= 0 and molecule.nst > 1):
            shutil.copy("dftb_in.hsd.double", "dftb_in.hsd")
            os.system(command)

        # Copy dftb_in.hsd for target state
        file_name = f"dftb_in.hsd.geom.{bo_list[0]}"
        shutil.copy(file_name, "dftb_in.hsd")

        # Run DFTB+ method for molecular dynamics
        os.system(command)

        # Copy detailed.out for target state
        file_name = f"detailed.out.{bo_list[0]}"
        shutil.copy("detailed.out", file_name)

        # Copy the output file to 'qm_log' directory
        tmp_dir = os.path.join(base_dir, "qm_log")
        if (os.path.exists(tmp_dir)):
            detailed_out_step = f"detailed.out.{istep + 1}.{bo_list[0]}"
            shutil.copy("detailed.out", os.path.join(tmp_dir, detailed_out_step))
            log_step = f"log.{istep + 1}.{bo_list[0]}"
            shutil.copy("log", os.path.join(tmp_dir, log_step))

    def extract_QM(self, molecule, base_dir, istep, bo_list, dt, calc_force_only):
        """ Read the output files to get BO information

            :param object molecule: Molecule object
            :param string base_dir: Base directory
            :param integer istep: Current MD step
            :param integer,list bo_list: List of BO states for BO calculation
            :param double dt: Time interval
            :param boolean calc_force_only: Logical to decide whether calculate force only
        """
        # Read 'detailed.out' file
        # TODO: the qmmm information is written in this file
        file_name = f"detailed.out.{bo_list[0]}"
        with open(file_name, "r") as f:
            detailed_out = f.read()

        # Read 'EXC.DAT' file
        if (molecule.nst > 1):
            file_name = "EXC.DAT"
            with open(file_name, "r") as f:
                exc_out = f.read()

        # Energy
        if (not calc_force_only):
            energy = re.findall('Total energy:\s+([-]\S+) H', detailed_out)
            energy = np.array(energy[0], dtype=np.float64)
            molecule.states[0].energy = energy

            if (molecule.nst > 1):
                tmp_e = f'[=]+\n' + ('\s+([-]*\S+)\s+\S+\s+\d+\s+->\s+\d+\s+\S+\s+\S+\s+[ST]') * molecule.nst
                energy = re.findall(tmp_e, exc_out)
                energy = np.array(energy[0], dtype=np.float64)
                energy *= eV_to_au
                for ist in range(1, molecule.nst):
                    molecule.states[ist].energy = molecule.states[0].energy + energy[ist - 1]

        # Force
        tmp_f = 'Total Forces' + '\n\s+\d*\s+([-]*\S+)\s+([-]*\S+)\s+([-]*\S+)' * molecule.nat_qm
        force = re.findall(tmp_f, detailed_out)
        force = np.array(force[0], dtype=np.float64)
        force = force.reshape(molecule.nat_qm, 3, order='C')
        molecule.states[bo_list[0]].force = np.copy(force)

        # NACME
        if (self.calc_coupling and not calc_force_only):
            if (istep >= 0):
                self.CI_overlap(molecule, istep, dt)

    def CI_overlap(self, molecule, istep, dt):
        """ Read the necessary files and calculate NACME from tdnac.c routine,
            note that only reading of several files is required in this method

            :param object molecule: Molecule object
            :param integer istep: Current MD step
            :param double dt: Time interval
        """
        # Read upper right block of 'oversqr.dat' file (< t | t+dt >)
        file_name_in = "oversqr.dat"

        self.ao_overlap = np.zeros((self.nbasis, self.nbasis))
        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            row = 0
            iline = 0
            for line in lines:
                # Skip first five lines and read upper block
                if (iline in range(5, 5 + self.nbasis)):
                    col = 0
                    count = False
                    field = line.split()
                    for element in field:
                        # Read right block
                        if (count):
                            ind_a = self.check_basis[row]
                            ind_b = self.check_basis[col]
                            if (ind_a == ind_b):
                                # Choose onsite (same-atom) block
                                # Sometimes NaN or too large values appear in the onsite block due to the slater-koster file
                                # The values set to 1 or 0 regardless of original elements
                                if (row == col):
                                    # Diagonal element in onsite block
                                    new_val = 1.
                                else:
                                    # Off-diagonal element in onsite block
                                    new_val = 0.
                            else:
                                # Choose offsite (different-atom) block
                                new_val = float(element)
                            # Set overlap matrix element
                            self.ao_overlap[row, col] = new_val
                        col += 1
                        # Read right block
                        if (col > self.nbasis - 1):
                            col -= self.nbasis
                            count = True
                    row += 1
                iline += 1
#        np.savetxt("test-over", self.ao_overlap, fmt=f"%6.3f")

        # Read 'eigenvec.bin.pre' file at time t
        if (istep == 0):
            file_name_in = "eigenvec.bin.pre"

            self.mo_coef_old = np.zeros((self.norb, self.nbasis))
            with open(file_name_in, "rb") as f_in:
                dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                for iorb in range(self.norb):
                    dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                    data = np.fromfile(f_in, dtype=np.float64, count=self.nbasis)
                    self.mo_coef_old[iorb] = data
#            np.savetxt("test-mo1", self.mo_coef_old, fmt=f"%12.6f")

        # Read 'eigenvec.bin' file at time t + dt
        file_name_in = "eigenvec.bin"

        self.mo_coef_new = np.zeros((self.norb, self.nbasis))
        with open(file_name_in, "rb") as f_in:
            dummy = np.fromfile(f_in, dtype=np.int32, count=1)
            for iorb in range(self.norb):
                dummy = np.fromfile(f_in, dtype=np.int32, count=1)
                data = np.fromfile(f_in, dtype=np.float64, count=self.nbasis)
                self.mo_coef_new[iorb] = data
#        np.savetxt("test-mo2", self.mo_coef_new, fmt=f"%12.6f")

        # The CI coefficients are arranged in order of single-particle excitations
        # Read 'SPX.DAT.pre' file at time t
        if (istep == 0):
            file_name_in = "SPX.DAT.pre"

            with open(file_name_in, "r") as f_in:
                lines = f_in.readlines()
                # Dimension for CI coefficients (number of excitations)
                ndim_old = int(lines[-2].strip().split()[0])
                get_wij_ind_old = np.zeros((ndim_old, 2), dtype=np.int32)
                iline = 0
                for line in lines:
                    # Skip first five lines
                    if (iline in range(5, 5 + ndim_old)):
                        # Column information: 1st = index, 4th = occ(i), 6th = virt(a)
                        field = line.split()
                        # Determine new limits for for-loops
                        if (int(field[5]) > self.orb_final[0]):
                            self.orb_final[0] = int(field[5])
                        if (int(field[3]) < self.orb_ini[0] + 1):
                            self.orb_ini[0] = int(field[3]) - 1
                        get_wij_ind_old[int(field[0]) - 1] = [int(field[3]), int(field[5])]
                    iline += 1

        # Read 'SPX.DAT' file at time t + dt
        file_name_in = "SPX.DAT"

        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            # Dimension for CI coefficients (number of excitations)
            ndim = int(lines[-2].strip().split()[0])
            get_wij_ind_new = np.zeros((ndim, 2), dtype=np.int32)
            iline = 0
            for line in lines:
                # Skip first five lines
                if (iline in range(5, 5 + ndim)):
                    # Column information: 1st = index, 4th = occ(i), 6th = virt(a)
                    field = line.split()
                    if (int(field[5]) > self.orb_final[0]):
                        self.orb_final[0] = int(field[5])
                    if (int(field[3]) < self.orb_ini[0] + 1):
                        self.orb_ini[0] = int(field[3]) - 1
                    get_wij_ind_new[int(field[0]) - 1] = [int(field[3]), int(field[5])]
                iline += 1

        # Read 'XplusY.DAT.pre' file at time t
        if (istep == 0):
            file_name_in = "XplusY.DAT.pre"

            self.ci_coef_old = np.zeros((molecule.nst, self.nocc, self.nvirt))
            with open(file_name_in, "r") as f_in:
                lines = f_in.readlines()
                iline = 0
                for line in lines:
                    if (iline == 0):
                        field = line.split()
                        assert (int(field[0]) == ndim_old)
                        assert (int(field[1]) >= molecule.nst - 1)
                        # nxply is number of lines for each excited state in 'XplusY.dat'
                        nxply = int(ndim_old / 6) + 1
                        if (ndim_old % 6 != 0):
                            nxply += 1
                    else:
                        field = line.split()
                        if (iline % nxply == 1):
                            ind = 0
                            ist = int(field[0]) - 1
                            # In general, TDDFTB calculate the excited states more than molecule.nst,
                            # so we do not need to read all data for 'XplusY.DAT'
                            if (ist == molecule.nst - 1):
                                break
                        else:
                            # Currently, elements for CI coefficients for S0 state are zero (not used values)
                            for element in field:
                                ind_occ = get_wij_ind_old[ind, 0] - 1
                                ind_virt = get_wij_ind_old[ind, 1] - self.nocc - 1
                                self.ci_coef_old[ist + 1, ind_occ, ind_virt] = float(element)
                                ind += 1
                    iline += 1
#            np.savetxt("test-ci1", self.ci_coef_old[1], fmt=f"%12.6f")

        # Read 'XplusY.DAT' file at time t + dt
        file_name_in = "XplusY.DAT"

        self.ci_coef_new = np.zeros((molecule.nst, self.nocc, self.nvirt))
        with open(file_name_in, "r") as f_in:
            lines = f_in.readlines()
            iline = 0
            for line in lines:
                if (iline == 0):
                    field = line.split()
                    assert (int(field[0]) == ndim)
                    assert (int(field[1]) >= molecule.nst - 1)
                    # nxply is number of lines for each excited state in 'XplusY.dat'
                    nxply = int(ndim / 6) + 1
                    if (ndim % 6 != 0):
                        nxply += 1
                else:
                    field = line.split()
                    if (iline % nxply == 1):
                        ind = 0
                        ist = int(field[0]) - 1
                        # In general, TDDFTB calculate the excited states more than molecule.nst,
                        # so we do not need to read all data for 'XplusY.DAT'
                        if (ist == molecule.nst - 1):
                            break
                    else:
                        # Currently, elements for CI coefficients for S0 state are zero (not used values)
                        for element in field:
                            ind_occ = get_wij_ind_new[ind, 0] - 1
                            ind_virt = get_wij_ind_new[ind, 1] - self.nocc - 1
                            self.ci_coef_new[ist + 1, ind_occ, ind_virt] = float(element)
                            ind += 1
                iline += 1
#        np.savetxt("test-ci2", self.ci_coef_new[1], fmt=f"%12.6f")

        # Calculate wavefunction overlap with orbital scheme
        # Reference: J. Phys. Chem. Lett. 2015, 6, 4200-4203
        wf_overlap(self, molecule, istep, dt)


# -*- coding: UTF-8 -*-
#   \file solvate_maths.py
#   \brief This file provides the specific computations functionality for other parts of solvate.
#
#   Copyright by the Authors and individual contributors. All rights reserved.
#
#   Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
#       1) Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
#       2) Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
#       3) Neither the name of Michal Tykac nor the names of this code's contributors may be used to endorse or promote products derived from this software without specific prior written permission.
#
#   This software is provided by the copyright holder and contributors "as is" and any express or implied warranties, including, but not limitted to, the implied warranties of merchantibility and fitness for a particular purpose are disclaimed. In no event shall the copyright owner or the contributors be liable for any direct, indirect, incidental, special, exemplary, or consequential damages (including, but not limitted to, procurement of substitute goods or services, loss of use, data or profits, or business interuption) however caused and on any theory of liability, whether in contract, strict liability or tort (including negligence or otherwise) arising in any way out of the use of this software, even if advised of the possibility of such damage.
#
#   \author    Michal Tykac
#   \author    Lada Biedermannová
#   \author    Jiří Černý
#   \version   0.1.0
#   \date      DEC 2020
######################################################

######################################################
# Imports

import numpy                                          ### Maths
import solvate.solvate_log as solvate_log             ### For writing the log

######################################################
# residueToResFragmentDistance_backbone ()
def residueToResFragmentDistance_backbone ( res, resFrag, settings ):
    """
    This function computes the procrustes optimal overlay of the backbone atoms of the input residue and
    hydrated residue fragment and then proceeds to obtain the RMSD of these backbone atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : resFrag
        A list containing the positions and names of the hydrated residue frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Reduce residue to N, Ca, C and O (backbone) atoms
    resAts                                            = []
    resAtsNames                                       = []
    for rIt in range( 0, len( res[1] ) ):
        if ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "O" ) or ( res[1][rIt][0] == "N" ):
            resAts.append                             ( [ res[1][rIt][1], res[1][rIt][2], res[1][rIt][3] ] )
            resAtsNames.append                        ( res[1][rIt][0] )

    ### Sanity check
    if len( resAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all backbone atoms (C, CA, N and O) in the residue described by " + str( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( settings )

    ### Reduce hydrated residue fragment to N, Ca, C and O (backbone) atoms
    fraAts                                            = []
    for fIt in range( 0, len( resFrag ) ):
        for rIt in range ( 0, len( resAtsNames ) ):
            if resFrag[fIt][0] == resAtsNames[rIt]:
                fraAts.append                         ( [ resFrag[fIt][1], resFrag[fIt][2], resFrag[fIt][3] ] )

    ### Sanity check
    if len( fraAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all backbone atoms (C, CA, N and O) in the supplied hydrated fragment " + str ( resFrag ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( settings )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAts ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fraAts ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAts, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated residue fragment distance of " + str( rmsdDist ), settings, 5 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# residueToResFragmentDistance_rotamer ()
def residueToResFragmentDistance_rotamer ( res, resFrag, settings ):
    """
    This function computes the procrustes optimal overlay of the side-chain atoms of the input residue and
    hydrated residue fragment and then proceeds to obtain the RMSD of these side-chain atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : resFrag
        A list containing the positions and names of the hydrated residue frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Initialise variables
    resAts                                            = []
    resAtsNames                                       = []
    fraAts                                            = []

    ### Parse out the side-chain atoms for the residue
    for rIt in range( 0, len( res[1] ) ):

        # Deal with various atom names and lengths for different amino acids
        if ( ( ( res[0] == "VAL" ) or ( res[0] == "ILE" ) ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "CG1" ) ) ) or \
           ( ( ( res[0] == "THR" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "OG1" ) ) ) or \
           ( ( ( res[0] == "SER" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "OG"  ) ) ) or \
           ( ( ( res[0] == "CYS" )                        ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "SG"  ) ) ) or \
           ( ( ( res[0] == "ARG" ) or ( res[0] == "ASN" ) or ( res[0] == "ASP" ) or ( res[0] == "GLN" ) or ( res[0] == "GLU" ) or ( res[0] == "HIS" ) or ( res[0] == "LEU" ) or \
               ( res[0] == "LYS" ) or ( res[0] == "MET" ) or ( res[0] == "PHE" ) or ( res[0] == "PRO" ) or ( res[0] == "TRP" ) or ( res[0] == "TYR" ) \
                                                          ) and ( ( res[1][rIt][0] == "C" ) or ( res[1][rIt][0] == "CA" ) or ( res[1][rIt][0] == "CB" ) or ( res[1][rIt][0] == "CG"  ) ) ):

            # Save
            resAts.append                             ( [ res[1][rIt][1], res[1][rIt][2], res[1][rIt][3] ] )
            resAtsNames.append                        ( res[1][rIt][0] )

    ### Sanity check
    if len ( resAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all side-chain atoms (C, CA, CB and ?) in the supplied residue " + str ( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Reduce fragment to N, Ca, Cb and X (as above) atoms
    for rIt in range( 0, len( resAtsNames ) ):

        # For each fragment atom
        for fIt in range( 0, len( resFrag ) ):

            # Get only the same atoms
            if resAtsNames[rIt] == resFrag[fIt][0]:

                # Save
                fraAts.append                         ( [ resFrag[fIt][1], resFrag[fIt][2], resFrag[fIt][3] ] )

    ### Sanity check
    if len ( fraAts ) != 4:
        solvate_log.writeLog                          ( "!!! ERROR !!! Could not find all side-chain atoms (C, CA, CB and ?) in the supplied hydrated residue fragment " + str ( resFrag ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAts ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fraAts ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAts, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated residue fragment distance of " + str( rmsdDist ), settings, 4 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# residueToResFragmentDistance_direct ()
def residueToResFragmentDistance_direct ( res, fragInfo, settings ):
    """
    This function computes the procrustes optimal overlay of all available atoms
    and then proceeds to obtain the RMSD of these atoms.

    Parameters
    ----------
    list : res
        A list containing the positions and names of the compared residue atoms.

    list : fragInfo
        A list containing the positions and names of the hydrated frament atoms.

    solvate_globals.globalSettings : settings
        Instance of the settings class contaning all the options and values.

    Returns
    -------
    list : dist, tForm
        First element: The RMSD distance between the supplied residues's side-chain atoms
        Second element: The dictionary of procrustes transformations.

    """
    ### Initialise variables
    resAtoms                                          = []
    fragAtoms                                         = []

    ### Get residue positions in same order as fragment
    for fIt in range ( 0, len ( fragInfo ) ):

        # For each residue atom name
        for atIt in range ( 0, len ( res[1] ) ):

            # Compare the atom names
            if res[1][atIt][0] == fragInfo[fIt][0]:

                # Save atoms in same order
                resAtoms.append                       ( [ res[1][atIt][1], res[1][atIt][2], res[1][atIt][3] ] )
                fragAtoms.append                      ( [ fragInfo[fIt][1], fragInfo[fIt][2], fragInfo[fIt][3] ] )

    ### Check for missing atoms
    if len ( fragInfo ) != len ( resAtoms ):
        solvate_log.writeLog                          ( "!!! ERROR !!! Failed to match the fragment atoms of fragment " + str ( fragInfo ) + " to residue " + str( res ), settings, 0 )

        # Terminate
        solvate_log.endLog                            ( )

    ### Convert to numpy arrays
    resNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in resAtoms ] )
    fraNumpy                                          = numpy.array ( [ numpy.array ( xi ) for xi in fragAtoms ] )

    ### Compute procrustes analysis
    ( dist, z, t )                                    = procrustes ( resNumpy, fraNumpy )

    ### Get RMSD
    rmsdDist                                          = getRMSD ( resAtoms, z )

    ### Report log
    solvate_log.writeLog                              ( "Found residue to hydrated fragment distance of " + str( rmsdDist ), settings, 5 )

    ### Done
    return                                            ( rmsdDist, t )

######################################################
# procrustes ()
def procrustes ( X, Y ):
    """
    A port of MATLAB's "procrustes" function to Numpy.

    Procrustes analysis determines a linear transformation (translation,
    reflection, orthogonal rotation and scaling) of the points in Y to best
    conform them to the points in matrix X, using the sum of squared errors
    as the goodness of fit criterion.

        d, Z, [tform] = procrustes ( X, Y )

    In this particular code, scaling has been disabled (in protein comparisons,
    scale plays a role) as was the reflections ( mirror immage of a molecule is
    a different molecule ).

    Parameters
    ----------
    numpy.ndarray : X, Y
        Matrices of target and input coordinates. They must have equal
        numbers of  points (rows), but Y may have fewer dimensions
        (columns) than X.

    Returns
    -------
    float : d
        The residual sum of squared errors, normalized according to a
        measure of the scale of X, ( ( X - X.mean(0))^2 ).sum( )

    matrix : Z
        The matrix of transformed Y-values

    dictionary : tform
        A dict specifying the rotation, translation and scaling that
        maps X --> Y

    """

    ### Find the dimensions of the input matrices
    nx , mx                                           = X.shape
    ny , my                                           = Y.shape

    ### Find the column means of the input matrices
    muX                                               = X.mean(0)
    muY                                               = Y.mean(0)

    ### Find the x - mean values of the input matrices
    X0                                                = X - muX
    Y0                                                = Y - muY

    ### Sum the squares of the x - mean values of the input matrices
    ssX                                               = ( numpy.power ( X0, 2.0 ) ).sum()
    ssY                                               = ( numpy.power ( Y0, 2.0 ) ).sum()

    ### Compute the centred Frobenius norm of the input matrices (i.e. square root of sum of x - mean squared)
    normX                                             = numpy.sqrt ( ssX )
    normY                                             = numpy.sqrt ( ssY )

    ### Scale x - mean to equal (unit) norm
    X0                                               /= normX
    Y0                                               /= normY

    ### If second matrix has more rows than the first, shorten it ( this should really be error... )
    if my < mx:
        Y0                                            = numpy.concatenate ( ( Y0, numpy.zeros ( nx, mx - my ) ), 0 )

    ### Find the optimal rotation matrix of Y using least squares
    A                                                 = numpy.dot ( X0.T, Y0 )
    U,s,Vt                                            = numpy.linalg.svd ( A, full_matrices=False )
    V                                                 = Vt.T
    T                                                 = numpy.dot(V, U.T)

    ### Does the current solution use a reflection?
    have_reflection                                   = numpy.linalg.det(T) < 0

    ### If that's the case, force the other reflection (i.e. solution with no reflections)
    if have_reflection != False:
        V[:,-1]                                      *= -1
        s[-1]                                        *= -1
        T                                             = numpy.dot(V, U.T)

    ### Find the trace of singular values
    traceTA                                           = s.sum()

    ### With no scaling
    b                                                 = 1
    d                                                 = 1 + ssY/ssX - 2 * traceTA * normY / normX
    Z                                                 = normY * numpy.dot ( Y0, T ) + muX

    ### If second matrix has more rows than the first, shorten it ( this should really be error... )
    if my < mx:
        T                                             = T[:my,:]

    ### Compute the translation vector
    c                                                 = muX - b * numpy.dot ( muY, T )

    ### Save all to transformation list
    tform                                             = {'rotation':T, 'scale':b, 'translation':c}

    ### Done
    return d, Z, tform

######################################################
# getRMSD ()
def getRMSD ( coo1, coo2 ):
    """
    This is a simple function for computing RMSD given two matrices with equal dimmensions (n rows and
    exactly 3 columns) of (presumably atomic) co-ordinates.

    Parameters
    ----------
    matrix : coo1, coo2
        matrices of target and input coordinates. They must have equal
        numbers of points (rows) and equal number of dimensions (columns
        = 3)

    Returns
    -------
    float : dist
        The RMSD distance value between the two matrices.

    """
    ### Initialise sum
    sum                                               = 0.0

    ### For each row
    for iter in range ( 0, len ( coo1 ) ):
        sum                                          += ( numpy.power ( coo1[iter][0] - coo2[iter][0], 2.0 ) +
                                                          numpy.power ( coo1[iter][1] - coo2[iter][1], 2.0 ) +
                                                          numpy.power ( coo1[iter][2] - coo2[iter][2], 2.0 ) )

    ### Done
    return                                            ( numpy.sqrt ( sum  / float ( len ( coo1 ) ) ) )

######################################################
# findSmallestDistance ()
def findSmallestDistance ( distances, threshold ):
    """
    This function takes a list of distances and a threshold and returns the index of the
    smallest distance or -1 if no distances is below the threshold.

    Parameters
    ----------
    list : distances
        The list of the distances to be searched.

    float : threshold
        The threshold below which any passing distance needs to be.

    Returns
    -------
    int : index
        The index of the smallest distance or -1 if no distance is below the threshold.

    """
    ### Find lowest RMSD and check if it is larger than the threshold
    minVal                                            = numpy.Inf
    minInd                                            = -1
    indCtr                                            = 0

    ### Find the best match (lowest distance)
    for val in distances:
        if ( val[0] < minVal ) and ( val[0] < threshold ):
            minVal                                    = val[0]
            minInd                                    = indCtr
        indCtr                                       += 1

    ### Return the index
    return                                            ( minInd )

######################################################
# findPassingDistances ()
def findPassingDistances ( distances, threshold ):
    """
    This function takes a list of distances and a threshold and returns all indices of
    distances which pass the threshold, or -1 if no distances is below the threshold.

    Parameters
    ----------
    list : distances
        The list of the distances to be searched.

    float : threshold
        The threshold below which any passing distance needs to be.

    Returns
    -------
    list : indices
        The indices of all passing distances or -1 if no distance is below the threshold.

    """
    ### Initialise variables
    minInd                                            = []
    indCtr                                            = 0

    ### Find all passing distances
    for val in distances:
        if val[0] < threshold:
            minInd.append                             ( indCtr )
        indCtr                                       += 1

    ### Check for at least one
    if len( minInd ) == 0:
        minInd                                        = -1

    ### Return the index
    return                                            ( minInd )

######################################################
# findPassingDistances ()
def rotateAndTranslateAtom ( atom, transform ):
    """
    This function takes a numpy.array of atom positions (X, Y, Z) and proceeds to
    apply the transform rotation and then transform translation to this position.
    It then returns the new position.

    Warning: Note that rotation is done before translation, if this is not the
    correct order of transformations, then please do not use the function!

    Parameters
    ----------
    numpy.array : atom
        Array with three numbers signifying the current atom position.

    dictionary : transform
        Dictionary containing the "rotation" key with the optimal rotation matrix and
        a "translation" key with the optimal translation. It should be the result of
        the procrustes analysis.

    Returns
    -------
    numpy.array : res
        Array containing the new position resulting from rotating and translation the
        input position.

    """
    ### Apply rotation
    atomPos                                           = numpy.matmul ( atom, transform["rotation"] )

    ### Apply translation
    atomPos                                           = atomPos + transform["translation"]

    ### Return the index
    return                                            ( atomPos )

######################################################
# getWaterClusterCOM ()
def getWaterClusterCOM ( clusterIndex, clusteredWaters ):
    """
    This function takes the index of a cluster, the list of all water molecules and  the
    appropriate labels list and proceeds to compute the occupancy weighted COM position
    for the cluster, which it then returns.

    Parameters
    ----------
    int : clusterIndex
        The index of the cluster for which the computation is to be made.

    list : clusteredWaters
        A list of cluster labels, cluster count, water positions and molecule occupancies and B factors for
        each structure type available in the noClashWaters input variable.

    list : watersNoClash
        List containing a list of water molecules which have no clashes (as defined by the settings)
        to the various versions of the input structure. The versions are: full structure, no hydrogens,
        no waters, no ligands. If some of the versions were not requested, the list for that version will
        be empty.

    Returns
    -------
    list : COM
        This list contains the COM position weighted by the occupancy of the contributing water molecules and the
        average occupancy and B factor as well.

    """
    ### Find all member indices
    members                                           = numpy.where ( clusteredWaters[0] == clusterIndex )[0]

    ### Initialise variables
    COM                                               = numpy.zeros ( 5 )
    occSum                                            = 0.0
    bFacSum                                           = 0.0

    ### For each member, find sum of positions weighted by occupancy
    for mem in members:

        ### Get occupancy weighted COM and occupancy sum
        COM[0]                                       += ( clusteredWaters[2][mem][0] * clusteredWaters[3][mem] )
        COM[1]                                       += ( clusteredWaters[2][mem][1] * clusteredWaters[3][mem] )
        COM[2]                                       += ( clusteredWaters[2][mem][2] * clusteredWaters[3][mem] )
        occSum                                       += clusteredWaters[3][mem]
        bFacSum                                      += clusteredWaters[4][mem]

    ### Average by occupancy sum
    COM[0]                                           /= occSum
    COM[1]                                           /= occSum
    COM[2]                                           /= occSum
    COM[3]                                            = occSum / float ( len ( members ) )
    COM[4]                                            = bFacSum / float ( len ( members ) )

    ### Done
    return ( COM )


#!/usr/bin/env python3
# -*- coding:utf-8 -*-

from typing import Dict, Optional, Union

import numpy
import torch
import torch.nn.functional as F
from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.data import Vocabulary
from allennlp.models.model import Model
from allennlp.modules import FeedForward, Maxout, Seq2SeqEncoder, TextFieldEmbedder, Seq2VecEncoder
from allennlp.nn import InitializerApplicator, RegularizerApplicator, util
from allennlp.training.metrics import CategoricalAccuracy, F1Measure
from overrides import overrides

@Model.register("semeval_classifier_attention")
class SemEvalClassifierAttention(Model):
    """This ``Model`` performs text classification for SemEval 2017 task 4 subset A.
    """

    def __init__(self,
                vocab: Vocabulary,
                text_field_embedder: TextFieldEmbedder,
                embedding_dropout: float,
                encoder: Seq2SeqEncoder,
                integrator: Seq2SeqEncoder,
                integrator_dropout: float,
                output_layer: Union[FeedForward, Maxout],
                initializer: InitializerApplicator = InitializerApplicator(),
                regularizer: Optional[RegularizerApplicator] = None) -> None:
        super().__init__(vocab, regularizer)
        # We need the embeddings to convert word IDs to their vector representations
        self.embedding_dropout = torch.nn.Dropout(embedding_dropout)
        self.text_field_embedder = text_field_embedder
        self.encoder = encoder

        self.integrator = integrator
        self.integrator_dropout = torch.nn.Dropout(integrator_dropout)

        self._self_attentive_pooling_projection = torch.nn.Linear(
                self.integrator.get_output_dim(), 1)
        self.output_layer = output_layer

        # Monitor the metrics - we use accuracy, as well as prec, rec, f1 for 4 (very positive)
        self.accuracy = CategoricalAccuracy()
        self.f1_measure_positive = F1Measure(
            vocab.get_token_index("positive", "labels"))
        self.f1_measure_negative = F1Measure(
            vocab.get_token_index("negative", "labels"))
        self.f1_measure_neutral = F1Measure(
            vocab.get_token_index("neutral", "labels"))

        # We use the cross entropy loss because this is a classification task.
        # Note that PyTorch's CrossEntropyLoss combines softmax and log likelihood loss,
        # which makes it unnecessary to add a separate softmax layer.
        self.loss_function = torch.nn.CrossEntropyLoss()
        initializer(self)

    @overrides
    def forward(self,
                tokens: Dict[str, torch.Tensor],
                label: torch.Tensor = None) -> torch.Tensor:
        # In deep NLP, when sequences of tensors in different lengths are batched together,
        # shorter sequences get padded with zeros to make them equal length.
        # Masking is the process to ignore extra zeros added by padding
        text_mask = util.get_text_field_mask(tokens).float()
        # Forward pass
        embedded_text = self.text_field_embedder(tokens)
        dropped_embedded_text = self.embedding_dropout(embedded_text)
        encoded_tokens = self.encoder(dropped_embedded_text, text_mask)

        # Compute biattention. This is a special case since the inputs are the same.
        attention_logits = encoded_tokens.bmm(encoded_tokens.permute(0, 2, 1).contiguous())
        attention_weights = util.masked_softmax(attention_logits, text_mask)
        encoded_text = util.weighted_sum(encoded_tokens, attention_weights)

        # Build the input to the integrator
        integrator_input = torch.cat([encoded_tokens,
                                      encoded_tokens - encoded_text,
                                      encoded_tokens * encoded_text], 2)
        integrated_encodings = self.integrator(integrator_input, text_mask)

        # Simple Pooling layers
        max_masked_integrated_encodings = util.replace_masked_values(
                integrated_encodings, text_mask.unsqueeze(2), -1e7)
        max_pool = torch.max(max_masked_integrated_encodings, 1)[0]
        min_masked_integrated_encodings = util.replace_masked_values(
                integrated_encodings, text_mask.unsqueeze(2), +1e7)
        min_pool = torch.min(min_masked_integrated_encodings, 1)[0]
        mean_pool = torch.sum(integrated_encodings, 1) / torch.sum(text_mask, 1, keepdim=True)

        # Self-attentive pooling layer
        # Run through linear projection. Shape: (batch_size, sequence length, 1)
        # Then remove the last dimension to get the proper attention shape (batch_size, sequence length).
        self_attentive_logits = self._self_attentive_pooling_projection(
                integrated_encodings).squeeze(2)
        self_weights = util.masked_softmax(self_attentive_logits, text_mask)
        self_attentive_pool = util.weighted_sum(integrated_encodings, self_weights)

        pooled_representations = torch.cat([max_pool, min_pool, mean_pool, self_attentive_pool], 1)
        pooled_representations_dropped = self.integrator_dropout(pooled_representations)

        logits = self.output_layer(pooled_representations_dropped)

        # In AllenNLP, the output of forward() is a dictionary.
        # Your output dictionary must contain a "loss" key for your model to be trained.
        output = {"logits": logits}
        if label is not None:
            self.accuracy(logits, label)
            self.f1_measure_positive(logits, label)
            self.f1_measure_negative(logits, label)
            self.f1_measure_neutral(logits, label)
            output["loss"] = self.loss_function(logits, label)

        return output

    @overrides
    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
        _, recall_positive, f1_measure_positive = self.f1_measure_positive.get_metric(
            reset)
        _, recall_negative, f1_measure_negative = self.f1_measure_negative.get_metric(
            reset)
        _, recall_neutral, _ = self.f1_measure_neutral.get_metric(reset)
        accuracy = self.accuracy.get_metric(reset)

        avg_recall = (recall_positive + recall_negative + recall_neutral) / 3.0
        macro_avg_f1_measure = (
            f1_measure_positive + f1_measure_negative) / 2.0

        results = {
            'accuracy': accuracy,
            'avg_recall': avg_recall,
            'f1_measure': macro_avg_f1_measure
        }
        return results

    @classmethod
    def from_params(cls, vocab: Vocabulary, params: Params) -> 'SemEvalClassifierAttention':  # type: ignore
        # pylint: disable=arguments-differ
        embedder_params = params.pop("text_field_embedder")
        text_field_embedder = TextFieldEmbedder.from_params(vocab=vocab, params=embedder_params)
        embedding_dropout = params.pop("embedding_dropout")
        encoder = Seq2SeqEncoder.from_params(params.pop("encoder"))
        integrator = Seq2SeqEncoder.from_params(params.pop("integrator"))
        integrator_dropout = params.pop("integrator_dropout")

        output_layer_params = params.pop("output_layer")
        if "activations" in output_layer_params:
            output_layer = FeedForward.from_params(output_layer_params)
        else:
            output_layer = Maxout.from_params(output_layer_params)

        elmo = params.pop("elmo", None)
        if elmo is not None:
            elmo = Elmo.from_params(elmo)
        use_input_elmo = params.pop_bool("use_input_elmo", False)
        use_integrator_output_elmo = params.pop_bool("use_integrator_output_elmo", False)

        initializer = InitializerApplicator.from_params(params.pop('initializer', []))
        regularizer = RegularizerApplicator.from_params(params.pop('regularizer', []))
        params.assert_empty(cls.__name__)

        return cls(vocab=vocab,
                   text_field_embedder=text_field_embedder,
                   embedding_dropout=embedding_dropout,
                   encoder=encoder,
                   integrator=integrator,
                   integrator_dropout=integrator_dropout,
                   output_layer=output_layer,
                   initializer=initializer,
                   regularizer=regularizer)


#! /usr/bin/env python3
import argparse
import gym
import roboschool
import scipy.signal

import numpy as np
import tensorflow as tf
import utils as U

from tabulate import tabulate
from discrete_a2c import DiscreteA2C

parser = argparse.ArgumentParser(description='discrete advantage actor critic algorithm')

parser.add_argument(
    '--lr', default=7e-4, type=float, help='learning rate')

parser.add_argument(
    '--ent_coef', default=0., type=float, help='the coefficient of entropy')

parser.add_argument(
    '--vf_coef', default=0.5, type=float, help='the coefficient of value function')

parser.add_argument(
    '--max_grad_norm', default=0.5, type=float, help='max gradients normalize')

parser.add_argument(
    '--gamma', default=.99, type=float, help='gamma')

parser.add_argument(
    '--seed', default=0, type=int, help='RNG seed')

parser.add_argument(
    '--num_steps', default=5, type=int, help='the number of steps')

parser.add_argument(
    '--num_procs', default=32, type=int, help='the number of processes')

parser.add_argument(
    '--max_steps', default=8e6, type=int, help='max steps of training')

parser.add_argument(
    '--animate', default=False, type=bool, help='whether to animate environment')

parser.add_argument(
    '--softmax', default=True, type=bool, help='whether to use softmax to sample action')

parser.add_argument(
    '--huber', default=False, type=bool, help='whether to use huber loss')

parser.add_argument(
    '--save_network', default=False, type=bool, help='whether to save network')

parser.add_argument(
    '--load_network', default=False, type=bool, help='whether to load network')

parser.add_argument(
    '--test_alg', default=False, type=bool, help='whether to test our algorithm')

parser.add_argument(
    '--gym_id', default='CartPole-v1', type=str, help='gym id')

parser.add_argument(
    '--model_name', default='discrete_a2c', type=str, help='save or load model name')

args = parser.parse_args()

def build_multi_envs():
    def make_env(rank):
        def _thunk():
            env = gym.make(args.gym_id)
            env.seed(args.seed+rank)
            return env
        return _thunk

    U.set_global_seeds(args.seed)
    env = U.SubprocVecEnv([make_env(i) for i in range(args.num_procs)])
    return env

class PlayGym(object):
    def __init__(self, args, env, agent):
        self.args = args
        self.env = env
        self.agent = agent
        self.test_env = gym.make(self.args.gym_id)

    def play(self, max_iters=100000):
        obs = self.env.reset()
        for i in range(max_iters):
            obses, acts, rews, values, obs = self._sample_trajs(obs)
            self.agent.update(obses, acts, rews, values)
            if i % 100 == 0:
                score = self.test()
                print ("iter: {} | score: {}".format(i, score))
                self.agent.score = score

    def test(self):
        env = self.test_env
        obs = env.reset()
        score = 0
        done = False
        while not done:
            act = self.agent.get_action([obs])
            obs, rew, done, _ = env.step(act)
            score += rew
        return score

    def _sample_trajs(self, obs):
        obses, acts, rews, values, dones = [], [], [], [], []

        for step in range(self.args.num_steps):
            obses.append(obs)
            act, value = self.agent.step(obs)
            obs, rew, done, _ = self.env.step(act)
            acts.append(act)
            rews.append(rew)
            values.append(value)
            dones.append(done)

        obses = np.asarray(obses, dtype=np.float32).swapaxes(1, 0)
        acts = np.asarray(acts, dtype=np.int32).swapaxes(1, 0)
        rews = np.asarray(rews, dtype=np.float32).swapaxes(1, 0)
        values = np.asarray(values, dtype=np.float32).swapaxes(1, 0)
        dones = np.asarray(dones, dtype=np.bool).swapaxes(1, 0)
        last_values = self.agent.get_value(obs)

        for n, (rew, done, value) in enumerate(zip(rews, dones, last_values)):
            rew = rew.tolist()
            done = done.tolist()
            if done[-1] == 0:
                rew = U.discount_with_dones(rew+[value], done+[0.], self.args.gamma)[:-1]
            else:
                rew = U.discount_with_dones(rew, done, self.args.gamma)
            rews[n] = rew

        obses = np.concatenate([obs for obs in obses])
        acts = acts.flatten()
        rews = rews.flatten()
        values = values.flatten()

        return obses, acts, rews, values, obs

if __name__ == '__main__':
    graph = tf.get_default_graph()
    config = tf.ConfigProto(allow_soft_placement=True,
        intra_op_parallelism_threads=args.num_procs,
        inter_op_parallelism_threads=args.num_procs)

    session = tf.Session(graph=graph, config=config)

    # build env
    env = build_multi_envs()
    obs_dim = env.observation_space.shape[0]
    act_dim = env.action_space.n

    # build agent
    agent = DiscreteA2C(session, args, obs_dim, act_dim)

    # build player
    player = PlayGym(args, env, agent)

    # start to play :)
    session.run(tf.global_variables_initializer())
    player.play()


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time

import numpy
import torch
import torchcontrol as toco
from polymetis import GripperInterface, RobotInterface
from torchcontrol.transform import Rotation as R
from torchcontrol.transform import Transformation as T

DEFAULT_MAX_ITERS = 3

# Sampling params
DEFAULT_GP_RANGE_UPPER = [0.7, 0.1, numpy.pi / 2]
DEFAULT_GP_RANGE_LOWER = [0.4, -0.1, -numpy.pi / 2]

# Grasp params
DEFAULT_REST_POSE = ([0.5, 0.0, 0.7], [1.0, 0.0, 0.0, 0.0])
DEFAULT_PREGRASP_HEIGHT = 0.55
DEFAULT_GRASP_HEIGHT = 0.15
DEFAULT_PLANNER_DT = 0.02
DEFAULT_TIME_TO_GO_SECS = 2.0

DEFAULT_GAIN_MULTIPLIER = 1.0

CARTESIAN_SPACE_CONTROLLER = 1
JOINT_SPACE_CONTROLLER = 2

class ManipulatorSystem:
    def __init__(self, robot_kwargs={}, gripper_kwargs={}):
        ip_address = robot_kwargs.get("ip_address", "localhost")
        enforce_version = robot_kwargs.get("enforce_version", "True") == "True"
        gripper_ip_address = gripper_kwargs.get("gripper_ip_address", "localhost")
        self.arm = RobotInterface(
            ip_address=ip_address, enforce_version=enforce_version
        )
        self.gripper = GripperInterface(ip_address=gripper_ip_address)
        self._rest_pose = DEFAULT_REST_POSE
        self._pregrasp_height = DEFAULT_PREGRASP_HEIGHT
        self._grasp_height = DEFAULT_GRASP_HEIGHT
        self._planner_dt = DEFAULT_PLANNER_DT
        self._gp_range_upper = torch.Tensor(DEFAULT_GP_RANGE_UPPER)
        self._gp_range_lower = torch.Tensor(DEFAULT_GP_RANGE_LOWER)
        self._up_time_seconds = DEFAULT_TIME_TO_GO_SECS
        self._down_time_seconds = DEFAULT_TIME_TO_GO_SECS
        self._gain_multiplier = DEFAULT_GAIN_MULTIPLIER
        self._controller_type = JOINT_SPACE_CONTROLLER
        time.sleep(0.5)

        # Set continuous control policy
        self.reset_policy()

        # Reset to rest pose
        self.rest_pos = torch.Tensor(self._rest_pose[0])
        self.rest_quat = torch.Tensor(self._rest_pose[1])
        self.reset()

    def __del__(self):
        self.arm.terminate_current_policy()

    def reset(self, time_to_go=DEFAULT_TIME_TO_GO_SECS):
        self.move_to(self.rest_pos, self.rest_quat, time_to_go=time_to_go)
        self.open_gripper()

    def reset_policy(self):
        # Go home
        self.arm.go_home()

        # Send PD controller
        if self._controller_type == CARTESIAN_SPACE_CONTROLLER:
            joint_pos_current = self.arm.get_joint_positions()
            policy = toco.policies.CartesianImpedanceControl(
                joint_pos_current=joint_pos_current,
                Kp=(self._gain_multiplier * torch.Tensor(self.arm.metadata.default_Kx)),
                Kd=(
                    self._gain_multiplier * torch.Tensor(self.arm.metadata.default_Kxd)
                ),
                robot_model=self.arm.robot_model,
            )
            self.arm.send_torch_policy(policy, blocking=False)
        elif self._controller_type == JOINT_SPACE_CONTROLLER:
            joint_pos_current = self.arm.get_joint_positions()
            policy = toco.policies.JointImpedanceControl(
                joint_pos_current=joint_pos_current,
                Kp=(self._gain_multiplier * torch.Tensor(self.arm.metadata.default_Kq)),
                Kd=(
                    self._gain_multiplier * torch.Tensor(self.arm.metadata.default_Kqd)
                ),
                robot_model=self.arm.robot_model,
            )
            self.arm.send_torch_policy(policy, blocking=False)
        else:
            raise Exception(f"Invalid controller type {self._controller_type}")

    def move_to(
        self, pos, quat, gather_arm_state_func=None, time_to_go=DEFAULT_TIME_TO_GO_SECS
    ):
        """
        Attempts to move to the given position and orientation by
        planning a Cartesian trajectory (a set of min-jerk waypoints)
        and updating the current policy's target to that
        end-effector position & orientation.

        Returns (num successes, num attempts, robot_states)
        """
        # Plan trajectory
        N = int(time_to_go / self._planner_dt)

        if self._controller_type == CARTESIAN_SPACE_CONTROLLER:
            pos_curr, quat_curr = self.arm.get_ee_pose()
            waypoints = toco.planning.generate_cartesian_space_min_jerk(
                start=T.from_rot_xyz(R.from_quat(quat_curr), pos_curr),
                goal=T.from_rot_xyz(R.from_quat(quat), pos),
                time_to_go=time_to_go,
                hz=1 / self._planner_dt,
            )
        elif self._controller_type == JOINT_SPACE_CONTROLLER:
            joint_pos_current = self.arm.get_joint_positions()
            waypoints = toco.planning.generate_cartesian_target_joint_min_jerk(
                joint_pos_start=joint_pos_current,
                ee_pose_goal=T.from_rot_xyz(R.from_quat(quat), pos),
                time_to_go=time_to_go,
                hz=1 / self._planner_dt,
                robot_model=self.arm.robot_model,
            )
        else:
            raise Exception(f"Invalid controller type {self._controller_type}")

        # Execute trajectory
        t0 = time.time()
        t_target = t0
        successes = 0
        error_detected = False
        robot_states = []
        for i in range(N):
            # Update traj
            try:
                if self._controller_type == CARTESIAN_SPACE_CONTROLLER:
                    ee_pos_desired = waypoints[i]["pose"].translation()
                    ee_quat_desired = waypoints[i]["pose"].rotation().as_quat()
                    # ee_twist_desired = waypoints[i]["twist"]
                    self.arm.update_current_policy(
                        {
                            "ee_pos_desired": ee_pos_desired,
                            "ee_quat_desired": ee_quat_desired,
                            # "ee_vel_desired": ee_twist_desired[:3],
                            # "ee_rvel_desired": ee_twist_desired[3:],
                        }
                    )
                    if gather_arm_state_func:
                        observed_state = gather_arm_state_func(
                            self.arm.get_robot_state()
                        )
                        observed_state["ee_pos_desired"] = ee_pos_desired
                        observed_state["ee_quat_desired"] = ee_quat_desired
                        robot_states.append(observed_state)
                elif self._controller_type == JOINT_SPACE_CONTROLLER:
                    joint_pos_desired = waypoints[i]["position"]
                    self.arm.update_current_policy(
                        {
                            "joint_pos_desired": joint_pos_desired,
                        }
                    )
                    if gather_arm_state_func:
                        observed_state = gather_arm_state_func(
                            self.arm.get_robot_state()
                        )
                        observed_state["joint_pos_desired"] = joint_pos_desired
                        robot_states.append(observed_state)
                else:
                    raise Exception(f"Invalid controller type {self._controller_type}")
            except Exception as e:
                error_detected = True
                print(f"Error updating current policy {str(e)}")

            # Check if policy terminated due to issues
            if self.arm.get_previous_interval().end != -1 or error_detected:
                error_detected = False
                print("Interrupt detected. Reinstantiating control policy...")
                time.sleep(3)
                self.reset_policy()
                break
            else:
                successes += 1

            # Spin once
            t_target += self._planner_dt
            t_remaining = t_target - time.time()
            time.sleep(max(t_remaining, 0.0))

        # Wait for robot to stabilize
        time.sleep(0.2)

        return successes, N, robot_states

    def close_gripper(self):
        self.gripper.grasp(speed=0.1, force=1.0)
        time.sleep(1.0)

        # Check state
        state = self.gripper.get_state()
        assert state.width < state.max_width
        return state

    def open_gripper(self):
        max_width = self.gripper.get_state().max_width
        self.gripper.goto(width=max_width, speed=0.1, force=1.0)
        time.sleep(0.5)

        # Check state
        state = self.gripper.get_state()
        assert state.width > 0.0
        return state

    def grasp_pose_to_pos_quat(self, grasp_pose, z):
        x, y, rz = grasp_pose
        pos = torch.Tensor([x, y, z])
        quat = (
            R.from_rotvec(torch.Tensor([0, 0, rz])) * R.from_quat(self.rest_quat)
        ).as_quat()

        return pos, quat

    def grasp(
        self,
        grasp_pose0,
        grasp_pose1,
        gather_arm_state_func=None,
        gather_gripper_state_func=None,
        reset_at_end=True,
    ):
        results = []
        traj_state = {}

        traj_state["grasp_pose"] = grasp_pose0

        # Move to pregrasp
        pos, quat = self.grasp_pose_to_pos_quat(grasp_pose0, self._pregrasp_height)
        successes, N, robot_states = self.move_to(
            pos, quat, gather_arm_state_func, time_to_go=self._up_time_seconds
        )
        results.append((successes, N))
        traj_state["move_to_grasp_states"] = robot_states

        # Lower (slower than other motions to prevent sudden collisions)
        pos, quat = self.grasp_pose_to_pos_quat(grasp_pose0, self._grasp_height)
        successes, N, robot_states = self.move_to(
            pos, quat, gather_arm_state_func, time_to_go=self._down_time_seconds
        )
        results.append((successes, N))
        traj_state["lower_to_grasp_states"] = robot_states

        # Grasp
        close_gripper_state = self.close_gripper()
        if gather_gripper_state_func:
            traj_state["close_gripper_state"] = gather_gripper_state_func(
                close_gripper_state
            )

        # Lift to pregrasp
        pos, quat = self.grasp_pose_to_pos_quat(grasp_pose0, self._pregrasp_height)
        successes, N, robot_states = self.move_to(
            pos, quat, gather_arm_state_func, time_to_go=self._up_time_seconds
        )
        results.append((successes, N))
        traj_state["lift_to_release_states"] = robot_states

        # Move to new pregrasp
        pos, quat = self.grasp_pose_to_pos_quat(grasp_pose1, self._pregrasp_height)
        successes, N, robot_states = self.move_to(
            pos, quat, gather_arm_state_func, time_to_go=self._up_time_seconds
        )
        results.append((successes, N))
        traj_state["move_to_release_states"] = robot_states

        # Read gripper state prior to release, in case object had fallen down
        if gather_gripper_state_func:
            preopen_gripper_state = self.gripper.get_state()
            traj_state["preopen_gripper_state"] = gather_gripper_state_func(
                preopen_gripper_state
            )

        # Release
        open_gripper_state = self.open_gripper()
        if gather_gripper_state_func:
            traj_state["open_gripper_state"] = gather_gripper_state_func(
                open_gripper_state
            )

        traj_state["release_pose"] = grasp_pose1

        # Reset
        if reset_at_end:
            self.reset()

        total_successes = sum([r[0] for r in results])
        total_tries = sum([r[1] for r in results])
        return total_successes, total_tries, traj_state

    def continuously_grasp(self, max_iters=1000):
        # Perform grasping
        i = 0
        total_successes, total_tries = 0, 0
        try:
            while True:
                # Sample grasp
                grasp_pose0 = uniform_sample(self._gp_range_lower, self._gp_range_upper)
                grasp_pose1 = uniform_sample(self._gp_range_lower, self._gp_range_upper)

                # Perform grasp
                print(f"Grasp {i + 1}: grasp={grasp_pose0}, release={grasp_pose1}")
                n_successes, n_tries, _ = self.grasp(grasp_pose0, grasp_pose1)
                total_successes += n_successes
                total_tries += n_tries

                # Loop termination
                i += 1
                if max_iters > 0 and i >= max_iters:
                    break

        except KeyboardInterrupt:
            print("Interrupted by user.")

        return total_successes, total_tries

    @property
    def pregrasp_height(self) -> float:
        return self._pregrasp_height

    @pregrasp_height.setter
    def pregrasp_height(self, pregrasp_height: float):
        self._pregrasp_height = pregrasp_height

    @property
    def grasp_height(self) -> float:
        return self._grasp_height

    @grasp_height.setter
    def grasp_height(self, grasp_height: float):
        self._grasp_height = grasp_height

    @property
    def planner_dt(self) -> float:
        return self._planner_dt

    @planner_dt.setter
    def planner_dt(self, planner_dt: float):
        self._planner_dt = planner_dt

    @property
    def gp_range_upper(self) -> torch.Tensor:
        return self._gp_range_upper

    @gp_range_upper.setter
    def gp_range_upper(self, gp_range_upper: torch.Tensor):
        self._gp_range_upper = gp_range_upper

    @property
    def gp_range_lower(self) -> torch.Tensor:
        return self._gp_range_lower

    @gp_range_lower.setter
    def gp_range_lower(self, gp_range_lower: torch.Tensor):
        self._gp_range_lower = gp_range_lower

    @property
    def up_time_seconds(self) -> float:
        return self._up_time_seconds

    @up_time_seconds.setter
    def up_time_seconds(self, up_time_seconds: float):
        self._up_time_seconds = up_time_seconds

    @property
    def down_time_seconds(self) -> float:
        return self._down_time_seconds

    @down_time_seconds.setter
    def down_time_seconds(self, down_time_seconds: float):
        self._down_time_seconds = down_time_seconds

    @property
    def gain_multiplier(self) -> float:
        return self._gain_multiplier

    @gain_multiplier.setter
    def gain_multiplier(self, gain_multiplier: float):
        self._gain_multiplier = gain_multiplier

    @property
    def controller_type(self) -> int:
        return self._controller_type

    @controller_type.setter
    def controller_type(self, type: int):
        self._controller_type = type

def uniform_sample(lower, upper):
    return lower + (upper - lower) * torch.rand_like(lower)


import torchvision
import torchvision.models as models

import abc
from abc import ABC
from typing import List, Tuple, Dict, Set, Union
import numpy as np
import torch.nn as nn
import torch
import torchnlp
import torch.nn.functional as F
import fasttext
from ...utils.sample import SampleList, Sample
from torchnlp.word_to_vector import CharNGram
from torchnlp.word_to_vector import BPEmb

from ...training import get_auc_dice_loss
from ...utils import init_fc, GaussianNoise, stack_and_pad_tensors, get_torchvision_classification_models, get_image_info_fn, LambdaLayer, get_device, \
    dict2sampleList, clean_memory, get_vgg_face_model, load_stored_params, get_regularization_layers
from ..classifiers import TransformerEnsembleFeaturizer
from ..text_models import Fasttext1DCNNModel, LangFeaturesModel
from ..external.detr import get_detr_model

class MultiImageMultiTextAttentionEarlyFusionModel(nn.Module):
    def __init__(self, image_models, num_classes,
                 text_models,
                 gaussian_noise, dropout,
                 internal_dims, classifier_dims,
                 final_layer_builder,
                 n_tokens_out, n_layers,
                 **kwargs):
        super(MultiImageMultiTextAttentionEarlyFusionModel, self).__init__()
        assert type(image_models) == list
        names, im_models, im_shapes, im_procs = [], [], [], []
        for imo in image_models:
            if type(imo) == dict:
                module_gaussian = imo["gaussian_noise"] if "gaussian_noise" in imo else 0.0
                module_dropout = imo["dropout"] if "dropout" in imo else 0.0
                large_rf = imo["large_rf"] if "large_rf" in imo else True
                imo = imo["model"]
            elif type(imo) == str:
                module_gaussian = 0.0
                module_dropout = 0.0
                large_rf = True
            else:
                raise NotImplementedError()

            if "torchvision" in imo:
                net = "_".join(imo.split("_")[1:])
                im_model, im_shape = get_torchvision_classification_models(net, large_rf)
                im_model = LambdaLayer(im_model, module_gaussian, module_dropout)
                self.torchvision_pool = nn.AdaptiveAvgPool2d(1)

                def global_view(x):
                    xv = self.torchvision_pool(x).expand(-1, -1, im_shape[1], -1)
                    return torch.cat([x, xv], 3)

                im_shape = list(im_shape)
                im_shape[-1] = im_shape[-1] + 1

                im_proc = LambdaLayer(global_view)

            elif imo == "caption_features":
                im_model = LambdaLayer(get_image_info_fn(enable_encoder_feats=True)["get_batch_encoder_feats"], module_gaussian, module_dropout)
                im_shape = (512, 100)
                im_proc = nn.Identity()

            elif imo == "faster_rcnn":
                im_shape = (2048, 100)
                im_model = LambdaLayer(get_image_info_fn(enable_encoder_feats=False)["get_batch_img_roi_features"], module_gaussian, module_dropout)
                im_proc = nn.Identity()

            elif imo == "lxmert_faster_rcnn":
                im_shape = (2048, 36)
                im_model = LambdaLayer(get_image_info_fn(enable_encoder_feats=False)["get_batch_lxmert_roi_features"], module_gaussian, module_dropout)
                im_proc = nn.Identity()
            elif "detr" in imo:
                im_shape = (256, 100)
                im_model = LambdaLayer(get_detr_model(get_device(), imo)["batch_detr_fn"], module_gaussian, module_dropout)
                im_proc = nn.Identity()
            elif "vgg_face" in imo:
                im_shape = (256, 1)
                im_model = LambdaLayer(get_vgg_face_model(), module_gaussian, module_dropout)
                im_proc = nn.Identity()
            else:
                raise NotImplementedError(imo)

            names.append(imo)
            im_models.append(im_model)
            im_shapes.append(im_shape)
            im_procs.append(im_proc)
        self.im_models = nn.ModuleDict(dict(zip(names, im_models)))
        self.post_procs = nn.ModuleDict(dict(zip(names, im_procs)))
        self.im_shapes = dict(zip(names, im_shapes))
        self.require_raw_img = {"detr_demo", 'detr_resnet50', 'detr_resnet50_panoptic', 'detr_resnet101', 'detr_resnet101_panoptic',
                                "ssd", "faster_rcnn", "lxmert_faster_rcnn", "caption_features"}

        assert type(text_models) == list
        tx_models, tx_names, tx_methods = [], [], []
        for i, tm in enumerate(text_models):
            assert type(tm) == dict
            text_model_class, text_model_params, text_in_channels, text_in_tokens = tm["cls"], tm["params"], tm["in_channels"], tm["in_tokens"]
            text_fwd = tm["forward"] if "forward" in tm else "__call__"
            assert text_fwd in ["__call__", "get_word_vectors"]
            text_model = text_model_class(**text_model_params)
            if text_fwd != "__call__":
                if hasattr(text_model, "featurizer"):
                    text_model.featurizer = None
                    del text_model.featurizer
                if hasattr(text_model, "final_layer"):
                    text_model.final_layer = None
                    del text_model.final_layer
            module_gaussian = tm["gaussian_noise"] if "gaussian_noise" in tm else 0.0
            module_dropout = tm["dropout"] if "dropout" in tm else 0.0
            tx_models.append(text_model)
            self.post_procs["tx_" + str(i)] = nn.Sequential(nn.Dropout(module_dropout), GaussianNoise(module_gaussian))
            tx_names.append("tx_" + str(i))
            tx_methods.append(text_fwd)

        self.tx_models = nn.ModuleDict(dict(zip(tx_names, tx_models)))
        self.text_models = dict(zip(tx_names, text_models))
        self.tx_methods = dict(zip(tx_names, tx_methods))

        ensemble_conf = {k: dict(is2d=len(v) == 3, n_tokens_in=(v[-1] * v[-2]) if len(v) == 3 else v[-1], n_channels_in=v[0]) for k, v in self.im_shapes.items()}
        text_ensemble_conf = {k: dict(is2d=False, n_tokens_in=v["in_tokens"], n_channels_in=v["in_channels"]) for k, v in
             self.text_models.items()}
        ensemble_conf.update(text_ensemble_conf)
        # ensemble_conf = text_ensemble_conf
        n_encoders = kwargs.pop("n_encoders", n_layers)
        n_decoders = kwargs.pop("n_decoders", n_layers)
        self.attention_drop_proba = kwargs["attention_drop_proba"] if "attention_drop_proba" in kwargs else 0.0
        self.featurizer = TransformerEnsembleFeaturizer(ensemble_conf, n_tokens_out, classifier_dims, internal_dims,
                                                        n_encoders, n_decoders, gaussian_noise, dropout, self.attention_drop_proba)

        self.final_layer = final_layer_builder(classifier_dims, n_tokens_out, num_classes, dropout, **kwargs)
        if "stored_model" in kwargs:
            load_stored_params(self, kwargs["stored_model"])
        self.reg_layers = get_regularization_layers(self)
        self.auc_loss_coef = kwargs.pop("auc_loss_coef", 0.0)
        self.dice_loss_coef = kwargs.pop("dice_loss_coef", 0.0)
        self.auc_method = kwargs.pop("auc_method", 1)
        self.auc_dice_loss = get_auc_dice_loss(num_classes, self.dice_loss_coef, self.auc_loss_coef, auc_method=self.auc_method)

    def get_vectors(self, sampleList: SampleList):
        sampleList = dict2sampleList(sampleList, device=get_device())
        img = sampleList.torchvision_image
        image = sampleList.image
        sample_weights = sampleList.sample_weight

        vectors = dict()
        for k, m in self.tx_models.items():
            r = getattr(m, self.tx_methods[k])(sampleList if self.tx_methods[k] == "__call__" else sampleList.text)
            text_repr = r[2] if self.tx_methods[k] == "__call__" else r
            text_repr = self.post_procs[k](text_repr)
            vectors[k] = text_repr.to(get_device())
            clean_memory()

        del sampleList

        if self.require_raw_img.isdisjoint(set(self.im_models.keys())):
            del image

        if len(set(self.im_models.keys()) - self.require_raw_img) > 0:
            img = img.to(get_device())
        else:
            del img

        clean_memory()
        for k, m in self.im_models.items():
            im_repr = m(image if k in self.require_raw_img else img)
            im_repr = self.post_procs[k](im_repr)
            vectors[k] = im_repr.to(get_device())
            clean_memory()

        vectors = self.featurizer(vectors)
        return vectors

    def forward(self, sampleList: SampleList):
        sampleList = dict2sampleList(sampleList, device=get_device())
        labels = torch.tensor(sampleList.label, dtype=float).to(get_device())
        vectors = self.get_vectors(sampleList)
        del sampleList
        clean_memory()
        logits, loss = self.final_layer(vectors, labels)
        if self.training:
            loss += self.auc_dice_loss(logits, labels)
        return logits, vectors.mean(1), vectors, loss