#!/usr/bin/env python3

"""Separe a molecule from a bulk and apply a vector to the molecule"""
from pyRDTP import geomio
import argparse

PARSER = argparse.ArgumentParser(description='Apply a vector to a molecule.')
PARSER.add_argument("-i", type=str, default='POSCAR',
                    help="Location of the POSCAR file.")
PARSER.add_argument("-o", type=str, default='POSCAR',
                    help="Location of the output with the new coordinates.")
PARSER.add_argument("-d", action='store_const',
                    const='direct', default='cartesian',
                    help="Use direct coordinates instead of cartesian.")
PARSER.add_argument("-e", type=str, default=False,
                    help="""Use elements to select the molecule instead of
                         layers""")
PARSER.add_argument("-ie", type=str, default=False,
                    help="""Use inverse elements to select the molecule
                         instead of layers""")
PARSER.add_argument("-layers", nargs='+',
                    help="Number of layers on the bulk")
PARSER.add_argument("x", type=float, default=0.0,
                    help="Distance vector (in x axis) to apply to molecule.")
PARSER.add_argument("y", type=float, default=0.0,
                    help="Distance vector (in y axis) to apply to molecule.")
PARSER.add_argument("z", type=float, default=0.0,
                    help="Distance vector (in z axis) to apply to molecule.")

ARGS = PARSER.parse_args()

VECTOR = [ARGS.x, ARGS.y, ARGS.z]
BULK = vaspio.read_from_file_vasp(ARGS.i, definition='bulk')
BULK = geomio.file_to_mol(ARGS.i, 'contcar', bulk=True)

if ARGS.e:
    CAT_SYS = BULK.split_by_element(ARGS.e, invert=False)
if ARGS.ie:
    CAT_SYS = BULK.split_by_element(ARGS.ie, invert=True)
else:
    BULK.layer_detect_smart(threshold=0.04)
    CAT_SYS = BULK.split_by_layer(ARGS.layers)

CAT_SYS.molecules[0].move_vector(VECTOR, ARGS.d)
vaspio.print_vasp(CAT_SYS, ARGS.o)
geomio.mol_to_file(
