#!/usr/bin/env python
#===========================================================================#
#                                                                           #
#  File:       v2openmx.py                                                  #
#  Dependence: parse.py,write.py                                            #
#  Usage:      convert the POSCAR file to the dat file for openmx           #      
#  Author:     Shunhong Zhang <szhang2@ustc.edu.cn>                         #
#  Date:       Apr 21, 2017                                                 #
#                                                                           #
#===========================================================================#

from __future__ import print_function
from pysupercell.pysupercell import cryst_struct,verbose_pkg_info
from pysupercell import __version__
import os

#==================================================#
# input file for OPENMX (only for SCF calculation) #
#==================================================#

def write_openmx_dat(prefix,struct,datapath):
    print( "the generated input file should be checked carefully before real calculation!")
    print( "make sure that it is consistent with what you want!")
    with open(prefix+'.dat','w') as filename:
        print ('#\n# .dat file for OPENMX,generated by v2openmx\n#', file=filename)
        print ('System.CurrrentDirectory         ./    # default=./', file=filename)
        print ('System.Name                   '+ prefix, file=filename)
        print ('level.of.stdout                  1     # default=1 (1-3)', file=filename)
        print ('level.of.fileout                 1     # default=1 (0-2)\n', file=filename)
        print ('\n#\n# Definition of Atomic Species\n#\n', file=filename)
        print ('Species.Number',len(struct._species) , file=filename)
        print ('<Definition.of.Atomic.Species', file=filename)
        for ispec in struct._species:
            get_vps='locate '+datapath+'/VPS/'+ispec+'_'
            if ispec=='E':
                vps='E.vps'
            else:
                vps=os.popen(get_vps+'|grep PBE13S').readline().split('/')[-1].rstrip('\n')
            if not vps:
                vps=os.popen(get_vps+'|grep PBE13').readline().split('/')[-1].rstrip('\n')
            print (ispec, vps)
            print ('{0:8s} {1:30s} {2:20s}'.format(ispec,pao_set.pao_set[ispec],vps.rstrip('.vps')), file=filename)
        print('Definition.of.Atomic.Species>\n', file=filename)
        print('Atoms.SpeciesAndCoordinates.Unit  FRAC    # Ang|AU|FRAC ', file=filename)
        print('Atoms.Number ','{0:4d}'.format(struct._natom), file=filename)
        print('<Atoms.SpeciesAndCoordinates', file=filename)
        for iat in range(struct._natom):
            if struct._symbols[iat]=='E':
                charge_up,charge_dn=0,0
            else:
                get_ve='grep valence.electron '+datapath+'/VPS/'+struct._symbols[iat]
            try:
                ve=float(os.popen(get_ve+'_PBE13.vps 2>/dev/null').readline().split()[1])
            except:
                ve=float(os.popen(get_ve+'_PBE13S.vps').readline().split()[1])
            charge_up=ve/2.0
            charge_dn=ve/2.0
            print ('{0:4d} {1:4<s}'.format(iat+1,struct._symbols[iat]),("%20.16f %20.16f %20.16f" % (tuple(struct._pos[iat]))), file=filename)
            print ('{0:4.2f} {1:4.2f}'.format(charge_up,charge_dn), file=filename)
        print ('Atoms.SpeciesAndCoordinates>\n', file=filename)
        print ('Atoms.UnitVectors.Unit            Ang # Ang|AU', file=filename)
        print ('<Atoms.UnitVectors', file=filename)
        print ('\n'.join([('{:20.16f}'*3).format(*tuple(item)) for item in struct._cell]), file=filename)
        print ('Atoms.UnitVectors>\n', file=filename)
        print ('\n#\n# SCF or Electronic System\n#\n', file=filename)
        print ('scf.XcType                 GGA-PBE     # LDA|LSDA-CA|LSDA-PW|GGA-PBE', file=filename)

def get_args():
    import argparse
    desc_str='convert the POSCAR to the dat file for OPENMX'
    parser = argparse.ArgumentParser(prog='v2openmx.py', description = desc_str)
    parser.add_argument('--prefix',type=str,default='openmx',help='the prefix for the input file')
    parser.add_argument('--datapath',type=str,default='DFT_DATA13',help='the directory containing pseudopotentials')
    args = parser.parse_args()
    return parser, args



def main(args):
    datapath=os.popen('locate {}'.format(args.datapath)).readlines()
    if len(datapath)==0: 
        print('Cannot find {} for pseudopotentials!'.format(args.datapath))
        exit()
    else:
        datapath = datapath[0].rstrip('\n')
        print ('find DATA.PATH: {}'.format(datapath))
        print ('please confirm that the PAO and VPS files are in this directory!\n')
    struct = cryst_struct.load_poscar('POSCAR')
    write_openmx_dat(args.prefix,struct,datapath)




if __name__=='__main__':
    verbose_pkg_info(__version__)
    parser, args=get_args()
    main(args)

