# This is a simplied version to get the local mode force constant 
# based on the formatted checkpoint file generated by Gaussian.

import math
import numpy.linalg 
import numpy
 

def func(path,atom1,atom2):
  infile= path

  def LT2Sqr(N,T):
      "Transform a lower triangular vector into a square matrix"
      S=[]
      for i in range(N):
          S.append( [] )
      for i in range(N):
          for j in range(N):
              S[i].append(0)


      k = 0 -1 # corrected for python
      for i in range(N):
          if i == 1:
              k = k + 1
              S[i][0] = T[k]
              S[0][i] = T[k]
              #print S
          if i != 1:
           for j in range(i):
              k = k + 1
              S[i][j] = T[k]
              S[j][i] = T[k]
              #print S
          k = k + 1
          S[i][i] = T[k]
          #print S

      return S



  ncar=0
  ncarline=0
  carcount=-1
  carreads=[]
  coord1=[]
  coord2=[]
  nat3=0

  with open(infile) as f:
      for line in f:
          if "Current cartesian coord" in line:
            ncar = int(line.split()[-1])
            #print ncar
            ncarline = int(math.ceil(ncar/5.0))
            carcount = 0
            continue
          if carcount > -1 and carcount < ncarline:
            carcount = carcount + 1 
            #print line
            carreadline =  line.split()
            for ii  in range(len(carreadline)):
                carreads.append( float( carreadline[ii])) 
            continue
  f.close()
  nat3=len(carreads)


  for i in range(3):
      coord1.append( carreads[3*(atom1-1)+i] )
      coord2.append( carreads[3*(atom2-1)+i] )


  bvec=[]
  for i in range(nat3):
     bvec.append( 0 ) 

  r = ((coord1[0]-coord2[0])**2 + (coord1[1]-coord2[1])**2 + (coord1[2]-coord2[2])**2)**0.5
  bvec[ 3*(atom1-1)+0 ] = (coord1[0]-coord2[0]) / r
  bvec[ 3*(atom1-1)+1 ] = (coord1[1]-coord2[1]) / r
  bvec[ 3*(atom1-1)+2 ] = (coord1[2]-coord2[2]) / r

  bvec[ 3*(atom2-1)+0 ] = -(coord1[0]-coord2[0]) / r 
  bvec[ 3*(atom2-1)+1 ] = -(coord1[1]-coord2[1]) / r
  bvec[ 3*(atom2-1)+2 ] = -(coord1[2]-coord2[2]) / r


  #print bvec

  ############################################
  nhess=0
  nhessline=0
  hesscount=-1
  hessreads=[]
  with open(infile) as f:
      for line in f:
          if "Cartesian Force Constants" in line:
            nhess  = int(line.split()[-1])
            nhessline = int(math.ceil(nhess/5.0))
            hesscount = 0
            continue
          if hesscount > -1 and hesscount < nhessline:
            hesscount = hesscount + 1 
            hessreadline = line.split()
            for ii in range(len(hessreadline)):
                hessreads.append( float( hessreadline[ii]))
            continue
  f.close()

  #print hessreads

             #print line 


  SHess = LT2Sqr(nat3,hessreads)
  #print SHess

  SHessInv = numpy.linalg.inv( SHess )
  #print SHessInv 
  b = numpy.array( bvec )
  tmp1 = SHessInv.dot(b)
  tmp2 = numpy.dot( tmp1, b)

  return str(round(1/tmp2*15.56893,3))

def get_ids(path):
    f=open(path,'r')
    lines=f.readlines()
    f.close()
    ref=0
    st=[]
    for line in lines:
        if ref==5:
            break
        if 'File format' in line:
            ref=1
        if len(line.strip().split())==0 and ref>0:
            ref+=1
        if ref==2 or ref==3:
            ref+=1
        if ref==4:
            s,e=line.index('('),line.index(')')
            l=line[s+1:e].split(',')
            st.append([int(l[0].strip()),int(l[1].strip())])
    return st

def make_lmode(txt_file):
  ids=get_ids(txt_file)
  g=open(txt_file.split('.')[0]+'.out','w')
  g.write(""" ======================================================================
 ==        Conversion of Normal into Local Vibrational Modes:        ==
 ==                        Program LOCALMODES                        ==
 ==                                                                  ==
 ==                     Ver. 1.3.0 (2016/11/08)                      ==
 ==                                                                  ==
 ==                   Wenli Zou and Dieter Cremer                    ==
 ==    Computational and Theoretical Chemistry Group (CATCO), SMU    ==
 ==                     Dallas, Texas 75275 USA                      ==
 ==               Webpage: sites.smu.edu/dedman/catco/               ==
 ======================================================================

 Contributions to this or earlier versions by Zoran Konkoli, Elfi Kraka,
 Andreas Larsson, Dmitry Izotov, Robert Kalescky, Yunwen Tao, Chengfei
 Wang are acknowledged.
Local mode properties:
 -----------------------------------------------------------------------------------------------
   No.   i   j   k   l    Name                   q_n       k^^a    k^^a(d)       m^^a       w^^a
 -----------------------------------------------------------------------------------------------
 """)
  no=1
  for i,j in ids:
    print i,j
    try:
      ka=func('1s4a-qm4.fchk',i,j)
      g.write('   '+str(no)+'    '+str(i)+'  '+str(j)+'   0   0    -             -      '+ka+'  -  -\n')
    except IndexError:
      g.write('   '+str(no)+'    '+str(i)+'  '+str(j)+'   0   0    -             -      '+'ER'+'  -  -\n')
    no+=1
  g.write("-----------------------------------------------------------------------------------------------\n\n")
  g.close()

make_lmode('1s4a-qm4.txt')
#print func('1s4a-qm4.fchk',2,31)






 
