#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright (c) 2012, tamanegi (tamanegi@users.sourceforge.jp)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

import sys
import re
import os

filedirname = os.path.dirname( os.path.abspath( __file__ ) )
commonpath = os.path.join( filedirname, "../common" )
if not commonpath in sys.path:
  sys.path.append( commonpath )
from vector3d import *
from basis import *

def skipLines( fh, num = 1 ):
  for i in range( 0, num ):
    line = fh.readline()
  return 0

class GamessOut:
  def __init__( self ):
    # following regex should be class variables?
    self.__reg_basis_def = re.compile( "SHELL TYPE  PRIMITIVE" )
    self.__reg_coordinate = re.compile( "COORDINATES \(BOHR\)" )
    self.__reg_num_basis = re.compile( "BASIS SET SHELLS" )
    self.__reg_eigen_vectors = re.compile( "EIGENVECTORS" )
    self.__reg_end_eigen_vectors = re.compile( "\.\.\." )

  # this function must return [ [ AtomicOribtal ], [ float ] ],
  # where the former is the list of AO info (including coordinate) of
  # AtomicOrbital type (see basis.py) and the latter is the list of
  # coefficients for each atomic orbital
  #
  # filename: gamess out file name
  # mo_index: index of MO
  def read( self, filename, mo_index ):
    self.__orbitals = []
    self.__positions = []
    self.__eigenvecs = []
    try:
      fh = open( filename, "r" )
      self.__read( fh, mo_index )
      ## debug: atom names in bases
      #for hoge in self.__orbitals:
      #  print hoge.atomname
      ## debug: atom names in coordinate
      #for hoge in self.__positions:
      #  print hoge[0]
    except IOError, ( errno, msg ):
      if len( filename ) == 0:
        print >> sys.stderr, "Error: filename is empty."
      else:
        print >> sys.stderr, "Error: cannot open %s for reading." % filename
      print >> sys.stderr, "Error: errno = [%d] msg = [%s]" % ( errno, msg )
      sys.exit(1)
    atoms = []
    # non-C1 case?; i do not know when this case occurs
    if len( self.__orbitals ) != len( self.__positions ):
      for pos in self.__positions:
        atom = AtomicOrbital()
        # i should use __orbitals as local variable ...
        atom.setOrbitals( self.__searchOrbital( pos[0] ) )
        atom.setPosition( pos[1] )
        atoms.append( atom )
      # this must not happen
      if len( atoms ) != len( self.__positions ):
        print >> sys.stderr, "Error: number of atoms mismatch."
        print >> sys.stderr, "Error: # from AO = %d # from pos = %d." % ( len( atoms ), len( self.__positions ) )
        sys.exit(2)
    else:
      for i in range( 0, len( self.__positions ) ):
        atom = AtomicOrbital()
        atom.setOrbitals( self.__orbitals[i] )
        atom.setPosition( self.__positions[i][1] )
        atoms.append( atom )
    # free useless objects
    self.__orbitals = []
    self.__positions = []
    return [ atoms, self.__eigenvecs ]

  # note: self.__orbitals are used
  def __searchOrbital( self, atomname ):
    for orb in self.__orbitals:
      if orb.atomname == atomname:
        return orb
    print sys.stderr, "Error: cannot find corresponding orbital for atom %s" % atomname
    sys.exit(7)

  def __read( self, fh, mo_index ):
    skipLines( fh, 1 )
    line = fh.readline()
    while line:
      if self.__reg_basis_def.search( line ):
        print >> sys.stderr, "Info: read basis"
        self.__readBases( fh )
      elif self.__reg_coordinate.search( line ):
        print >> sys.stderr, "Info: read coordinates"
        self.__readCoordinates( fh )
      elif self.__reg_eigen_vectors.search( line ):
        print >> sys.stderr, "Info: read eigenvectors"
        self.__readEigenVectors( fh, mo_index )
      line = fh.readline()

  def __readEigenVectors( self, fh, mo_index ):
    self.__eigenvecs = []
    if mo_index <= 0:
      print >> sys.stderr, "Error: index of MO must be large than 0."
      sys.exit(3)
    skipSets = int( ( mo_index - 1 ) / 5 )
    column = ( mo_index - 1 ) % 5
    curSet = 1
    skipLines( fh, 1 )
    line = fh.readline()
    #print >> sys.stderr, "Info: skipSets =", skipSets, "column =", column
    while line:
      skipLines( fh, 3 )
      line = fh.readline()
      while( len( line ) > 20 and not self.__reg_end_eigen_vectors.search( line )):
        sline = line[15:].split()
        if skipSets < curSet:
          self.__eigenvecs.append( sline[column] )
        line = fh.readline()
      if skipSets < curSet or self.__reg_end_eigen_vectors.search( line ):
        break
      curSet += 1
    if len( self.__eigenvecs ) == 0:
      print >> sys.stderr, "Error: error reading eigenvectors."
      sys.exit(5)

  def __readCoordinates( self, fh ):
    self.__positions = []
    skipLines( fh, 1 )
    line = fh.readline()
    while line:
      sline = line.split()
      if len(sline) < 5:
        break
      pos = Vector3D( float( sline[2] ), float( sline[3] ), float( sline[4] ) )
      self.__positions.append( [ sline[0], pos ] )
      line = fh.readline()

  def __readBases( self, fh ):
    self.__orbitals = []
    line = fh.readline()
    myao = AtomicOrbital()
    while line:
      sline = line.split()
      if self.__reg_num_basis.search( line ):
        break
      if len(sline) == 1:
        # maybe atomname
        if not myao.empty():
          self.__orbitals.append( myao )
          myao = AtomicOrbital()
        # set atomname
        myao.atomname = sline[0]
        skipLines( fh, 1 )
      elif len(sline) != 0:
        # maybe list of gaussian basis
        temporalgaussians = []
        while len( sline ) > 0:
          shelltype = sline[0]
          orbital_type = sline[1]
          primitive_num = sline[2]
          exp = sline[3]
          conts = sline[4:]
          temporalgaussians.append( [ exp, orbital_type, conts ] )
          line = fh.readline()
          sline = line.split()
        orbs = self.createOrbitals( temporalgaussians )
        for orbital in orbs:
          myao.addOrbital( orbital )
      line = fh.readline()
    if not myao.empty():
      self.__orbitals.append( myao )

  def createOrbitals( self, basis_data ):
    # list of orbitals; s, px, py, pz, dxx, dyy, dzz, dxy, dxz, dyz
    container = [ [], [], [], [], [], [], [], [], [], [] ]
    for a_basis in basis_data:
      if a_basis[1] == "L": # split valence
        if len( a_basis[2] ) < 2:
          print >> sys.stderr, "Error: L type orbital must have two contraction coeffs."
          sys.exit(1)
        # S
        container[0].append( Basis( a_basis[0], a_basis[2][0], 0, 0, 0 ) )
        # P
        container[1].append( Basis( a_basis[0], a_basis[2][1], 1, 0, 0 ) )
        container[2].append( Basis( a_basis[0], a_basis[2][1], 0, 1, 0 ) )
        container[3].append( Basis( a_basis[0], a_basis[2][1], 0, 0, 1 ) )
      else: # S, P, D
        if a_basis[1] == "S":
          container[0].append( Basis( a_basis[0], a_basis[2][0], 0, 0, 0 ) )
        elif a_basis[1] == "P":
          container[1].append( Basis( a_basis[0], a_basis[2][0], 1, 0, 0 ) )
          container[2].append( Basis( a_basis[0], a_basis[2][0], 0, 1, 0 ) )
          container[3].append( Basis( a_basis[0], a_basis[2][0], 0, 0, 1 ) )
        elif a_basis[1] == "D":
          container[4].append( Basis( a_basis[0], a_basis[2][0], 2, 0, 0 ) )
          container[5].append( Basis( a_basis[0], a_basis[2][0], 0, 2, 0 ) )
          container[6].append( Basis( a_basis[0], a_basis[2][0], 0, 0, 2 ) )
          container[7].append( Basis( a_basis[0], a_basis[2][0], 1, 1, 0 ) )
          container[8].append( Basis( a_basis[0], a_basis[2][0], 1, 0, 1 ) )
          container[9].append( Basis( a_basis[0], a_basis[2][0], 0, 1, 1 ) )
    ret = []
    for component in container:
      if len(component) == 0:
        continue
      orbital = Orbital()
      for gaussian in component:
        orbital.addGaussian( gaussian )
      ret.append( orbital )
    return ret
