#!/usr/bin/env python3

"""Move a molecule over the center of a certain number of atoms in a bulk"""
import argparse
from math import radians
from pyRDTP import vaspio
from pyRDTP import molecule

PARSER = argparse.ArgumentParser(description="""Read two different files, one
                                 with the geometry of the bulk and another
                                 with the geometry of a molecule molecule,
                                 then put the molecule over the center of the
                                 atoms with the index given""")

PARSER.add_argument("-o", type=str, default='POSCAR',
                    help="Location of the output with the new coordinates.")
PARSER.add_argument("-rx", type=float, default=None,
                    help="Apply a rotation to the molecule in the x axis.")
PARSER.add_argument("-ry", type=float, default=None,
                    help="Apply a rotation to the molecule in the y axis.")
PARSER.add_argument("-rz", type=float, default=None,
                    help="Apply a rotation to the molecule in the z axis.")
PARSER.add_argument("-l", action='store_true', default=False,
                    help="""Use the lower atom instead of the centroid of the
                          molecule to move the molecule.""")
PARSER.add_argument("bulk", type=str,
                    help="CONTCAR/POSCAR containing the geometry of the bulk.")
PARSER.add_argument("molecule", type=str,
                    help="CONTCAR/POSCAR containing the geometry of the molecule.")
PARSER.add_argument("distance", type=float,
                    help="Distance between the molecule and the atoms.")
PARSER.add_argument("atoms", type=int, nargs='+',
                    help="""Atoms indexes (IN THE BULK) from which the
                         center will be computed.""")

ARGS = PARSER.parse_args()

CATY = molecule.CatalyticSystem('MyCat')
BULK = vaspio.read_from_file_vasp(ARGS.bulk, definition='bulk')
MOL = vaspio.read_from_file_vasp(ARGS.molecule, definition='molecule')

if ARGS.rx is not None:
    MOL.rotate((1, 0, 0), radians(ARGS.rx))
if ARGS.ry is not None:
    MOL.rotate((0, 1, 0), radians(ARGS.ry))
if ARGS.rz is not None:
    MOL.rotate((0, 0, 1), radians(ARGS.rz))
if ARGS.l:
    ORIGIN = MOL.atom_obtain_lowest(dimension='z').coords
else:
    ORIGIN = 'centroid'
ATOMS = tuple([BULK.atoms[int(index) - 1] for index in ARGS.atoms])

CATY.surface_set(BULK)
CATY.molecule_add(MOL)

CATY.move_over_multiple_atoms(MOL, ATOMS, ARGS.distance,
                              dimension='z', origin=ORIGIN)

vaspio.print_vasp(CATY, ARGS.o)
