molrecog

#!/usr/bin/env python
import os
import sys
import numpy as np
import itertools as it
import argparse
import copy
import subprocess 
import json

stdBondLengths = {'cc': 1.60, 'ch': 1.10, 'cn': 1.50, 'co': 1.50,
               'cs': 1.90, 'nh': 1.10, 'hh': 2.6, 'sh': 1.40,
               'sn': 1.70, 'nn': 1.50}

grad2rad = np.pi/180.
rad2grad = 180./np.pi
b2A = 0.529177249
A2b = 1./b2A

class molGraph(object):
    # Simple molecular graph class, where vertices are
    # atoms and edges are bonds 
    def __init__(self,graphinput=None):
        if graphinput == None:
            graphinput = {}
        self.__graph_dict = graphinput
    
    def add_vertex(self,vertex):
        if vertex not in self.__graph_dict:
            self.__graph_dict[vertex] = []  
            return True
        else:
            return False
    
    def add_edge(self,edge):
        (vertex1, vertex2) = (edge[0], edge[1])
        if vertex1 in self.__graph_dict:
            self.__graph_dict[vertex1].append(vertex2)
        else:
            self.__graph_dict[vertex1] = [vertex2]

    def get_vertices(self):
        return list(self.__graph_dict.keys())

    def get_edges(self):
        edges = []
        for node in self.__graph_dict:
            for neighbour in self.__graph_dict[node]:
                edges.append({node, neighbour})

        return edges

    def get_neighbours(self,vertex):
        return self.__graph_dict[vertex]

    def get_dict(self):
        return self.__graph_dict

    def depth_first_search(self,vertex,paths,discovered,curPathLength,
                           maxPathLength=4):
        """ This is a modified depth first search method that determines 
            all simple paths of a certain length. In the usual DFS one
            discovers vertices only once, while here multiple discoveries
            are allowed. The list paths is passed by reference and contains 
            the output of the function.
            Inputs: vertex (str) and maximum pathlength (int)
            Outputs: list of paths with maximum pathlengt (list of strs)    """
        curPathLength += 1   # each recursion, increase current path length 
        neighbours = self.get_neighbours(vertex)
        # only recurse if maximum pathlength is not yet reached
        if curPathLength < maxPathLength: 
            for neighbour in neighbours:
                # only recurse if the vertex has not yet been discovered
                if neighbour not in discovered:
                    # add current neighbour to the discovered list
                    discovered.append(neighbour)
                    # recurse 
                    test = self.depth_first_search(neighbour,paths,discovered,
                                            curPathLength, maxPathLength)
                # only if the bottom is actually reached will test a value
                    if test != None:
                        tmp = test[:]
                        paths.append(tmp)
                    discovered.remove(neighbour)
        elif curPathLength == maxPathLength:
            # bottom is reached 
            return discovered

class Cartesians(object):
    def __init__(self, filename):
        self.filename = filename 

    def get(self, filetype):
        return getattr(self, filetype)()

    def tm(self):
        f = open(self.filename)
        try:
            atomName = []
            atomCoord = []
            i = 1
            for line in f:
                if line[0] != "$":
                    #print(line)
                    atomName.append((line.strip().split()[3], str(i)))
                    atomCoord.append([float(line.strip().split()[i])
                                      for i in range(3)])
                    i += 1
                if "$user-defined bonds" in line:
                    break
        finally:
            f.close()
        return atomName, atomCoord 

    def xyz(self):
        f = open(self.filename)
        try:
            atomName = []
            atomCoord = []
            iLine = 1
            for line in f:
                if iLine < 3:  
                    iLine += 1
                    continue
                else:
                    atomName.append((line.strip().split()[0], str(iLine-2)))
                    atomCoord.append([float(line.strip().split()[i])
                                      for i in range(1,4)])
                    iLine += 1
        finally:
            f.close()
        return atomName, atomCoord 

    def movie_abin(self):
        def _gen_movie(stepSize):
            step = 0
            geom = {} 
            frame = {}
            time = 0
            initNames, initCoord = self.xyz()
            geom["coords"] = initCoord
            geom["atomNames"] = initNames 
            frame["time"] = time
            frame["geom"] = geom
            yield frame
            frame = {}
            with open("movie.xyz", "r") as lines:
                nrAtoms = 0
                readGeom = False
                nrLines = 0
                for line in lines:
                    if nrAtoms == 0:
                        nrAtoms = int(line.strip())

                    if "Time step:" in line:
                        tmpTime = float(line.strip().split()[-1])
                        if not(tmpTime % stepSize):
                            readGeom = True
                            nrLines += 1
                            frame["time"] = tmpTime 
                            atomName = []
                            atomCoord = []
                            continue
                    
                    if readGeom:
                        nrLines += 1
                        atomName.append((line.strip().split()[0], str(nrLines-1)))
                        atomCoord.append([float(line.strip().split()[i]) for i in range(1,4)])
                        if nrLines > nrAtoms:
                            geom = {}
                            geom["coords"] = atomCoord
                            geom["atomNames"] = atomName
                            frame["geom"] = geom
                            yield frame
                            frame = {}
                            readGeom = False
                            nrLines = 0
                        
                

        return _gen_movie

class Writer(object):
    def __init__(self, inputCoords, atomNames, outputDir):
        self.inputCoords = inputCoords
        self.outputDir = outputDir
        self.atomNames = atomNames

    def write(self, filetype):
        return getattr(self, filetype)()

    def tm(self):
        fileName = self.outputDir + '/coord' 
        subprocess.run(['touch',fileName]) 

        with open(fileName,"w") as f:
            # write first three atoms
            f.write("$coord\n")
            for i, atomCoord in enumerate(self.inputCoords):
                f.write('{x:20.14f}  {y:20.14f}  {z:20.14f}'
                        '      {n:3}\n'.format(x = atomCoord[0], y=atomCoord[1],
                                          z=atomCoord[2], n=self.atomNames[i][0]))
            f.write("$end")
    
    def xyz(self):
        fileName = self.outputDir + '/coord.xyz' 
        subprocess.run(['touch',fileName]) 

        with open(fileName,"w") as f:
            f.write(str(len(self.inputCoords)) + '\n \n')
            for i, atomCoord in enumerate(self.inputCoords):
                f.write('{n:3} {x:20.14f}  {y:20.14f}  {z:20.14f}'
                        '\n'.format(x = atomCoord[0], y=atomCoord[1],
                                    z=atomCoord[2], n=self.atomNames[i][0]))

    def gaussian(self):
        fileName = self.outputDir + '/go.com' 
        templateFile = './go.com'
        try:
            subprocess.run(['cp',templateFile, fileName])
        except:
            raise ValueError('Template file ./go.com not found') 

        with open(fileName,"a") as f:
            for i, atomCoord in enumerate(self.inputCoords):
                f.write('{n:3} {x:20.14f}  {y:20.14f}  {z:20.14f}'
                        '\n'.format(x = atomCoord[0], 
                                    y=atomCoord[1],
                                    z=atomCoord[2],
                                    n=self.atomNames[i][0]))
            f.write('\n')


    def bagel(self):
        fileName = self.outputDir + '/input.json' 
        templateFile = './input.json'
        try:
            fin = open(templateFile, 'r')
        except:
            raise ValueError('Template file ./input.json not found') 

        templateObject = json.load(fin)
        for i, atomCoord in enumerate(self.inputCoords):
            if type(atomCoord) != type([]):
                atomCoord = atomCoord.tolist()
            templateObject['bagel'][0]['geometry']\
            .append( {"atom": 
                    self.atomNames[i][0],
                 "xyz": 
                    #atomCoord.tolist()
                    atomCoord
                }
            )

        fin.close()
        fout = open(fileName, 'w')
        json.dump(templateObject, fout, indent=4)
        fout.close()

    def bagel_dyson(self): 
        input_files = [
            './input_files/' + x
            for x in os.listdir('./input_files')
            if 'QM.in' not in x
        ]
        for input_file in input_files:
            subprocess.run(['cp', '-r', input_file, self.outputDir])
        fileName = self.outputDir + '/QM.in' 
        templateFile = './input_files/QM.in'
        template = open(templateFile, 'r')    
        templateLines = template.readlines() 
        template.close()
        with open(fileName, "w") as f:
            f.write(str(len(self.inputCoords)) + '\n Dyson\n')
            for i, atomCoord in enumerate(self.inputCoords):
                f.write(
                    '{n:3} {x:20.14f}  {y:20.14f}'
                    '  {z:20.14f}'.format(
                        x=atomCoord[0], 
                        y=atomCoord[1],
                        z=atomCoord[2],
                        n=self.atomNames[i][0]
                    )
                )
                f.write(
                    '  {vx:20.14f}  {vy:20.14f}'
                    '  {vz:20.14f}\n'.format(
                        vx=0.0, 
                        vy=0.0,
                        vz=0.0
                    )
                )
            for templateLine in templateLines:
                if 'savedir' in templateLine:
                    modTemplateLine = templateLine.split()[0] + ' '
                    modTemplateLine += self.outputDir + '/temp\n'
                    f.write(modTemplateLine)
                    continue
                f.write(templateLine)
                
        #    print(self.atomNames[i][0])

    def Z_matrix(self,outputFile):
        fileName = outputFile 
        subprocess.run(['touch',fileName]) 
        with open(fileName,"w") as f:
            # write first three atoms
            f.write("# Output\n\n")
            f.write(" \n")
            f.write("\n")
            f.write("0  1\n")
            line = self.atomNames[0][0].upper() + "\n"
            f.write(line)
            line = self.atomNames[1][0].upper() + "  " + str(self.BLConnectivity[1]+1) + "  r2\n"
            f.write(line)
            line = (self.atomNames[2][0].upper() +  "  " + str(self.BLConnectivity[2]+1) +  "  r3  " +
                    str(self.BAConnectivity[2]+1) + "  a3\n")
            f.write(line)
            for i in range(3,self.nrAtoms):
                line = (self.atomNames[i][0].upper() + "  " + str(self.BLConnectivity[i]+1) + "  r" +
                str(i+1) + "  " + str(self.BAConnectivity[i]+1) + "  a" + str(i+1) + "  "
                + str(self.TDConnectivity[i]+1) + "  d" + str(i+1) +"\n") 
                f.write(line)
            f.write("Variables:\n")
            for i in range(self.nrAtoms):
                if self.bondLengths[i] != 0:
                    f.write('r{strInd:1d}= {strBL:6.4f}\n'.format(strInd = i + 1, strBL = self.bondLengths[i]))
                if self.bondAngles[i] != 0:
                    f.write('a{strInd:1d}= {strBA:6.2f}\n'.format(strInd = i + 1, strBA = self.bondAngles[i]))
                if self.torsions[i] != 0:
                    f.write('d{strInd:1d}= {strTD:6.2f}\n'.format(strInd = i + 1, strTD = self.torsions[i]))
            f.write("\n") 

class redundantInternals(molGraph):
    """ This class is inteded for conversion of Cartesian coordinates 
        to redundant internals. Basically the cartesians are read from
        some file (right now only Turbomole is supported -> change in
        future) and the relevant bond lengths are calculated first. 
        From this the molecular graph molGraph class is constucted.
        Based on this connectivity all bond angles and torsional 
        dihedrals are determined. """
    def __init__(self, fname=None, inpMolGraph=None, mollist=None,
                 filetype = None):
        if fname != None and mollist == None and inpMolGraph==None:
            self.setupGraph = True
            self.atomNames, self.atomCoords = self.getcartesians(fname, filetype)
            self.nrAtoms = len(self.atomNames)
            self.distMat = self.initDistanceMatrix()
            self.BLConnectivity, self.bondLengths = self.initBondLengths()
        elif fname != None and mollist == None and inpMolGraph!=None:
            self.setupGraph = False 
            self.atomNames, self.atomCoords = self.getcartesians(fname, filetype)
            self.nrAtoms = len(self.atomNames)
            self.distMat = self.initDistanceMatrix()
        elif fname == None and mollist != None:
            self.atomNames, self.atomCoords = mollist 
            self.nrAtoms = len(self.atomNames)
            self.distMat = self.initDistanceMatrix()

        if inpMolGraph == None:
            molGraph.__init__(self)
            for indBL, BL in enumerate(self.BLConnectivity):
                self.add_vertex(indBL)
                self.add_edge([indBL , BL])
                self.add_edge([BL , indBL])
        else:
            molGraph.__init__(self,graphinput=inpMolGraph)
            if self.atomCoords != None:
                self.BLConnectivity, self.bondLengths = self.initBondLengths()

        self.BAConnectivity, self.bondAngles = self.setBondAngles()
        self.TDConnectivity, self.torsions = self.setTorsionalDihedrals()

    def getcartesians(self,filename,filetype):
        return Cartesians(filename).get(filetype)
    
    def initDistanceMatrix(self):
        distMat = np.zeros((self.nrAtoms,self.nrAtoms))
        for i in range(self.nrAtoms):
            for j in range(i+1,self.nrAtoms):
                posi = np.array(self.atomCoords[i])
                posj = np.array(self.atomCoords[j])
                vecji = posi - posj 
                distji = np.sqrt(np.dot(vecji,vecji))
                distMat[i,j] = distji
                distMat[j,i] = distji

        return distMat
                 
    def getNumberedAtom(self, ind):
        return self.atomNames[ind][0] +  self.atomNames[ind][1]

    def initBondLengths(self):
        print("Bond lengths:")
        bondLengths = []
        distconnect = np.zeros(self.nrAtoms).astype(int)
        dist = np.zeros(self.nrAtoms)
        if self.setupGraph:
            for i in range(1,self.nrAtoms):
                indMin, distMin = (np.argmin(self.distMat[i,:i]),
                                   np.amin(self.distMat[i,:i]))
                atomInit = self.getNumberedAtom(indMin)
                atomFinal = self.getNumberedAtom(i)
                print('{strAtomi:3} - {strAtomj:4}: {strDistji:.3f} pm'.format(
                      strAtomi = atomInit, strAtomj = atomFinal,
                      strDistji = distMin*100))
                distconnect[i] = indMin 
                dist[i] = distMin
        else:
            for i in range(1,self.nrAtoms):
                bondPartner = self.get_neighbours(i)[0]
                bondLength = self.distMat[i, self.get_neighbours(i)[0]]
                distconnect[i] = bondPartner
                dist[i] = bondLength 
                atomInit = self.getNumberedAtom(bondPartner)
                atomFinal = self.getNumberedAtom(i)
                print('{strAtomi:3} - {strAtomj:4}: {strDistji:.3f} pm'.format(
                      strAtomi = atomInit, strAtomj = atomFinal,
                      strDistji = bondLength*100))
        
        return distconnect, dist


    def setBondAngles(self):
        print("Bond Angles:")
        anglesConnectivity = np.zeros(self.nrAtoms).astype(int)
        angles = np.zeros(self.nrAtoms)
        for i in range(2,self.nrAtoms):
            anglevert = [i]
            anglevert.append(self.get_neighbours(anglevert[-1])[0])
            neighbours = self.get_neighbours(anglevert[-1])
            if neighbours[0] == anglevert[-1]:
                for neighbour in neighbours:
                    if neighbour not in anglevert and neighbour < i:
                        anglevert.append(neighbour)
                        break
            else:
                anglevert.append(neighbours[0])
            anglesConnectivity[i] = anglevert[2]
            atomi = self.getNumberedAtom(anglevert[0]) 
            atomj = self.getNumberedAtom(anglevert[1])
            atomk = self.getNumberedAtom(anglevert[2])
            posi = np.array(self.atomCoords[anglevert[0]])
            posj = np.array(self.atomCoords[anglevert[1]])
            posk = np.array(self.atomCoords[anglevert[2]])
            vecji = posi - posj 
            vecjk = posk - posj 
            normji = np.sqrt(np.dot(vecji,vecji))
            normjk = np.sqrt(np.dot(vecjk,vecjk))
            dotijk =  np.dot(vecji,vecjk)
            angleijk = np.arccos(dotijk/(normji*normjk))*rad2grad
            angles[i] = angleijk
            print('{strAtomi:3} -  {strAtomj:3} -  {strAtomk:4}: {strAngleijk:7.3f} deg'.format(
                  strAtomi = atomi, strAtomj = atomj, strAtomk = atomk,
                  strAngleijk =  angleijk))
        return anglesConnectivity, angles 

    def setTorsionalDihedrals(self):
        torsConnectivity = np.zeros(self.nrAtoms).astype(int)
        tors = np.zeros(self.nrAtoms)
        dihedralVertices = []
        print("Proper dihedrals:")
        for i in self.get_vertices()[3:]:
            torsvert = [i]
            torsvert.append(self.get_neighbours(torsvert[-1])[0])
            torsvert.append(self.BAConnectivity[i])
            nextneighbours = self.get_neighbours(self.BAConnectivity[i]) 
            if nextneighbours[0] in torsvert:
                if len(nextneighbours) == 1:
                    newneighbours = self.get_neighbours(nextneighbours[0])
                    for neighbour in newneighbours:
                        if neighbour not in torsvert and neighbour < i:
                            torsvert.append(neighbour)
                            break
                else:
                    for neighbour in nextneighbours:
                        if neighbour not in torsvert and neighbour < i:
                            torsvert.append(neighbour)
                            break
            else:
                torsvert.append(nextneighbours[0]) 
                        
            atomi = self.getNumberedAtom(torsvert[0])
            atomj = self.getNumberedAtom(torsvert[1])
            atomk = self.getNumberedAtom(torsvert[2])
            #print(torsvert)
            atoml = self.getNumberedAtom(torsvert[3])
            posi = np.array(self.atomCoords[torsvert[0]])
            posj = np.array(self.atomCoords[torsvert[1]])
            posk = np.array(self.atomCoords[torsvert[2]])
            posl = np.array(self.atomCoords[torsvert[3]])
            vecij = posj - posi 
            veckj = posk - posj 
            veclk = posl - posk 
            nrmlijk = np.cross(vecij,veckj)
            nrmljkl = np.cross(veckj,veclk)
            ccijkl = np.dot(veckj,np.cross(nrmlijk,nrmljkl))
            ddijkl = np.linalg.norm(veckj)*np.dot(nrmlijk,nrmljkl)
            torsi = np.arctan2(ccijkl,ddijkl)*rad2grad
            #if torsi < 0:
            #    torsi += 360
            tors[i] = torsi 
            torsConnectivity[i] = torsvert[-1] 
            print('{strAtomi:3} -  {strAtomj:3} -  {strAtomk:3}'
                  ' -  {strAtoml:4}: {strDihedralijkl:7.3f} deg'.format(
                  strAtomi = atomi, strAtomj = atomj, strAtomk = atomk,
                  strAtoml = atoml, strDihedralijkl = torsi))

        return torsConnectivity, tors

    def write(self, outputDir, filetype, inputCoords=None):
        if inputCoords == None:
            Writer(self.atomCoords, self.atomNames, outputDir).write(filetype)
        else:
            Writer(inputCoords, self.atomNames, outputDir).write(filetype)


    def _setup_3Drotation_matrix(self,rotationAxis,angle):
        angle = angle * grad2rad
        antisymmetricTerm = np.array([[              0., -rotationAxis[2],  rotationAxis[1]],
                                      [ rotationAxis[2],               0., -rotationAxis[0]],
                                      [-rotationAxis[1],  rotationAxis[0],               0.]])
        M = (np.cos(angle) * np.eye(3) + np.outer(rotationAxis, rotationAxis)
             * (1. - np.cos(angle)) + antisymmetricTerm * np.sin(angle))
        return M
        

    def transform_2_xyz(self, overWrite=False):
        """
        Implementation of the natural extension reference frame (NeRF) method
        first presented by Parsons et. al J. Comput. Chem. 26; 1063-1068; 2005
        The idea is to assume that a triad of atoms, A, B, C are located on the
        xy plane such that C lies at the origin, B on the negative x axis and 
        A in the xy plane. Then a new atom D is placed in R3 accoriding to sph-
        erical coordinates, defined via its internals with respect to the triad.
        Since the triad A B C can be situated anywhere in R3 in the last step we
        have to transform the vector D to the triad reference frame.   
        """

        positions = []
        pos0 = np.array([0., 0., 0.])
        positions.append(pos0)
        pos1 = np.array([self.bondLengths[1], 0., 0.])
        positions.append(pos1)
        if self.BLConnectivity[2] == 0:
            rotMat = self._setup_3Drotation_matrix([0.,0.,1.], self.bondAngles[2])
            pos2 = np.matmul(rotMat,np.array([self.bondLengths[2],0., 0.])) 
            positions.append(pos2)
        elif self.BLConnectivity[2] == 1:
            rotMat = self._setup_3Drotation_matrix([0.,0.,1.], self.bondAngles[2])
            pos2 = np.matmul(rotMat,np.array([-self.bondLengths[2],0., 0.])) 
            pos2[0] = pos2[0] + pos1[0]
            positions.append(pos2)
         
        for i in range(3,self.nrAtoms):
            A = positions[self.TDConnectivity[i]]
            B = positions[self.BAConnectivity[i]]
            C = positions[self.BLConnectivity[i]]
            R = self.bondLengths[i]
            theta = self.bondAngles[i]*grad2rad
            phi = self.torsions[i]*grad2rad
            # -R*np.cos(theta) because the x element of the vector in the above
            # paper is not R*np.cos(theta) but R*np.cos(pi-theta)=-R*np.cos(theta)
            Dref = np.array([-R*np.cos(theta), R*np.cos(phi)*np.sin(theta),
                          R*np.sin(phi)*np.sin(theta)])
            unitbc = (C-B)/np.linalg.norm(C-B)
            vecAB = B-A
            nrmlabc = np.cross(vecAB,unitbc)
            nrmlabc = nrmlabc/np.linalg.norm(nrmlabc)
            molx = unitbc
            moly = np.cross(nrmlabc,unitbc) 
            moly = moly/np.linalg.norm(moly)
            molz = nrmlabc 
            molFrame = np.array([molx,moly,molz]).T
            D = np.matmul(molFrame,Dref) + C 
            if len(positions) > 0:
                # Hack to prevent atoms from bumping into each other: change sign of torsional dihedral
                # Not sure if this is any good... 
                if len([D for i in range(len(positions)) if (x := np.linalg.norm(positions[i]-D)) < 1.]) != 0:
                    Dref = np.array([-R*np.cos(theta), -R*np.cos(phi)*np.sin(theta),
                                     -R*np.sin(phi)*np.sin(theta)])
                    D = np.matmul(molFrame,Dref) + C
            positions.append(D)

        if overWrite:
            self.atomCoords = positions
        else:
            return positions

        
    def _mkdirs(self, step, CWD, other=None, second=None):
        if other != None: 
            initialDir = CWD + "step0"  
            finalDir   = CWD + "step" + str(step)  
            subprocess.run(['mkdir',initialDir]) 
            subprocess.run(['mkdir',  finalDir]) 
            return initialDir, finalDir
        elif second != None:
            newDir = CWD + "step" + str(step) + "_" + str(second)
            subprocess.run(['mkdir',  newDir]) 
            return newDir
        else:
            newDir = CWD + "step" + str(step)
            subprocess.run(['mkdir',  newDir]) 
            return newDir
        

    def interpolate_in_internals(self, other, nrSteps, CWD, filetype):
        diffBL = []
        diffBA = []
        diffTD = []
        if len(self.bondLengths) == len(other.bondLengths):
            for i in range(len(self.bondLengths)):
                diffBL.append(other.bondLengths[i] - self.bondLengths[i])
        if len(self.bondAngles) == len(other.bondAngles):
            for i in range(len(self.bondAngles)):
                diffBA.append(other.bondAngles[i] - self.bondAngles[i])
        if len(self.torsions) == len(other.torsions):
            for i in range(len(self.torsions)):
                tmpDiff = other.torsions[i] - self.torsions[i]
                if abs(tmpDiff) > 180.: 
                    if tmpDiff < 0.:
                        tmpDiff += 360 
                    elif tmpDiff > 0.:
                        tmpDiff -= 360 
                diffTD.append(tmpDiff)
        stepSizeBL = np.array(diffBL)/nrSteps
        stepSizeBA = np.array(diffBA)/nrSteps
        stepSizeTD = np.array(diffTD)/nrSteps
        initDir, finDir = self._mkdirs(nrSteps, CWD, other)
        self.transform_2_xyz(overWrite=True)
        other.transform_2_xyz(overWrite=True)
        self.write(initDir, filetype)
        other.write(finDir, filetype)
        copyBLs = copy.deepcopy(self.bondLengths)
        copyBAs = copy.deepcopy(self.bondAngles)
        copyTDs = copy.deepcopy(self.torsions)
        for i in range(1,nrSteps):
            for j in range(len(self.bondLengths)):
                self.bondLengths[j] = self.bondLengths[j] + stepSizeBL[j]
            for j in range(len(self.bondAngles)):
                self.bondAngles[j] = self.bondAngles[j] + stepSizeBA[j]
            for j in range(len(self.torsions)):
                self.torsions[j] = self.torsions[j] + stepSizeTD[j]
            tmpCoords = self.transform_2_xyz()
            newDir = self._mkdirs(i, CWD)
            self.write(newDir, filetype, inputCoords = tmpCoords)

    def incrementInternal(self, internalType, internalIndex, stepSize):
        if internalType == "bl":
            self.bondLengths[internalIndex] = self.bondLengths[internalIndex] + stepSize
        if internalType == "ba":
            self.bondAngles[internalIndex] = self.bondAngles[internalIndex] + stepSize
        if internalType == "td":
            self.torsions[internalIndex] = self.torsions[internalIndex] + stepSize
            #print(self.torsions[internalIndex])
            #print(internalIndex)
        print(self.torsions)

    def getInternal(self, internalType, internalIndex):
        if internalType == "bl":
            outInternal = self.bondLengths[internalIndex]
        if internalType == "ba":
            outInternal = self.bondAngles[internalIndex]
        if internalType == "td":
            outInternal = self.torsions[internalIndex]
        return outInternal

    def getInternals(self):
        return (copy.deepcopy(self.bondLengths),
                copy.deepcopy(self.bondAngles),
                copy.deepcopy(self.torsions))

    def setInternal(self, internalType, internalIndex, inpInternal):
        if internalType == "bl":
            self.bondLengths[internalIndex] = inpInternal
        if internalType == "ba":
            self.bondAngles[internalIndex] = inpInternal
        if internalType == "td":
            self.torsions[internalIndex] = inpInternal
    
    def setInternals(self, inpInternals):
        self.bondLenghts = inpInternals[0]
        self.bondAngles  = inpInternals[1]
        self.torsions    = inpInternals[2]
    
    def rigidscan1D(self, internalType, atomsInvolved, stepSize, nrSteps, CWD,
                    filetype):
        tmpAtoms = atomsInvolved.split("-")[:]
        if internalType != 'bl': 
            internalIndex = self.atomNames.index(tuple(tmpAtoms[0]))
        else:
            internalIndex = self.atomNames.index(tuple(tmpAtoms[-1]))
        self.transform_2_xyz(overWrite=True)
        initDir = self._mkdirs(0, CWD)
        self.write(initDir, filetype)
        for i in range(1,nrSteps+1):
            self.incrementInternal(internalType, internalIndex, stepSize)
            tmpCoords = self.transform_2_xyz()
            newDir = self._mkdirs(i, CWD)
            self.write(newDir, filetype, inputCoords = tmpCoords)

    def atomTuple(self, string):
        for i in range(len(string)):
            try:
                int(string[i:])
                break
            except:
                pass
        return (string[:i], string[i:])

    def rigidscan2D(self, internalType, atomsInvolved, stepSize1, stepSize2,
                    nrSteps1, nrSteps2, CWD, filetype):
        tmpAtoms1 = atomsInvolved[0].split("-")[:]
        tmpAtoms2 = atomsInvolved[1].split("-")[:]
        internalType1 = internalType[0] 
        internalType2 = internalType[1] 
        if internalType != 'bl': 
            internalIndex1 = self.atomNames.index(self.atomTuple(tmpAtoms1[0]))
            internalIndex2 = self.atomNames.index(self.atomTuple(tmpAtoms2[0]))
        else:
            internalIndex1 = self.atomNames.index(self.atomTuple(tmpAtoms1[-1]))
            internalIndex2 = self.atomNames.index(self.atomTuple(tmpAtoms2[-1]))
        initDir = self._mkdirs(0, CWD, second=0)
        self.transform_2_xyz(overWrite=True)
        self.write(initDir, filetype)
        for i in range(0,nrSteps1+1):
            copyInternals2 = self.getInternals() 
            for j in range(0,nrSteps2+1):
                if (i == 0) and ( j == 0):
                    continue
                self.incrementInternal(internalType2, internalIndex2, stepSize2)
                tmpCoords2 = self.transform_2_xyz()
                newDir = self._mkdirs(i, CWD, second=j)
                self.write(newDir, filetype, inputCoords = tmpCoords2)
            self.setInternals(copyInternals2)
            self.incrementInternal(internalType1, internalIndex1, stepSize1)

    def __sub__(self, other):
        meanDiffBL  = 0.
        meanDiffBA  = 0.
        meanDiffTD  = 0.
        meanDiffTot = 0.
        totNrInternals = (len(self.bondLengths) + len(self.bondAngles)
                         + len(self.torsions) - 6)
        if len(self.bondLengths) == len(other.bondLengths):
            for i in np.arange(len(self.bondLengths)): 
                meanDiffBL += np.abs(self.bondLengths[i] 
                                   - other.bondLengths[i])
            meanDiffTot += meanDiffBL/totNrInternals
            meanDiffBL = meanDiffBL/(len(self.bondLengths)-1)
            print("Bond length deviation")
            print(meanDiffBL*100)
        if len(self.bondAngles) == len(other.bondAngles):
            for i in np.arange(len(self.bondAngles)): 
                meanDiffBA += np.abs(self.bondAngles[i] 
                                   - other.bondAngles[i])
            meanDiffTot += meanDiffBA/totNrInternals
            meanDiffBA = meanDiffBA/(len(self.bondAngles)-2)
            print("Bond angle deviation")
            print(meanDiffBA)
        if len(self.torsions) == len(other.torsions):
            for i in np.arange(len(self.torsions)): 
                tmpDiff = np.abs(self.torsions[i] 
                                  - other.torsions[i])
                if tmpDiff > 180.:
                    tmpDiff -= 360 
                meanDiffTD += tmpDiff 
            meanDiffTot += meanDiffTD/totNrInternals
            meanDiffTD = meanDiffTD/(len(self.torsions)-2)
            print("Torsional dihedral deviation")
            print(meanDiffTD)
        
        print("Total internals distance") 
        print(meanDiffTot)

class Movie():
    def __init__(self, fname=None, code=None, init=False, initFile=None):
        if (fname != None) and (code != None):
            self.fname = fname
            self.code  = code 
        else:
            raise ValueError("fname and code variables can't be none!")
        self._movie = Cartesians(self.fname).get("movie_" + self.code)

    def sample(self, stepSize, CWD, filetype): 
        step = 0
        for frame in self._movie(stepSize):
            tmpDir = CWD + "step" + str(step)  
            step += 1
            subprocess.run(['mkdir',tmpDir]) 
            Writer(frame["geom"]["coords"], 
                   frame["geom"]["atomNames"],
                   tmpDir).write(filetype)

def main():
    CWD = os.getcwd() + "/"
    parser = argparse.ArgumentParser()
    parser.add_argument("-sh","--show",nargs=2)
    parser.add_argument("-c","--compare",nargs=3)
    parser.add_argument("-cp","--convert",nargs=4)
    parser.add_argument("-ls1d","--linscan1d",nargs=7)
    parser.add_argument("-ls2d","--linscan2d",nargs=11)
    parser.add_argument("-liic","--lininterp",nargs=5)
    parser.add_argument("-post","--postprocess",nargs=4)
    parser.add_argument("-b","--bohr",action='store_true')
    args = parser.parse_args()

    if args.show:
        tmpFileName = CWD + args.show[0]
        geom = redundantInternals(tmpFileName,
                                  filetype=args.show[1]) 

    if args.compare:
        tmpFileName1 = CWD + args.compare[0]
        tmpFileName2 = CWD + args.compare[1]
        geom1 = redundantInternals(tmpFileName1, 
                                   filetype=args.compare[2])
        gGraph1 = geom1.get_dict()
        geom2 = redundantInternals(tmpFileName2,
                                   inpMolGraph=gGraph1,
                                   filetype=args.compare[2])
        diff = geom1 - geom2
    
    if args.convert: 
        tmpFileName1 = CWD + args.convert[0]
        atomNames, atomCoords = Cartesians(tmpFileName1).get(args.convert[1])
        Writer(atomCoords, atomNames, args.convert[3]).write(args.convert[2]) 
        
    if args.linscan1d:
        tmpFileName   = CWD + args.linscan1d[0]
        filetype = args.linscan1d[1]
        internalType  = args.linscan1d[2]
        atomsInvolved = args.linscan1d[3]
        stepSize      = float(args.linscan1d[4])
        nrSteps       = int(args.linscan1d[5])
        geom = redundantInternals(tmpFileName, filetype=filetype) 
        geom.rigidscan1D(internalType, atomsInvolved, 
                         stepSize, nrSteps, CWD, args.linscan1d[6])
        
    if args.linscan2d:
        tmpFileName   = CWD + args.linscan2d[0]
        filetype      = args.linscan2d[1]
        internalType  = args.linscan2d[2:4]
        atomsInvolved = args.linscan2d[4:6]
        stepSize1     = float(args.linscan2d[6])
        stepSize2     = float(args.linscan2d[7])
        nrSteps1      = int(args.linscan2d[8])
        nrSteps2      = int(args.linscan2d[9])
        geom = redundantInternals(tmpFileName, filetype=filetype) 
        geom.rigidscan2D(internalType, atomsInvolved, stepSize1, stepSize2,
                         nrSteps1, nrSteps2, CWD, args.linscan2d[10])

    if args.lininterp:
        tmpFileName1 = CWD + args.lininterp[0]
        tmpFileName2 = CWD + args.lininterp[1]
        filetype = args.lininterp[2]
        geom1 = redundantInternals(fname=tmpFileName1, filetype=filetype) 
        gGraph1 = geom1.get_dict()
        geom2 = redundantInternals(fname=tmpFileName2, inpMolGraph=gGraph1,
                                   filetype=filetype)
        geom1.interpolate_in_internals(geom2,int(args.lininterp[3]),CWD,args.lininterp[4])

    if args.postprocess:
        tmpFileName = CWD + args.postprocess[0]
        code = args.postprocess[1]
        stepSize = float(args.postprocess[2])
        movie = Movie(fname=tmpFileName, code=code) 
        movie.sample(stepSize, CWD, args.postprocess[3])

if __name__ == '__main__':
    main()