#!/usr/bin/env python
import argparse
import yaml
import sys
from cosymlib import file_io
from cosymlib.molecule.geometry import Geometry
from cosymlib.cosym_api import Cosymlib
from cosymlib.shape import shape_tools
import os

parser = argparse.ArgumentParser(description='Cosym')
parser.add_argument(type=str, dest='input_file', help='input file name(+extension)')
parser.add_argument(type=str, dest="yaml_input", nargs='?', default=None,
                    help='Perform the calculations with the command file')
parser.add_argument('-o', '--output_name', dest='output_name', default=None, help='save in file name')
parser.add_argument('-info', action='store_true', default=False,
                    help='return information about the input geometries')

parser.add_argument('-c', '--central_atom', action='store', dest='central_atom',
                    type=int, default=0, help='position of the central atom if exist')
parser.add_argument('-custom_ref', action='store', dest='custom_ref', default=None,
                    help='take a given structure from the file and use it as reference')
parser.add_argument('-fix_perm', '--fix_permutation', dest='fix_permutation', action='store_true', default=False,
                    help='use the given permutation to perform a calculation')
# parser.add_argument('-connectivity', dest='connectivity', action='store', default=None,
#                     help='Connect a set of atoms by...')


# Shape input flags
group_shape = parser.add_argument_group('Shape')
group_shape.add_argument('-m', '--measure',
                         dest='measure',
                         action='store',
                         default=None,
                         help='Shape measure of input structure with reference polyhedra')
group_shape.add_argument('-l', '--labels', action='store_true',
                         dest='labels',
                         default=False,
                         help='show the reference labels for a given structure')
group_shape.add_argument('-s', '--structure',
                         dest='structure',
                         action='store_true',
                         default=False,
                         help='return the closes input structure to the reference shape')
group_shape.add_argument('-r', '--references',
                         dest='references',
                         action='store_true',
                         default=False,
                         help='return a file with the coordinates of reference polyhedra')


args = parser.parse_args(sys.argv[1:])
print('Starting...')

if args.yaml_input:
    with open(args.yaml_input, 'r') as stream:
        input_parameters = yaml.load(stream, Loader=yaml.FullLoader)

    for key, value in input_parameters.items():
        if key.lower() in args:
            setattr(args, key.lower(), value)
        else:
            raise KeyError("Key %s is not valid" % key)

if args.input_file is not None:
    structures = file_io.read_input_file(args.input_file)
    symobj = Cosymlib(structures)

    try:
        n_atoms = structures[0].geometry.get_n_atoms()
    except AttributeError:
        try:
            n_atoms = structures[0].get_n_atoms()
        except AttributeError:
            raise AttributeError('Molecule object not found')
    if args.central_atom != 0:
        n_atoms -= 1

if args.references:
    test_structure = []
    symbols = ['H' for _ in range(n_atoms)]
    symbols.append('N')
    for label in shape_tools.get_structure_references(n_atoms):
        test_structure.append(Geometry(symbols=symbols,
                                       positions=shape_tools.get_test_structure(label, central_atom=args.central_atom),
                                       name=label))
    input_file, file_extension = os.path.splitext(args.input_file)
    if args.central_atom == 0:
        output = open('{}/L{}.xyz'.format(os.path.dirname(input_file), n_atoms), 'w')
    else:
        output = open('{}/ML{}.xyz'.format(os.path.dirname(input_file), n_atoms), 'w')
    output.write(file_io.write_file_xyz(test_structure))
    # file_io.write_file_xyz(test_structure, output_name='ML{}_ref'.format(n_atoms))

if args.info:
    file_io.write_input_info(structures, output_name=args.output_name)

# Shape's commands
if args.labels:
    if args.central_atom == 0:
        print(shape_tools.get_shape_label_info(n_atoms))
    else:
        print(shape_tools.get_shape_label_info(n_atoms))

reference_polyhedron = []
if args.measure == 'custom':
    reference_polyhedron = file_io.get_molecule_from_file_xyz(args.custom_ref)
    [x.set_positions(args.central_atom - 1) for x in reference_polyhedron]
elif not reference_polyhedron:
    if args.measure == 'all':
        try:
            n_atoms = structures[0].geometry.get_n_atoms()
        except AttributeError:
            n_atoms = structures[0].get_n_atoms()
        c = int(bool(args.central_atom))
        reference_polyhedron = shape_tools.get_structure_references(n_atoms - c)
    else:
        reference_polyhedron = [args.measure]

if args.structure:
    symobj.write_shape_structure_2file(reference_polyhedron,
                                       central_atom=args.central_atom,
                                       output_name=args.output_name,
                                       fix_permutation=args.fix_permutation)

elif args.measure or args.custom_ref:
    symobj.write_shape_measure_2file(reference_polyhedron,
                                     central_atom=args.central_atom,
                                     output_name=args.output_name,
                                     fix_permutation=args.fix_permutation)

print('\nEnd of cosym calculation')