#!/usr/bin/env python

#-------------------------------------------------------------------------------
#
#   volume.py
#   converts volumetric data from a3Dr to Gaussian Cube or XCrySDen XSF
#   written by G. Samsonidze (November 2008)
#
#-------------------------------------------------------------------------------

from common import format_index2mat, boolean_input
from matter import mat_check, mat_species, mat_format, mat_replicate, mat_read, paratec_read, vasp_read, espresso_read, siesta_read, tbpw_read, xyz_read, xsf_read

def main(argv = None):
   if argv is None:
      argv = sys.argv
   argc = len(argv)
   self = "volume.py"
   if argv[0][-len(self):] != self:
      sys.exit("\n   Rename script to %s\n" % self)
   if not (argc == 9 or argc == 10):
      sys.exit("\n   Usage: %s imfn imff ivfn ivff ovfn ovff phas cplx [hole]\n" % self +
               "   imfn = input matter file name\n" +
               "   imff = input matter file format\n" +
               "          (mat|paratec|vasp|espresso|siesta|tbpw|xyz|xsf)\n" +
               "   ivfn = input volumetric file name\n" +
               "   ivff = input volumetric file format (a3dr)\n" +
               "   ovfn = output volumetric file name\n" +
               "   ovff = output volumetric file format (cube|xsf)\n" +
               "   phas = remove wavefunction phase (true|false)\n" +
               "   cplx = complex wavefunction (re|im|abs|abs2)\n" +
               "   hole = plot hole (true|false)\n")
   imfn = argv[1]
   imff = argv[2].lower()
   ivfn = argv[3]
   ivff = argv[4].lower()
   ovfn = argv[5]
   ovff = argv[6].lower()
   if imff == 'mat':
      ierr, mat = mat_read(imfn)
   elif imff == 'paratec':
      ierr, mat = paratec_read(imfn)
   elif imff == 'vasp':
      ierr, mat = vasp_read(imfn)
   elif imff == 'espresso':
      ierr, mat = espresso_read(imfn)
   elif imff == 'siesta':
      ierr, mat = siesta_read(imfn)
   elif imff == 'tbpw':
      ierr, mat = tbpw_read(imfn)
   elif imff == 'xyz':
      ierr, mat = xyz_read(imfn)
   elif imff == 'xsf':
      ierr, mat = xsf_read(imfn)
   else:
      sys.exit("\n   Error: invalid input format %s\n" % imff)
   if ierr != 0:
      sys.exit("\n   Error: unable to read file %s\n" % imfn)
   ierr = mat_check(mat)
   if ierr != 0:
      sys.exit("\n   Error: invalid content in file %s\n" % imfn)
   if ivff == 'a3dr':
      ierr, state, energy, holpos, latvec, grid, psi = a3dr_read(ivfn)
   else:
      sys.exit("\n   Error: invalid input format %s\n" % ivff)
   if ierr != 0:
      sys.exit("\n   Error: unable to read file %s\n" % ivfn)
   format = format_index2mat[0]
   so = [0.0, 0.0, 0.0]
# temporary hack for a3dr format
# lattice vectors in a3dr file are written in %.6f format
# in paratec benzene example, lattice vectors read from a3dr file
# differ from lattice vectors read from paratec input file
# this results in identical atoms written to rho.cube
# (total of 16 atoms instead of 12 atoms)
# this hack computes the supercell dimensions and
# replaces the lattice vectors read from a3dr file
# with the lattice vectors read from the input file
# multiplied by the supercell dimensions
   mat = mat_format(mat, mat['fc'], mat['fo'], format, mat['fp'])
   a1 = mat['lv'][0]
   a2 = mat['lv'][1]
   a3 = mat['lv'][2]
   b1 = latvec[0]
   b2 = latvec[1]
   b3 = latvec[2]
   n1 = int(round((a1[0] * b1[0] + a1[1] * b1[1] + a1[2] * b1[2]) / (a1[0]**2 + a1[1]**2 + a1[2]**2)))
   n2 = int(round((a2[0] * b2[0] + a2[1] * b2[1] + a2[2] * b2[2]) / (a2[0]**2 + a2[1]**2 + a2[2]**2)))
   n3 = int(round((a3[0] * b3[0] + a3[1] * b3[1] + a3[2] * b3[2]) / (a3[0]**2 + a3[1]**2 + a3[2]**2)))
   b1 = [n1 * a1[0], n1 * a1[1], n1 * a1[2]]
   b2 = [n2 * a2[0], n2 * a2[1], n2 * a2[2]]
   b3 = [n3 * a3[0], n3 * a3[1], n3 * a3[2]]
   latvec = [b1, b2, b3]
# end of hack
   sv = latvec
   smin, smax, si, mat = mat_replicate(mat, so, sv, format)
   if mat['fo'] != format and mat['fv'] != format and mat['fp'] != format:
      sys.exit("\n   Error: unable to make supercell\n")
   if argc > 9 and holpos != []:
      ierr, hole = boolean_input(argv[9])
      if ierr != 0:
         sys.exit("\n   Error: invalid hole %s\n" % argv[9])
      if hole != 0:
         mat['as'].append(0)
         mat['ap'].append(holpos)
         mat['na'] += 1
         mat['ns'] = len(mat_species(mat['na'], mat['as']))
   ierr, phas = boolean_input(argv[7])
   if ierr != 0:
      sys.exit("\n   Error: invalid phase %s\n" % argv[7])
   if phas != 0:
      rmax = 0.0
      for i in range(grid[0]):
         for j in range(grid[1]):
            for k in range(grid[2]):
               r = math.sqrt(math.pow(psi[k][j][i][0], 2) + math.pow(psi[k][j][i][1], 2))
               if r > rmax:
                  rmax = r
                  phase = math.atan2(psi[k][j][i][1], psi[k][j][i][0])
      cs = math.cos(phase)
      sn = math.sin(phase)
      for i in range(grid[0]):
         for j in range(grid[1]):
            for k in range(grid[2]):
               x = psi[k][j][i][0]
               y = psi[k][j][i][1]
               psi[k][j][i][0] = cs * x + sn * y
               psi[k][j][i][1] = -sn * x + cs * y
   ierr, cplx = cplx_input(argv[8])
   if ierr != 0:
      sys.exit("\n   Error: invalid cplx %s\n" % argv[8])
   phi = []
   for i in range(grid[0]):
      phi.append([])
      for j in range(grid[1]):
         phi[i].append([])
         for k in range(grid[2]):
            if cplx == 0:
               chi = psi[k][j][i][0]
            elif cplx == 1:
               chi = psi[k][j][i][1]
            elif cplx == 2:
               chi = math.sqrt(math.pow(psi[k][j][i][0], 2) + math.pow(psi[k][j][i][1], 2))
            elif cplx == 3:
               chi = math.pow(psi[k][j][i][0], 2) + math.pow(psi[k][j][i][1], 2)
            else:
               chi = 0.0
            if phas != 0 and (cplx == 2 or cplx == 3) and psi[k][j][i][0] < 0.0:
               chi = -chi
            phi[i][j].append(chi)
   if ovff == 'cube':
      i = 0
   elif ovff == 'xsf':
      i = 1
   else:
      sys.exit("\n   Error: invalid output format %s\n" % ovff)
   fc = format_index2mat[i]
   fo = format_index2mat[i]
   fv = format_index2mat[i]
   fp = format_index2mat[i]
   mat = mat_format(mat, fc, fo, fv, fp)
   if mat['fc'] != fc or mat['fo'] != fo or mat['fv'] != fv or mat['fp'] != fp:
      sys.exit("\n   Error: invalid content in file %s\n" % imfn)
   na = mat['na']
   aspecies = []
   for i in range(na):
      aspecies.append(mat['as'][i])
   acharge = []
   for i in range(na):
      acharge.append(float(mat['as'][i]))
   aposition = []
   for i in range(na):
      aposition.append([])
      for j in range(3):
         aposition[i].append(mat['ap'][i][j])
   origin = []
   for j in range(3):
      origin.append(mat['lo'][j])
   latvec = []
   for i in range(3):
      latvec.append([])
      for j in range(3):
         latvec[i].append(mat['lv'][i][j])
   step = []
   for i in range(3):
      step.append([])
      for j in range(3):
         step[i].append(latvec[i][j] / float(grid[i]))
   if ovff == 'cube':
      ierr = cub_write(ovfn, state, energy, na, aspecies, acharge, aposition, grid, origin, step, phi)
   elif ovff == 'xsf':
      ierr = xsf_write(ovfn, state, energy, na, aspecies, aposition, grid, origin, latvec, phi)
   else:
      sys.exit("\n   Error: invalid output format %s\n" % ovff)
   if ierr != 0:
      sys.exit("\n   Error: unable to write file %s\n" % ovfn)
   return 0

def a3dr_read(file):
   ierr = 0
   exciton_state = 0
   exciton_energy = 0.0
   hole_position = []
   lattice_vectors = []
   grid_size = []
   wavefunction = []
   try:
      h = open(file, 'r')
   except:
      ierr = 1
   if ierr == 0:
      while 1:
         s = h.readline()
         t = s.split()
         t.append('')
         if t[1] == 'ie':
            exciton_state = int(t[3])
         if t[1] == 'e':
            exciton_energy = float(t[3])
         if t[1] == 'rh':
            for j in range(3):
               hole_position.append(float(t[j + 3]))
         if t[1] == 'unit':
            for i in range(3):
               lattice_vectors.append([])
               s = h.readline()
               t = s.split()
               for j in range(3):
                  lattice_vectors[i].append(float(t[j + 3]))
         if t[1] == 'ni':
            for i in range(3):
               grid_size.append(int(t[i + 3]))
         if t[1] == 'real':
            break
      for k in range(grid_size[2]):
         wavefunction.append([])
         for j in range(grid_size[1]):
            wavefunction[k].append([])
            for i in range(grid_size[0]):
               while 1:
                  s = h.readline()
                  t = s.split()
                  if len(t) != 0:
                     break
               if len(t) > 1:
                  wavefunction[k][j].append([])
                  for l in range(2):
                     wavefunction[k][j][i].append(float(t[l]))
# temporary hack for a3dr format
# in case of real wavefunctions, paratec does not write imaginary column
# to a3dr file, while plotxct does write zero imaginary column
               else:
                  wavefunction[k][j].append([float(t[0]), 0.0])
# end of hack
      h.close()
      for i in range(3):
         if grid_size[i] < 1:
            ierr = 1
   return ierr, exciton_state, exciton_energy, hole_position, lattice_vectors, grid_size, wavefunction

def cub_write(file, state, energy, na, aspecies, acharge, aposition, grid, origin, step, vol):
   ierr = 0
   try:
      h = open(file, 'w')
   except:
      ierr = 1
   if ierr == 0:
      if state == 0 and abs(energy) < 1.0e-6:
         s = 'This file was produced with PARATEC\n'
         s += '\n'
      else:
         s = 'exciton state = %i\n' % state
         s += 'exciton energy = %.4f eV\n' % energy
      h.write(s)
      s = ' %4i' % na
      for j in range(3):
         s += ' %11.6f' % origin[j]
      s += '\n'
      h.write(s)
      for i in range(3):
         s = ' %4i' % grid[i]
         for j in range(3):
            s += ' %11.6f' % step[i][j]
         s += '\n'
         h.write(s)
      for i in range(na):
         s = ' %4i' % aspecies[i]
         s += ' %11.6f' % acharge[i]
         for j in range(3):
            s += ' %11.6f' % aposition[i][j]
         s += '\n'
         h.write(s)
      s = ''
      for i in range(grid[0]):
         for j in range(grid[1]):
            for k in range(grid[2]):
               s += ' %12.5E' % vol[i][j][k]
               if (k + 1) % 6 == 0 or k + 1 == grid[2]:
                  s += '\n'
                  h.write(s)
                  s = ''
      h.close()
   return ierr

def xsf_write(file, state, energy, na, aspecies, aposition, grid, origin, latvec, vol):
   ierr = 0
   try:
      h = open(file, 'w')
   except:
      ierr = 1
   if ierr == 0:
      h.write('CRYSTAL\n')
      h.write('PRIMVEC\n')
      for i in range(3):
         s = ''
         for j in range(3):
            s += ' %11.6f' % latvec[i][j]
         s += '\n'
         h.write(s)
      h.write('CONVVEC\n')
      for i in range(3):
         s = ''
         for j in range(3):
            s += ' %11.6f' % latvec[i][j]
         s += '\n'
         h.write(s)
      h.write('PRIMCOORD\n')
      s = ' %4i %i\n' % (na, 1)
      h.write(s)
      for i in range(na):
         s = ' %4i' % aspecies[i]
         for j in range(3):
            s += ' %11.6f' % aposition[i][j]
         for j in range(3):
            s += ' %11.6f' % 0.0
         s += '\n'
         h.write(s)
      h.write('BEGIN_BLOCK_DATAGRID_3D\n')
      t = 'exciton_state_%i_energy_%.4f_eV' % (state, energy)
      s = '%s\n' % t
      h.write(s)
      s = 'BEGIN_DATAGRID_3D_%s\n' % t
      h.write(s)
      s = ''
      for i in range(3):
         j = grid[i] + 1
         s += ' %11i' % j
      s += '\n'
      h.write(s)
      s = ''
      for i in range(3):
         s += ' %11.6f' % origin[i]
      s += '\n'
      h.write(s)
      for i in range(3):
         s = ''
         for j in range(3):
            s += ' %11.6f' % latvec[i][j]
         s += '\n'
         h.write(s)
      s = ''
      for k in range(grid[2] + 1):
         if k < grid[2]:
            pk = k
         else:
            pk = 0
         for j in range(grid[1] + 1):
            if j < grid[1]:
               pj = j
            else:
               pj = 0
            for i in range(grid[0] + 1):
               if i < grid[0]:
                  pi = i
               else:
                  pi = 0
               s += ' %12.5E' % vol[pi][pj][pk]
               if (i + 1) % 6 == 0 or i == grid[0]:
                  s += '\n'
                  h.write(s)
                  s = ''
      h.write('END_DATAGRID_3D\n')
      h.write('END_BLOCK_DATAGRID_3D\n')
      h.close()
   return ierr

def cplx_input(cplx):
   ierr = 0
   index = -1
   cplx = cplx.lower()
   if cplx == 'r' or cplx == 're' or cplx == 'real':
      index = 0
   elif cplx == 'i' or cplx == 'im' or cplx == 'imag' or cplx == 'imaginary':
      index = 1
   elif cplx == 'a' or cplx == 'abs' or cplx == 'absolute':
      index = 2
   elif cplx == 'a2' or cplx == 'abs2' or cplx == 'absolute2':
      index = 3
   else:
      ierr = 1
   return ierr, index

if __name__ == "__main__":
   import sys
   import math
   sys.exit(main())

