#!/usr/bin/env python3

import common
import sys,os,time
from functools import reduce

from resultsFile import *

# Check command line
det_thr = 1.e-6

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)
file.convert_to_cartesian()
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.electrons_elec_alpha_num = res.num_alpha
  ezfio.electrons_elec_beta_num = res.num_beta
  ezfio.electrons_elec_num = res.num_beta + res.num_alpha

# Nuclei
  ezfio.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.nuclei_nucl_charge = charge
  ezfio.nuclei_nucl_coord = coord_x+coord_y+coord_z

# AO Basis
  import string
  at = []
  num_prim = []
  magnetic_number = []
  angular_number = []
  power_x = []
  power_y = []
  power_z = []
  coefficient = []
  exponent = []
  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))
    s = b.sym
    power_x.append( s.count("x") ) 
    power_y.append( s.count("y") )
    power_z.append( s.count("z") )
    coefficient.append(  b.coef )
    exponent.append( [ p.expo for p in b.prim ] )
  ezfio.ao_basis_ao_num = len(res.basis)
  ezfio.ao_basis_ao_nucl = at
  ezfio.ao_basis_ao_prim_num = num_prim
  ezfio.ao_basis_ao_power = power_x+power_y+power_z
  prim_num_max = max(num_prim)
  ezfio.ao_basis_ao_prim_num_max = 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.ao_basis_ao_coef = coef
  ezfio.ao_basis_ao_expo = expo

# MOs
  MoTag = res.mo_types[-1]
  if MoTag == 'Natural':
    res.mo_types.pop()
    MoTag = res.mo_types[-1]
  mo_tot_num = len(res.mo_sets[MoTag])
  ezfio.mo_basis_mo_tot_num = mo_tot_num

  try:
    newocc = [0. for i in range(mo_tot_num)]
    while len(res.occ_num[MoTag]) < mo_tot_num:
      res.occ_num[MoTag].append(newmo)
    ezfio.mo_basis_mo_occ = res.occ_num[MoTag]
  except:
    pass

  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.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.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.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.mo_basis_mo_coef = MoMatrix

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

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

  ezfio.mo_basis_mo_closed_num  = len(res.closed_mos)
  ezfio.mo_basis_mo_active_num  = len(res.active_mos)
  ezfio.mo_basis_mo_virtual_num = len(res.virtual_mos)

  if len(dets_a[0]) > 0:
   for i in range(len(coef),0,-1):
     i -= 1
     if abs(coef[i]) < det_thr :
       dets_a.pop(i)
       dets_b.pop(i)
       coef.pop(i)
   ezfio.determinants_det_num = len(coef)
   ezfio.determinants_det_coef = coef
   ezfio.determinants_det_occ = dets_a+dets_b
  else:
   ezfio.determinants_det_num = 1
   ezfio.determinants_det_coef = [1.]
   ezfio.determinants_det_occ = dets_a+dets_b


  ezfio.compute_eplf = True
  for i in "density density_lapl elf elf_grad eplf_lapl density_grad elf_grad elf_lapl eplf_grad".split():
    exec("ezfio.compute_%s = False" % i)


def main():
  write_ezfioFile(file,firstArg+".ezfio")

if __name__ == '__main__':
  main()