#!/usr/bin/env python

import sys,os,time

wd = os.path.dirname(__file__)
sys.path += [ wd+"../EZFIO/Python" ]
sys.path += [ "/home/scemama/resultsFile" ]

from resultsFile import *

# Check command line

if len(sys.argv) == 2:
   State=0
elif len(sys.argv) == 3:
   State=int(sys.argv[2])
else:
   print "usage: "+sys.argv[0]+" file.out"
   sys.exit(2)

firstArg = sys.argv[1]

file = getFile(firstArg)
print firstArg, 'recognized as', str(file).split('.')[-1].split()[0]

from ezfio import ezfio

def write_ezfioFile(res,filename):
  ezfio.set_file(filename)

# Electrons
  ezfio.set_electrons_elec_alpha_num(res.num_alpha)
  ezfio.set_electrons_elec_beta_num(res.num_beta)

# Nuclei
  ezfio.set_nuclei_nucl_num(len(res.geometry))
  charge = []
  coord = []
  coord_x = []
  coord_y = []
  coord_z = []
  for a in res.geometry:
    charge.append(a.charge)
    if res.units == 'BOHR':
      coord_x.append(a.coord[0])
      coord_y.append(a.coord[1])
      coord_z.append(a.coord[2])
    else:
      coord_x.append(a.coord[0]/a0)
      coord_y.append(a.coord[1]/a0)
      coord_z.append(a.coord[2]/a0)
  ezfio.set_nuclei_nucl_charge(charge)
  ezfio.set_nuclei_nucl_coord(coord_x+coord_y+coord_z)

# AO Basis
  import string
  is_cartesian = True
  at = []
  num_prim = []
  magnetic_number = []
  angular_number = []
  power_x = []
  power_y = []
  power_z = []
  coefficient = []
  exponent = []
  for b in res.basis:
    if '+' in b.sym or '-' in b.sym:
      is_cartesian = False
  names = ["s","p","d","f","g","h","i","j"]
  for b in res.basis:
    c = b.center
    for i,atom in enumerate(res.geometry):
      if atom.coord == c:
         at.append(i+1)
    num_prim.append(len(b.prim))
    if is_cartesian:
      s = b.sym
      power_x.append( string.count(s,"x") ) 
      power_y.append( string.count(s,"y") )
      power_z.append( string.count(s,"z") )
    else:
      magnetic_number.append(names.index(b.sym[0]))
      angular_number.append(int(b.sym[1:]))
    coefficient.append(  b.coef )
    exponent.append( [ p.expo for p in b.prim ] )
  if not is_cartesian:
     print 'Only cartesian basis functions work...'
     sys.exit(0)
  ezfio.set_ao_basis_ao_num(len(res.basis))
  ezfio.set_ao_basis_ao_nucl(at)
  ezfio.set_ao_basis_ao_prim_num(num_prim)
  ezfio.set_ao_basis_ao_power(power_x+power_y+power_z)
  prim_num_max = ezfio.get_ao_basis_ao_prim_num_max()
  len_res_basis = len(res.basis)
  for i in range(len(res.basis)):
    coefficient[i] += [ 0. for j in range(len(coefficient[i]),prim_num_max) ]
    exponent[i]    += [ 0. for j in range(len(exponent[i]),prim_num_max) ]
  coefficient = reduce(lambda x, y: x+y, coefficient, [])
  exponent = reduce(lambda x, y: x+y, exponent, [])
  coef = []
  expo = []
  for i in range(prim_num_max):
   for j in range(i,len(coefficient),prim_num_max):
    coef.append ( coefficient[j] )
    expo.append ( exponent[j] )
  ezfio.set_ao_basis_ao_coef(coef)
  ezfio.set_ao_basis_ao_expo(expo)

# MOs
  NumOrbSym = [ s[1] for s in res.symmetries ]
  mo_tot_num = sum(NumOrbSym)
  ezfio.set_mo_basis_mo_tot_num(mo_tot_num)
  if res.occ_num is not None:
    ezfio.set_mo_basis_mo_occ(res.occ_num)

  MoTag = res.mo_types[-1]
  mo = res.mo_sets[MoTag]
  if len(mo) < mo_tot_num:
    newmo = orbital()
    newmo.eigenvalue = 0.
    newmo.vector = [0. for i in range(mo_tot_num)]
    while len(mo) < mo_tot_num:
      mo.append(newmo)
  Energies  = [ m.eigenvalue for m in mo ]
  ezfio.set_mo_basis_mo_energy(Energies)

  if res.occ_num is not None:
    OccNum    = res.occ_num[MoTag]
    while len(OccNum) < mo_tot_num:
      OccNum.append(0.)
    ezfio.set_mo_basis_mo_occ(OccNum)

  cls       = [ "v" for i in mo ]
  for i in res.closed_mos:
    cls[i] = 'c'
  for i in res.active_mos:
    cls[i] = 'a'
  ezfio.set_mo_basis_mo_classif(cls)

  MoMatrix = []
  for m in mo:
     for coef in m.vector:
         MoMatrix.append(coef)
  while len(MoMatrix) < len(mo[0].vector)**2:
    MoMatrix.append(0.)
  ezfio.set_mo_basis_mo_coef(MoMatrix)

# Determinants
  closed_mos = res.closed_mos
  nactive = ezfio.get_mo_basis_mo_active_num()
  dets_a = []
  dets_b = []
  for d in res.determinants:
    dnew_a = []
    dnew_b = []
    for x in d['alpha']:
      if x not in closed_mos:
        dnew_a.append(x+1)
    for x in d['beta']:
      if x not in closed_mos:
        dnew_b.append(x+1)
    for x in range(nactive-len(dnew_b)):
      dnew_b.append(0)
    dets_a.append( dnew_a )
    dets_b.append( dnew_b )

  coef = reduce(lambda x, y: x+y,res.det_coefficients,[])

  if len(dets_a[0]) > 0:
   ezfio.set_determinants_det_num(len(dets_a))
   ezfio.set_determinants_det_coef(coef)
   ezfio.set_determinants_det_occ(dets_a+dets_b)
  else:
   ezfio.set_determinants_det_num(1)
   ezfio.set_determinants_det_coef([1.])
   ezfio.set_determinants_det_occ(dets_a+dets_b)


write_ezfioFile(file,firstArg+".ezfio")