#  Imports
import argparse
import sys
import re

#  Take arguments
parser = argparse.ArgumentParser(description='Convert last frame of geopt file to POSCAR file so that ASE can read it.')
parser.add_argument('-f',
                    '--file',
                    type=str,
                    required=True,
                    metavar='',
                    help='the name of the geopt file')
parser.add_argument('-i',
                    '--ion',
                    type=str,
                    required=True,
                    metavar='',
                    help='the name of the ion file associated with the geopt file')
args = parser.parse_args()
geopt_file_name = args.file
ion_file_name = args.ion

#  Find number of each species in ion file
species_dict = {}
atom_type_regex = re.compile(r'ATOM_TYPE: ([a-zA-Z]{1,2})')
num_atoms_regex = re.compile(r'N_TYPE_ATOM: ([0-9]+)')
ion_file = open(ion_file_name, 'r')
ion_line = ion_file.readline()
while ion_line:
    atom_type_match = atom_type_regex.match(ion_line)
    if atom_type_match:
        atom_type = atom_type_match.group(1)
        ion_line = ion_file.readline()
        num = num_atoms_regex.match(ion_line).group(1)
        species_dict[atom_type] = num
    ion_line = ion_file.readline()
ion_file.close()

#  Find lattice vectors in ion file (commented if generated by SPARC Python API)
cell_regex = re.compile(r'#CELL:\s*([0-9\.\-]+)\s+([0-9\.\-]+)\s+([0-9\.\-]+)\s*')
vec_regex = re.compile(r'#\s*([0-9\.\-]+)\s+([0-9\.\-]+)\s+([0-9\.\-]+)\s*')
ion_file = open(ion_file_name, 'r')
ion_line = ion_file.readline()
cell_found = False
while ion_line:
    cell_match = cell_regex.match(ion_line)
    if cell_match:
        cell_found = True
        cell_x = float(cell_match.group(1))
        cell_y = float(cell_match.group(2))
        cell_z = float(cell_match.group(3))
        ion_line = ion_file.readline()
        ion_line = ion_file.readline()
        vec_1 = [str(float(vec_regex.match(ion_line).group(1))*cell_x),
                 str(float(vec_regex.match(ion_line).group(2))*cell_y),
                 str(float(vec_regex.match(ion_line).group(3))*cell_z)]
        ion_line = ion_file.readline()
        vec_2 = [str(float(vec_regex.match(ion_line).group(1))*cell_x),
                 str(float(vec_regex.match(ion_line).group(2))*cell_y),
                 str(float(vec_regex.match(ion_line).group(3))*cell_z)]
        ion_line = ion_file.readline()
        vec_3 = [str(float(vec_regex.match(ion_line).group(1))*cell_x),
                 str(float(vec_regex.match(ion_line).group(2))*cell_y),
                 str(float(vec_regex.match(ion_line).group(3))*cell_z)]
    ion_line = ion_file.readline()
ion_file.close()
if not cell_found:
    print('It appears that your ion file does not include a comment of the lattice vectors.\n'
          'This script is intended for ion files generated by the SPARC python API,\n'
          'which include a comment of the lattice vectors. Aborting.')
    sys.exit(2)


#  Find number of frames in geopt file
num_frames = 0
geopt_file = open(geopt_file_name, 'r')
geopt_line = geopt_file.readline()
while geopt_line:
    if geopt_line == ':R(Bohr):\n':
        num_frames += 1
    geopt_line = geopt_file.readline()
geopt_file.close()

#  Extract relevant coordinate block from geopt file
coord_block = ''
frame = 0
coord_regex = re.compile(r'\s*[0-9\.\-]+\s+[0-9\.\-]+\s+[0-9\.\-]+\s*')
geopt_file = open(geopt_file_name, 'r')
geopt_line = geopt_file.readline()
while geopt_line:
    if geopt_line == ':R(Bohr):\n':
        frame += 1
        if frame == num_frames:
            geopt_line = geopt_file.readline()
            while coord_regex.match(geopt_line):
                coord_block += (geopt_line[:-1] + ' F F F\n')
                geopt_line = geopt_file.readline()
    geopt_line = geopt_file.readline()
geopt_file.close()

#  Build POSCAR text and write to POSCAR file
poscar_text = ''
for species in species_dict:
    poscar_text += (species + ' ')
poscar_text += '\n0.52917721056\n'
poscar_text += ' '.join(vec_1) + '\n'
poscar_text += ' '.join(vec_2) + '\n'
poscar_text += ' '.join(vec_3) + '\n'
for species in species_dict:
    poscar_text += (str(species_dict[species]) + ' ')
poscar_text += '\nSelective dynamics\nCartesian\n'
poscar_text += coord_block
poscar_file = open('POSCAR', 'w')
poscar_file.write(poscar_text)
poscar_file.close()
print('POSCAR file generated!')
