def read_molecule(r):
    """ (reader) -> Molecule

    Read a single molecule from r and return it,
    or return None to signal end of file.
    """

    # If there isn't another line, we're at the end of the file.
    line = r.readline()
    if not line:
        return None

    # Name of the molecule: "COMPND   name"
    key, name = line.split()

    # Other lines are either "END" or "ATOM num kind x y z"
    molecule = Molecule(name)
    reading = True

    while reading:
        line = r.readline()
        if line.startswith('END'):
            reading = False
        else:
            key, num, kind, x, y, z = line.split()
            molecule.add(Atom(num, kind, float(x), float(y), float(z)))

    return molecule

def parse_lines():
    #f = open(filename,'r')
    isParsing = True
    mol = None
    line_counter = 0
    n_atoms = 0
    mol = Molecule("")
    title = ""
    to_angs = False
    datas = []
    for line in fileinput.input():
        data = line.split()
        if isParsing:
            line_counter += 1
            if line_counter == 1: n_atoms = int(data[0])
            if line_counter == 2:
                uline = line.upper()
                to_angs = "AU" in uline
            if line_counter > 2:
                toangs = 1.0
                if to_angs: toangs = 0.529177249
                char = data[0]
                data = map(float, data[1:])
                atom = Atom(char, data[0]*toangs, data[1]*toangs, data[2]*toangs)
                if mol is not None:
                    mol.addAtom(atom)
    #f.close()
    return mol

class LeastSquaresCharges(LeastSquaresBasic):

    def __init__(self, data):
        self.molecule = Molecule(atoms=data['atoms'])
        self.grid = Grid(data)
        self.setA()
        reference = self.grid.get_properties(property_name=data['property'],\
                theory='%s_%s' % (data['theory'], data['basis']))
        LeastSquaresBasic.__init__(self, A=self.A, b=reference)

    def setA(self):
        n_p = len(self.grid.points)
        n_s = len(self.molecule.sites_noneq)
        self.A =  np.zeros((n_p, n_s))
        for i in xrange(n_p):
            proton = Site(coordinates=self.grid.points[i].coordinates, name='H+',index=1)
            for j, name in enumerate(self.molecule.sites_names_noneq):
                for site in self.molecule.get_sites(name=name):
                    self.A[i,j] += 1/site.distance_to(proton)

    def setA_fast(self):
        grid_coordinates = self.grid.get_coordinates()
        sites_coordinates = self.molecule.get_coordinates()
        self.A = fast.set_inversed(grid_coordinates, sites_coordinates, \
                len(self.molecule.sites_names_eq), self.molecule.sym_sites)

    @property
    def charges(self):
        charges = {}
        for q, name in zip(self.solution, self.molecule.sites_names):
            charges[name] = q
        return charges

def createMolecules(arr,params):
    retarr=[]
    for elem in arr:
        mol=Molecule()
        mol.RDMol=elem["RDMol"]
        retarr.append(mol)
    return retarr

 def test_unit_conversion_symmetry(self):
     """ 
     Does converting back and forth between bohrs and angstroms introduce and compound rounding errors? Each 
     operation should exactly reverse its counterpart.
     """
             
     with open('../../extra-files/molecule.xyz', 'r') as file1:
         mol1 = Molecule(file1.read())
     mol2 = mol1.copy()
     for count in range(10000) :
         mol2.to_bohr()
         mol2.to_angstrom()
     self.assertEqual(mol1.geom, mol2.geom)

    def test_unit_conversion_accuracy(self):
        """ 
        1.0 Angstrom is approximately 1.889725989 Bohr. Is this conversion (and its reverse) carried out correctly?
        """

        with open('../../extra-files/molecule.xyz', 'r') as file1:
            mol1 = Molecule(file1.read())
        mol2 = mol1.copy()
        mol2.to_bohr()
        for i in range(mol1.natom) :
            self.assertAlmostEqual(mol1.geom[i][0] * 1.889725989, mol2.geom[i][0])
            self.assertAlmostEqual(mol1.geom[i][1] * 1.889725989, mol2.geom[i][1])
            self.assertAlmostEqual(mol1.geom[i][2] * 1.889725989, mol2.geom[i][2])

 def read_cartesian(self,fname):
     molc=Molecule()
     f = open(fname,'r')
     while (True):
         str = f.readline()
         if (str==''):
             break
         sp = str.split()
         if (len(sp) != 4):
             continue
         molc.add_atom(Atom(sp[0].lower(),0,float(sp[1]),float(sp[2]),float(sp[3])))
     self.mol = molc
     return  molc

def solve(mol_path, hess_path):

    with open(mol_path, 'r') as f:
        molecule = Molecule(f.read())
        molecule.to_angstrom()

    with open(hess_path, 'r') as f:
        str = (f.read()).replace("\n",";")
        while str[-1] == ';' :
            str = str[:-1]
        mat = matrix(str)

    output_frequencies(molecule, mat)

 def __init__(self, data):
     self.molecule = Molecule(atoms=data['atoms'])
     self.grid = Grid(data)
     self.setA()
     reference = self.grid.get_properties(property_name=data['property'],\
             theory='%s_%s' % (data['theory'], data['basis']))
     LeastSquaresBasic.__init__(self, A=self.A, b=reference)

 def __init__(self, name=None, ff=None, sym=False, g_settings=None, m_settings=None, multipoles=None, ls_jobs=['w', 'non-w']):
     GM_logger.info("Creating MoleculeOnGrid instance")
     Molecule.__init__(self, name=name, ff=ff, sym=sym, settings=m_settings, multipoles=multipoles)
     self.grid = grid.Grid(molecule_name=name, settings=g_settings)
     self.grid_coordinates = self.grid.get_coordinates()
     reference = {}
     reference['non-w'] = self.grid.get_properties(property_name='esp', theory='reference')
     try:
         reference['w']= reference['non-w']*np.sqrt(self.grid.weights)
     except AttributeError:
         pass
     self.A = {}
     self.set_A(ls_jobs)
     self.LS = {}
     for key, A in self.A.items():
         self.LS[key] = LeastSquares(name=key, A=A, b=reference[key])

    def test_random_molecule(self):
        m = Molecule.random_molecule(12, 16)

        assert m.__class__ == Molecule

        assert m.r.x <= 12
        assert m.r.y <= 16
        assert m.v.x <= 3
        assert m.v.y <= 4

    def test_copy_function(self):
        """ 
        Does the copy function correctly initialize all variables of a new molecule? Also, is the new molecule
        truly a different object? (ie: changing properties of the original does not affect the copy and vice versa)?
        """
                
        with open('../../extra-files/molecule.xyz', 'r') as file1:
            mol1 = Molecule(file1.read())
        mol2 = mol1.copy()
        self.assertEqual(mol1.units, mol2.units, 'checking units')
        self.assertEqual(mol1.natom, mol2.natom, 'checking natom')
        self.assertEqual(mol1.labels, mol2.labels, 'checking labels')
        self.assertEqual(mol1.masses, mol2.masses, 'checking masses')
        self.assertEqual(mol1.charges, mol2.charges, 'checking charges')
        self.assertEqual(mol1.geom, mol2.geom, 'checking geometry')

        mol2.to_bohr()
        self.assertNotEqual(mol1.units, mol2.units)
        self.assertNotEqual(mol1.geom, mol2.geom)

    def test_random_molecules(self):
        ms = Molecule.random_molecules(100, 12, 16)

        assert len(ms) == 100
        for m in ms:
            assert m.__class__ == Molecule

            assert m.r.x <= 12
            assert m.r.y <= 16
            assert m.v.x <= 3
            assert m.v.y <= 4

def filterByActivity(arr,params):
    retarr=[]
    discard=[0,0,0,0,0,0,0,0]
    mols={}
    for line in arr:
        #print(line)
        if not line[u'bioactivity_type'] == u'IC50':
            discard[0]+=1
            continue
        if line[u"target_confidence"]<7:
            discard[1]+=1
            continue
        if line[u"units"]!=u"nM":
            discard[2]+=1
            continue
        value=0.0
        try:
            value=float(line[u"value"])
        except ValueError:
            discard[3]+=1
            continue
        if value > 1000:
            discard[4]+=1
            continue
        mol=Molecule(line[u'parent_cmpd_chemblid'])

        #mol.id=line[u'ingredient_cmpd_chemblid']
        if(mol.id in mols):
            discard[5]+=1
            continue
        else:
            mols[mol.id]=True           
        mol.pIC=-math.log(value)
        retarr.append(mol)
        discard[6]+=1
    print(discard)
    return retarr

def harvest_zmat(zmat):
    """Parses the contents of the Cfour ZMAT file into array and
    coordinate information. The coordinate info is converted into a
    rather dinky Molecule (no fragment, but does read charge, mult,
    unit). Return qcdb.Molecule. Written for findif zmat* where
    geometry always Cartesian and Bohr.

    """
    zmat = zmat.splitlines()[1:]  # skip comment line
    Nat = 0
    readCoord = True
    isBohr = ''
    charge = 0
    mult = 1
    molxyz = ''
    cgeom = []
    for line in zmat:
        if line.strip() == '':
            readCoord = False
        elif readCoord:
            lline = line.split()
            molxyz += line + '\n'
            Nat += 1
        else:
            if line.find('CHARGE') > -1:
                idx = line.find('CHARGE')
                charge = line[idx + 7:]
                idxc = charge.find(',')
                if idxc > -1:
                    charge = charge[:idxc]
                charge = int(charge)
            if line.find('MULTIPLICITY') > -1:
                idx = line.find('MULTIPLICITY')
                mult = line[idx + 13:]
                idxc = mult.find(',')
                if idxc > -1:
                    mult = mult[:idxc]
                mult = int(mult)
            if line.find('UNITS=BOHR') > -1:
                isBohr = ' bohr'

    molxyz = '%d%s\n%d %d\n' % (Nat, isBohr, charge, mult) + molxyz
    mol = Molecule.init_with_xyz(molxyz, no_com=True, no_reorient=True, contentsNotFilename=True)

    return mol

def jajo2mol(jajodic):
    """Returns a Molecule from entries in dictionary *jajodic* extracted
    from JAINDX and JOBARC.

    """
    map = jajodic['MAP2ZMAT']
    elem = jajodic['ATOMCHRG']
    coord = jajodic['COORD   ']
    Nat = len(elem)

    molxyz = '%d bohr\n\n' % (Nat)
    # TODO chgmult, though not really necessary for reorientation
    for at in range(Nat):
        posn = map[at] - 1
        el = 'GH' if elem[posn] == 0 else z2el[elem[posn]]
        posn *= 3
        molxyz += '%s %21.15f %21.15f %21.15f\n' % (el, coord[posn], coord[posn + 1], coord[posn + 2])
    mol = Molecule.init_with_xyz(molxyz, no_com=True, no_reorient=True, contentsNotFilename=True)

    return mol

def harvest_GRD(grd):
    """Parses the contents *grd* of the Cfour GRD file into the gradient
    array and coordinate information. The coordinate info is converted
    into a rather dinky Molecule (no charge, multiplicity, or fragment),
    but this is these coordinates that govern the reading of molecule
    orientation by Cfour. Return qcdb.Molecule and gradient array.

    """
    grd = grd.splitlines()
    Nat = int(grd[0].split()[0])
    molxyz = '%d bohr\n\n' % (Nat)

    grad = []
    for at in range(Nat):
        mline = grd[at + 1].split()
        el = 'GH' if int(float(mline[0])) == 0 else z2el[int(float(mline[0]))]
        molxyz += '%s %16s %16s %16s\n' % (el, mline[-3], mline[-2], mline[-1])
        lline = grd[at + 1 + Nat].split()
        grad.append([float(lline[-3]), float(lline[-2]), float(lline[-1])])
    mol = Molecule.init_with_xyz(molxyz, no_com=True, no_reorient=True, contentsNotFilename=True)

    return mol, grad

    def __init__(self, file, reportInterval, simulation, xyzfile, pbc=True, pmegrid=None, tinkerpath = ''):
        """Create a TinkerReporter.

        Parameters:
         - tinkerpath (string) The
         - file (string) The file to write to
         - reportInterval (int) The interval (in time steps) at which to write frames
        """
        print "Initializing Tinker Reporter"
        self._reportInterval = reportInterval
        self._openedFile = isinstance(file, str)
        if self._openedFile:
            self._out = open(file, 'w')
        else:
            self._out = file
        self._pbc = pbc
        self._pmegrid = pmegrid
        self._tinkerpath = tinkerpath
        self._simulation = simulation

        self.xyzfile = xyzfile
        self.pdbfile = os.path.splitext(xyzfile)[0] + '.pdb'
        self.keyfile = os.path.splitext(xyzfile)[0] + '.key'
        #self.prmfile = os.path.splitext(xyzfile)[0] + '.prm'
        self.dynfile = os.path.splitext(xyzfile)[0] + '.dyn'

        print "Loading Tinker xyz file"
        if not os.path.exists(self.xyzfile):
            raise IOError('You need a Tinker .xyz')
        if not os.path.exists(self.keyfile):
            raise IOError('You need a Tinker .key file with the same base name as the .xyz file')
        #if not os.path.exists(self.prmfile):
        #    raise IOError('You need a Tinker .prm file with the same base name as the .xyz file')
        if not os.path.exists(self.pdbfile):
            raise IOError('You need a .pdb file with the same base name as the .xyz file (the same one you used to start the simulation)')
        self.M = Molecule(self.xyzfile,ftype='tinker')
        self.comm = open(self.xyzfile).readlines()[0]
        print "Done"

    def fgrad(x, indicate = False):
        """ Calculate the objective function and its derivatives. """
        # If the optimization algorithm tries to calculate twice for the same point, do nothing.
        # if x == fgrad.x0: return

        xyz = independent_vars_to_xyz(x)
        # these methods require 3d input
        xyzlist = np.array([xyz])
        my_bonds = core.bonds(xyzlist, ibonds).flatten()
        my_angles = core.angles(xyzlist, iangles).flatten()
        my_dihedrals = core.dihedrals(xyzlist, idihedrals, anchor=dihedrals).flatten()

        # Deviations of internal coordinates from ideal values.
        d1 = w1*(my_bonds - bonds)
        d2 = w2*(my_angles - angles)
        d3 = w3*(my_dihedrals - dihedrals)

        # Include an optional term if we have an anchor point.
        if xrefi != None:
            d4 = (x - xrefi).flatten() * w1 * w_xref
            fgrad.error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3)), d4]
        else:
            fgrad.error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3))]

        # The objective function contains another contribution from the Morse potential.
        fgrad.X = np.dot(fgrad.error, fgrad.error)
        d1s = np.dot(d1, d1)
        d2s = np.dot(d2, d2)
        d3s = np.dot(d3, d3)
        M = Molecule()
        M.elem = elem
        M.xyzs = [np.array(xyz)*10]
        if w_morse != 0.0:
            EMorse, GMorse = PairwiseMorse(M)
            EMorse = EMorse[0]
            GMorse = GMorse[0]
        else:
            EMorse = 0.0
            GMorse = np.zeros((n_atoms, 3), dtype=float)
        if indicate: 
            if fgrad.X0 != None:
                print ("LSq: %.4f (%+.4f) Distance: %.4f (%+.4f) Angle: %.4f (%+.4f) Dihedral: %.4f (%+.4f) Morse: % .4f (%+.4f)" % 
                       (fgrad.X, fgrad.X - fgrad.X0, d1s, d1s - fgrad.d1s0, d2s, d2s - fgrad.d2s0, d3s, d3s - fgrad.d3s0, EMorse, EMorse - fgrad.EM0)), 
            else:
                print "LSq: %.4f Distance: %.4f Angle: %.4f Dihedral: %.4f Morse: % .4f" % (fgrad.X, d1s, d2s, d3s, EMorse), 
                fgrad.X0 = fgrad.X
                fgrad.d1s0 = d1s
                fgrad.d2s0 = d2s
                fgrad.d3s0 = d3s
                fgrad.EM0 = EMorse
        fgrad.X += w_morse*EMorse

        # Derivatives of internal coordinates w/r.t. Cartesian coordinates.
        d_bonds = core.bond_derivs(xyz, ibonds) * w1
        d_angles = core.angle_derivs(xyz, iangles) * w2
        d_dihedrals = core.dihedral_derivs(xyz, idihedrals) * w3

        if xrefi != None:
            # the derivatives of the internal coordinates wrt the cartesian
            # this is 2d, with shape equal to n_internal x n_cartesian
            d_internal = np.vstack([gxyz_to_independent_vars(d_bonds.reshape((len(ibonds), -1))),
                                    gxyz_to_independent_vars(d_angles.reshape((len(iangles), -1))),
                                    gxyz_to_independent_vars(d_dihedrals.reshape((len(idihedrals), -1))),
                                    np.eye(len(x)) * w1 * w_xref])
        else:
            # the derivatives of the internal coordinates wrt the cartesian
            # this is 2d, with shape equal to n_internal x n_cartesian
            d_internal = np.vstack([gxyz_to_independent_vars(d_bonds.reshape((len(ibonds), -1))),
                                    gxyz_to_independent_vars(d_angles.reshape((len(iangles), -1))),
                                    gxyz_to_independent_vars(d_dihedrals.reshape((len(idihedrals), -1)))])

        # print fgrad.error.shape, d_internal.shape
        # print d_internal.shape
        # print xyz_to_independent_vars(d_internal).shape
        fgrad.G = 2*np.dot(fgrad.error, d_internal)
        fgrad.G += xyz_to_independent_vars(w_morse*GMorse.flatten())

    def make_Hessian(self):

        self.run_disps()
        
        h, N = self.h, self.N
        E0 = self.find_E(0,0,0,0)
        self.H = np.zeros((3*self.N, 3*self.N))

        for i in range(3*N):
            for i in range(3*N):
                self.H[i,i]= (self.find_E(i,0,1,0)+self.find_E(i,0,-1,0)-2*E0)/(h**2)
                for j in range(0,i):
                    self.H[i,j] = (self.find_E(i,j,1,1)+self.find_E(i,j,-1,-1)-self.find_E(i,0,1,0)-self.find_E(j,0,1,0)-self.find_E(j,0,-1,0)-self.find_E(i,0,-1,0)+2*E0)
                    self.H[i,j] /= 2*h**2
                    self.H[j,i] = self.H[i,j]


    def write_Hessian(self):
        """
        write Hessian matrix to hessian.dat file
        """
        self.make_Hessian()
        np.savetxt("hessian.dat",self.H,"%15.7f"," ","\n")

if __name__ == "__main__":

    mol = Molecule(open("/Users/avery/git/summer-program/extra-files/molecule.xyz","r").read() )
    mol.bohr()
    hessian = Hessian(mol,"template.dat")
    hessian.write_Hessian()