# This is the script to connect ABCluster with VASP.
#
# Usage: runCP2K.py $inp$ $out$ $xxx$
#

import sys
import os

def XYZ2POSCAR(fn, pbc):
    with open(fn, "r") as fd: lines = fd.read().splitlines()
    sym_last = None
    sym_num = 0
    symbols = []
    coords = []
    for line in lines[2:]:
        words = line.split()
        symbol = words[0]
        coords.append([float(word) for word in words[1:]])
        if sym_last != symbol:
            if sym_last is not None:
                symbols.append((sym_last, sym_num))            
            sym_last = symbol
            sym_num = 1
        else:
            sym_num += 1
    symbols.append((sym_last, sym_num))           
    with open("POSCAR", "w") as fd:
        fd.write("%s\n" % (fn))
        fd.write("1.0\n")
        for v in pbc:
            fd.write("%15.8f %15.8f %15.8f\n" % (v[0], v[1], v[2]))
        fd.write("".join(["%6s " % symbol[0] for symbol in symbols])+"\n")
        fd.write("".join(["%6d " % symbol[1] for symbol in symbols])+"\n")
        fd.write("Carti\n")
        for coord in coords:
            fd.write("%15.8f %15.8f %15.8f\n" % (coord[0], coord[1], coord[2]))

def CONTCAR2XYZ(fn):
    energy = float(os.popen("grep \"next E\" OUTCAR | tail -n 1 | awk '{print $4}'").read())
    with open("CONTCAR", "r") as fd: lines = fd.read().splitlines()
    pbc = [[float(word) for word in line.split()] for line in lines[2:5]]
    elements = lines[5].split()
    nums = [int(word) for word in lines[6].split()]
    symbols = []
    for n,e in zip(nums, elements):
        symbols += [e]*n   
    coords = [[float(word) for word in line.split()] for line in lines[8:8+len(symbols)]]
    direct = True if lines[7] == "Direct" else False
    with open(fn, "w") as fd:
        fd.write("%d\n" % (len(symbols)))
        fd.write("%15.8f\n" % (energy))
        for symbol, coord in zip(symbols, coords):
            if direct:
                xyz = [coord[0]*pbc[0][i]+coord[1]*pbc[1][i]+coord[2]*pbc[2][i] for i in range(3)]
            else: # Cartesian.
                xyz = coord
            fd.write("%5s %15.8f %15.8f %15.8f\n" % (symbol, xyz[0], xyz[1], xyz[2]))
    return pbc

if __name__ == "__main__":
    vaspcmd = "module load intel_mpi/2018_u1 gcc/8.3.0 ; mpirun /lustre/home/junz/bin/vasp-5.4.1-patch/vasp_std"
    pbc = [[8.1426336762816440, 0.0000000000000000, 0.0000000000000000],
            [-4.0677955560985355, 7.0495917923679352, 0.0000000000000000],
            [0.0000000000000000, 0.0000000000000000, 30.0000000000000000]]
    # Arguments.
    inp_fn = sys.argv[1] 
    out_fn = sys.argv[2]
    sys_fn = sys.argv[3]
    vasplog = "vasp.log"
    # Run the job.
    XYZ2POSCAR(inp_fn, pbc)
    os.system("%s > %s" % (vaspcmd, vasplog))
    pbc = CONTCAR2XYZ(out_fn)
    # Keep or delete files.
    os.system("rm -rf %s CHG CHGCAR DOSCAR EIGENVAL IBZKPT OSZICAR PCDAT REPORT vasprun.xml WAVECAR XDATCAR" % (vasplog))

