#! No Comments

import os
import shutil
from Atom import *
from operator import attrgetter


print("This script is to automatically generate the input files from .xyz file")
print("for vacuum relaxation (6.0 A * 6.0 A * 30.0 A; no e-field)")
xyz_name0 = raw_input("Enter xyz filename: ") # Ignore the error message in Windows. Unix uses "raw_input"
xyz_name = xyz_name0 + ".xyz"
sys_name = xyz_name0
latt_parameter = ("%.14f" % float(1))
current_dir = os.getcwd()
os.mkdir(sys_name)
pos_name = "POSCAR"
xyz = open(xyz_name, 'r')
pos = open(current_dir + "/" + sys_name + "/" + pos_name, 'w')
a = ("%.16f" % 30)
b = ("%.16f" % 0)
c = ("%.16f" % 6)
pos.write(sys_name + '\n')
pos.write("   " + latt_parameter + '\n')
pos.write(c.rjust(23) + b.rjust(22) + b.rjust(22) + '\n')
pos.write(b.rjust(23) + c.rjust(22) + b.rjust(22) + '\n')
pos.write(b.rjust(23) + b.rjust(22) + a.rjust(22) + '\n')


num = xyz.readline()
print("Total Atoms: " + num)
Atom_list = []
cnt = 0

while True:
    line = xyz.readline().split()
    print(line)
    if len(line) < 1:
        break
    Atom_list.append(Atom())
    Atom_list[cnt].info_input(line[1], float(line[2]), float(line[3]), float(line[4]))
    cnt += 1
xyz.close()

tag_count = []
Atom_list.sort(key=attrgetter('tag'))
for x in range(int(num)):
    print(Atom_list[x].coord_info())  # For debug
    tag_count.append(Atom_list[x].tag)
tag_names = list(set(tag_count))
tag_names.sort(key=tag_count.index)
print(tag_count)
print(tag_names)
pos.write(" "*3)
for name in tag_names:
    pos.write(name.ljust(5))
pos.write('\n' + " "*5)
for name in tag_names:
    times = tag_count.count(name)
    pos.write(repr(times).ljust(6))
pos.write('\n' + "Selective dynamics" + '\n' + "Cartesian" + '\n')
for x in range(int(num)):
    pos.write("  " + Atom_list[x].coord_info() + "   T" * 3 + '\n')

# Construct POTCAR file
target = open(current_dir + "/" + sys_name + "/POTCAR", "w")
for tag in tag_names:
    pot = open(current_dir + "/POTLib/" + tag + "/POTCAR", "r")
    pot_data = pot.read()
    target.write(pot_data)
    pot.close()
target.close()


# Copy INCAR file
shutil.copy("INCAR", current_dir + "/" + sys_name)


# Copy KPOINTS file
shutil.copy("KPOINTS", current_dir + "/" + sys_name)


# Construct submit.sh file
target = open(current_dir + "/" + sys_name + "/submit.sh", "w")
target.write("#\\!/bin/bash\n#\n#MOAB -N " + sys_name + "(B)")
source = open(r"msub_generic", "r")
target.write(source.read())
target.close()
source.close()


# Submit job
os.chdir(current_dir + '/' + sys_name + '/')
os.system('msub submit.sh')
