#!/usr/bin/env python3

"""Create a direct vector between two atoms and move the selected atoms
to a certain distance"""
import argparse
import numpy as np
from pyRDTP import geomio

PARSER = argparse.ArgumentParser(description='''Select atoms and move it to a
                                 certain distance''')
PARSER.add_argument("-i", type=str, default='POSCAR',
                    help="Location of the POSCAR file.")
PARSER.add_argument("-o", type=str, default='POSCAR_new',
                    help="Location of the output with the new coordinates.")
PARSER.add_argument("-s", nargs='+', type=int,
                    help="Selected atoms")
PARSER.add_argument("-oa", type=int,
                    help="Use an origin atom instead of a centroid")
PARSER.add_argument("-zsort", action='store_true', default=False,
                    help="Sort by z coordinate")
PARSER.add_argument("-elemsort", action='store_true', default=False,
                    help="Sort by element alphabetically")
PARSER.add_argument("end_atom", type=int,
                    help="End atom.")
PARSER.add_argument("distance", type=float,
                    help="""Distance between the selected atoms and the
                    endpoint""")

ARGS = PARSER.parse_args()
ARGS.oa -= 1
ARGS.end_atom -= 1
SELECTED = np.asarray(ARGS.s) - 1

VASPFILE = geomio.VaspContcar(ARGS.i)
COORDS = VASPFILE.universal_convert()
MOL_OBJ = geomio.MolObj(COORDS)
MOL_OBJ = MOL_OBJ.write()

SELECTION = MOL_OBJ.select(SELECTED, mol_obj=True)

if ARGS.oa:
    ORIGIN = MOL_OBJ.atoms[ARGS.oa].coords
else:
    ORIGIN = SELECTION.centroid

END = MOL_OBJ[ARGS.end_atom].coords

MOVE_VECT = END - ORIGIN
NORM = np.linalg.norm(MOVE_VECT)
NEW_NORM = NORM - ARGS.distance
MOVE_VECT /= NORM
MOVE_VECT *= NEW_NORM

SELECTION.move_vector(MOVE_VECT)

NEW_COORDS = geomio.MolObj(MOL_OBJ)
NEW_COORDS = NEW_COORDS.universal_convert()
OUTPUT = geomio.VaspContcar(NEW_COORDS)
OUTPUT.file_write(ARGS.o, z_sort=ARGS.zsort, elem_sort=ARGS.elemsort)
