# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

from ctypes import c_double, c_int
from mpi4py import MPI
import numpy as np
import pickle
import sys
from lammps import lammps

__author__ = "Sarath Menon, Jan Janssen"
__copyright__ = (
    "Copyright 2020, Max-Planck-Institut für Eisenforschung GmbH - "
    "Computational Materials Design (CM) Department"
)
__version__ = "1.0"
__maintainer__ = "Jan Janssen"
__email__ = "janssen@mpie.de"
__status__ = "production"
__date__ = "Feb 28, 2020"

# dict for extract atom methods
atom_properties = {
    "x": {"type": 3, "gtype": 1, "dim": 3},
    "mass": {"type": 2, "gtype": 1, "dim": 1},
    "id": {"type": 0, "gtype": 0, "dim": 1},
    "type": {"type": 0, "gtype": 0, "dim": 1},
    "mask": {"type": 0, "gtype": 0, "dim": 1},
    "v": {"type": 3, "gtype": 1, "dim": 3},
    "f": {"type": 3, "gtype": 1, "dim": 3},
    "molecule": {"type": 0, "gtype": 0, "dim": 1},
    "q": {"type": 2, "gtype": 1, "dim": 1},
    "mu": {"type": 3, "gtype": 1, "dim": 3},
    "omega": {"type": 3, "gtype": 1, "dim": 3},
    "angmom": {"type": 3, "gtype": 1, "dim": 3},
    "torque": {"type": 3, "gtype": 1, "dim": 3},
    "radius": {"type": 2, "gtype": 1, "dim": 1},
    "image": {"type": 2, "gtype": 0, "dim": 3},
    # we can add more quantities as needed
    # taken directly from atom.cpp -> extract()
}

# Lammps executable
args = ["-screen", "none"]
if len(sys.argv) > 1:
    args.extend(sys.argv[1:])
job = lammps(cmdargs=args)


def extract_compute(funct_args):
    # if MPI.COMM_WORLD.rank == 0:
    id = funct_args[0]
    style = funct_args[1]
    type = funct_args[2]
    length = funct_args[3]
    width = funct_args[4]
    filtered_args = [id, style, type]

    val = job.extract_compute(*filtered_args)

    # now process
    # length should be set
    data = []
    if style == 1:
        length = job.get_natoms()
    if type == 2:
        for i in range(length):
            dummy = []
            for j in range(width):
                dummy.append(val[i][j])
            data.append(np.array(dummy))
        data = np.array(data)
    elif type == 1:
        for i in range(length):
            data.append(val[i])
        data = np.array(data)
    else:
        data = val
    return data


def get_version(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.version()


def get_file(funct_args):
    job.file(*funct_args)
    return 1


def commands_list(funct_args):
    job.commands_list(*funct_args)
    return 1


def commands_string(funct_args):
    job.commands_string(*funct_args)
    return 1


def extract_setting(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.extract_setting(*funct_args)


def extract_global(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.extract_global(*funct_args)


def extract_box(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.extract_box(*funct_args)


def extract_atom(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        # extract atoms return an internal data type
        # this has to be reformatted
        name = str(funct_args[0])
        if not name in atom_properties.keys():
            return []

        # this block prevents error when trying to access values
        # that do not exist
        try:
            val = job.extract_atom(name, atom_properties[name]["type"])
        except ValueError:
            return []
        # this is per atom quantity - so get
        # number of atoms - first dimension
        natoms = job.get_natoms()
        # second dim is from dict
        dim = atom_properties[name]["dim"]
        data = []
        if dim > 1:
            for i in range(int(natoms)):
                dummy = [val[i][x] for x in range(dim)]
                data.append(dummy)
        else:
            data = [val[x] for x in range(int(natoms))]

        return np.array(data)


def extract_fix(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.extract_fix(*funct_args)


def extract_variable(funct_args):
    # in the args - if the third one,
    # which is the type is 1 - a lammps array is returned
    if MPI.COMM_WORLD.rank == 0:
        # if type is 1 - reformat file
        try:
            data = job.extract_variable(*funct_args)
        except ValueError:
            return []
        if funct_args[2] == 1:
            data = np.array(data)
        return data


def get_natoms(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.get_natoms()


def set_variable(funct_args):
    return job.set_variable(*funct_args)


def reset_box(funct_args):
    job.reset_box(*funct_args)
    return 1


def gather_atoms(funct_args):
    # extract atoms return an internal data type
    # this has to be reformatted
    name = str(funct_args[0])
    if name not in atom_properties.keys():
        return []

    # this block prevents error when trying to access values
    # that do not exist
    try:
        val = job.gather_atoms(
            name, atom_properties[name]["gtype"], atom_properties[name]["dim"]
        )
    except ValueError:
        return []
    # this is per atom quantity - so get
    # number of atoms - first dimension
    val = list(val)
    dim = atom_properties[name]["dim"]
    if dim > 1:
        data = [val[x : x + dim] for x in range(0, len(val), dim)]
    else:
        data = list(val)
    return np.array(data)


def gather_atoms_concat(funct_args):
    # extract atoms return an internal data type
    # this has to be reformatted
    name = str(funct_args[0])
    if not name in atom_properties.keys():
        return []

    # this block prevents error when trying to access values
    # that do not exist
    try:
        val = job.gather_atoms_concat(
            name, atom_properties[name]["gtype"], atom_properties[name]["dim"]
        )
    except ValueError:
        return []
    # this is per atom quantity - so get
    # number of atoms - first dimension
    val = list(val)
    dim = atom_properties[name]["dim"]
    if dim > 1:
        data = [val[x : x + dim] for x in range(0, len(val), dim)]
    else:
        data = list(val)
    return np.array(data)


def gather_atoms_subset(funct_args):
    # convert to ctypes
    name = str(funct_args[0])
    lenids = int(funct_args[1])
    ids = funct_args[2]

    # prep ids
    cids = (lenids * c_int)()
    for i in range(lenids):
        cids[i] = ids[i]

    if not name in atom_properties.keys():
        return []

    # this block prevents error when trying to access values
    # that do not exist
    try:
        val = job.gather_atoms_subset(
            name,
            atom_properties[name]["gtype"],
            atom_properties[name]["dim"],
            lenids,
            cids,
        )
    except ValueError:
        return []
    # this is per atom quantity - so get
    # number of atoms - first dimension
    val = list(val)
    dim = atom_properties[name]["dim"]
    if dim > 1:
        data = [val[x : x + dim] for x in range(0, len(val), dim)]
    else:
        data = list(val)
    return np.array(data)


def create_atoms(funct_args):
    # we have to process the input items
    # args are natoms, ids, type, x, v, image, shrinkexceed
    natoms = funct_args[0]
    ids = funct_args[1]
    type = funct_args[2]
    x = funct_args[3]
    v = funct_args[4]
    image = funct_args[5]
    shrinkexceed = funct_args[6]

    id_lmp = (c_int * natoms)()
    id_lmp[:] = ids

    type_lmp = (c_int * natoms)()
    type_lmp[:] = type

    image_lmp = (c_int * natoms)()
    image_lmp[:] = image

    args = [natoms, id_lmp, type_lmp, x, v, image_lmp, shrinkexceed]
    job.create_atoms(*args)
    return 1


def has_exceptions(funct_args):
    return job.has_exceptions


def has_gzip_support(funct_args):
    return job.has_gzip_support


def has_png_support(funct_args):
    return job.has_png_support


def has_jpeg_support(funct_args):
    return job.has_jpeg_support


def has_ffmpeg_support(funct_args):
    return job.has_ffmpeg_support


def installed_packages(funct_args):
    return job.installed_packages


def set_fix_external_callback(funct_args):
    job.set_fix_external_callback(*funct_args)


def get_neighlist(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.get_neighlist(*funct_args)


def find_pair_neighlist(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.find_pair_neighlist(*funct_args)


def find_fix_neighlist(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.find_fix_neighlist(*funct_args)


def find_compute_neighlist(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.find_compute_neighlist(*funct_args)


def get_neighlist_size(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.get_neighlist_size(*funct_args)


def get_neighlist_element_neighbors(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return job.get_neighlist_element_neighbors(*funct_args)


def get_thermo(funct_args):
    if MPI.COMM_WORLD.rank == 0:
        return np.array(job.get_thermo(*funct_args))


def scatter_atoms(funct_args):
    name = str(funct_args[0])
    py_vector = funct_args[1]
    # now see if its an integer or double type- but before flatten
    py_vector = np.array(py_vector).flatten()

    if atom_properties[name]["gtype"] == 0:
        c_vector = (len(py_vector) * c_int)(*py_vector)
    else:
        c_vector = (len(py_vector) * c_double)(*py_vector)

    job.scatter_atoms(
        name, atom_properties[name]["gtype"], atom_properties[name]["dim"], c_vector
    )
    return 1


def scatter_atoms_subset(funct_args):
    name = str(funct_args[0])
    lenids = int(funct_args[2])
    ids = funct_args[3]

    # prep ids
    cids = (lenids * c_int)()
    for i in range(lenids):
        cids[i] = ids[i]

    py_vector = funct_args[1]
    # now see if its an integer or double type- but before flatten
    py_vector = np.array(py_vector).flatten()

    if atom_properties[name]["gtype"] == 0:
        c_vector = (len(py_vector) * c_int)(*py_vector)
    else:
        c_vector = (len(py_vector) * c_double)(*py_vector)

    job.scatter_atoms_subset(
        name,
        atom_properties[name]["gtype"],
        atom_properties[name]["dim"],
        lenids,
        cids,
        c_vector,
    )
    return 1


def command(funct_args):
    job.command(funct_args)
    return 1


def select_cmd(argument):
    """
    Select a lammps command

    Args:
        argument (str): [close, extract_compute, get_thermo, scatter_atoms, command, gather_atoms]

    Returns:
        function: the selected function
    """
    switcher = {
        f.__name__: f
        for f in [
            extract_compute,
            get_version,
            get_file,
            commands_list,
            commands_string,
            extract_setting,
            extract_global,
            extract_box,
            extract_atom,
            extract_fix,
            extract_variable,
            get_natoms,
            set_variable,
            reset_box,
            gather_atoms_concat,
            gather_atoms_subset,
            scatter_atoms_subset,
            create_atoms,
            has_exceptions,
            has_gzip_support,
            has_png_support,
            has_jpeg_support,
            has_ffmpeg_support,
            installed_packages,
            set_fix_external_callback,
            get_neighlist,
            find_pair_neighlist,
            find_fix_neighlist,
            find_compute_neighlist,
            get_neighlist_size,
            get_neighlist_element_neighbors,
            get_thermo,
            scatter_atoms,
            command,
            gather_atoms,
        ]
    }
    return switcher.get(argument)


if __name__ == "__main__":
    while True:
        if MPI.COMM_WORLD.rank == 0:
            input_dict = pickle.load(sys.stdin.buffer)
            # with open('process.txt', 'a') as file:
            #     print('Input:', input_dict, file=file)
        else:
            input_dict = None
        input_dict = MPI.COMM_WORLD.bcast(input_dict, root=0)
        if input_dict["c"] == "close":
            job.close()
            break
        output = select_cmd(input_dict["c"])(input_dict["d"])
        if MPI.COMM_WORLD.rank == 0 and output is not None:
            # with open('process.txt', 'a') as file:
            #     print('Output:', output, file=file)
            pickle.dump(output, sys.stdout.buffer)
            sys.stdout.flush()
