import logging
import numpy as np
import ase

logger = logging.getLogger('DockOnSurf')


def get_vect_angle(v1, v2, degrees=False):
    """Computes the angle between two vectors.

    @param v1: The first vector.
    @param v2: The second vector.
    @param degrees: Whether the result should be in radians (True) or in
        degrees (False).
    @return: The angle in radians if degrees = False, or in degrees if
        degrees =True
    """
    v1_u = v1 / np.linalg.norm(v1)
    v2_u = v2 / np.linalg.norm(v2)
    angle = np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
    return angle if not degrees else angle * 180 / np.pi


def vect_avg(vects):
    """Computes the element-wise mean of a set of vectors.

    @param vects: list of lists-like: containing the vectors (num_vectors,
        length_vector).
    @return: vector average computed doing the element-wise mean.
    """
    from utilities import try_command
    err = "vect_avg parameter vects must be a list-like, able to be converted" \
          " np.array"
    array = try_command(np.array, [(ValueError, err)], vects)
    if len(array.shape) == 1:
        return array
    else:
        num_vects = array.shape[1]
        return np.array([np.average(array[:, i]) for i in range(num_vects)])


def get_atom_coords(atoms: ase.Atoms, ctrs_list):
    """Gets the coordinates of the specified atoms from a ase.Atoms object.

    Given an ase.Atoms object and a list of atom indices specified in ctrs_list
    it gets the coordinates of the specified atoms. If the element in the
    ctrs_list is not an index but yet a list of indices, it computes the
    element-wise mean of the coordinates of the atoms specified in the inner
    list.
    @param atoms: ase.Atoms object for which to obtain the coordinates of.
    @param ctrs_list: list of (indices/list of indices) of the atoms for which
                      the coordinates should be extracted.
    @return: np.ndarray of atomic coordinates.
    """
    coords = []
    for i, elem in enumerate(ctrs_list):
        if isinstance(elem, list):
            coords.append(vect_avg(np.array([atoms[c].position for c in elem])))
        elif isinstance(elem, int):
            coords.append(atoms[elem].position)
        else:
            err = f"'ctrs_list must be a list of ints or lists, {type(elem)} " \
                  "found."
            logger.error(err)
            raise ValueError
    return np.array(coords)


def add_adsorbate(slab, adsorbate, site_coord, ctr_coord, height, offset=None,
                  norm_vect=(0, 0, 1)):
    """Add an adsorbate to a surface.

    This function extends the functionality of ase.build.add_adsorbate
    (https://wiki.fysik.dtu.dk/ase/ase/build/surface.html#ase.build.add_adsorbate)
    by enabling to change the z coordinate and the axis perpendicular to the
    surface.
    @param slab: ase.Atoms object containing the coordinates of the surface
    @param adsorbate: ase.Atoms object containing the coordinates of the
        adsorbate.
    @param site_coord: The coordinates of the adsorption site on the surface.
    @param ctr_coord: The coordinates of the adsorption center in the molecule.
    @param height: The height above the surface where to adsorb.
    @param offset: Offsets the adsorbate by a number of unit cells. Mostly
        useful when adding more than one adsorbate.
    @param norm_vect: The vector perpendicular to the surface.
    """
    info = slab.info.get('adsorbate_info', {})
    pos = np.array([0.0, 0.0, 0.0])  # part of absolute coordinates
    spos = np.array([0.0, 0.0, 0.0])  # part relative to unit cell
    norm_vect_u = np.array(norm_vect) / np.linalg.norm(norm_vect)
    if offset is not None:
        spos += np.asarray(offset, float)
    if isinstance(site_coord, str):
        # A site-name:
        if 'sites' not in info:
            raise TypeError('If the atoms are not made by an ase.build '
                            'function, position cannot be a name.')
        if site_coord not in info['sites']:
            raise TypeError('Adsorption site %s not supported.' % site_coord)
        spos += info['sites'][site_coord]
    else:
        pos += site_coord
    if 'cell' in info:
        cell = info['cell']
    else:
        cell = slab.get_cell()
    pos += np.dot(spos, cell)
    # Convert the adsorbate to an Atoms object
    if isinstance(adsorbate, ase.Atoms):
        ads = adsorbate
    elif isinstance(adsorbate, ase.Atom):
        ads = ase.Atoms([adsorbate])
    else:
        # Assume it is a string representing a single Atom
        ads = ase.Atoms([ase.Atom(adsorbate)])
    pos += height * norm_vect_u
    # Move adsorbate into position
    ads.translate(pos - ctr_coord)
    # Attach the adsorbate
    slab.extend(ads)


def ads_euler(molec, surf, site, ctr, pts_angle, neigh_ctr, norm_vect):
    from random import random
    plane_vect = np.cross(norm_vect, [random() for i in range(3)])
    add_adsorbate()

    return


def ads_chemcat(site, ctr, pts_angle):
    return "TO IMPLEMENT"


def adsorb_confs(conf_list, surf, ads_ctrs, sites, algo, num_pts, neigh_ctrs,
                 norm_vect, coll_bttm):
    """Generates a number of adsorbate-surface structure coordinates.

    Given a list of conformers, a surface, a list of atom indices (or list of
    list of indices) of both the surface and the adsorbate, it generates a
    number of adsorbate-surface structures for every possible combination of
    them at different orientations.
    @param conf_list: list of ase.Atoms of the different conformers
    @param surf: the ase.Atoms object of the surface
    @param ads_ctrs: the list atom indices of the adsorbate.
    @param sites: the list of atom indices of the surface.
    @param algo: the algorithm to use for the generation of adsorbates.
    @param num_pts: the number of points per angle orientation to sample
    @param neigh_ctrs: the indices of the neighboring atoms to the adsorption
    atoms.
    @param norm_vect: The vector perpendicular to the surface.
    @param coll_bttm: The lowermost height for which to detect a collision
    @return: list of ase.Atoms for the adsorbate-surface structures
    """
    surf_ads_list = []
    sites_coords = get_atom_coords(surf, sites)
    for conf in conf_list:
        molec_ctr_coords = get_atom_coords(conf, ads_ctrs)
        molec_neigh_coords = get_atom_coords(conf, neigh_ctrs)
        for site in sites_coords:
            for i, molec_ctr in enumerate(molec_ctr_coords):
                if algo == 'euler':
                    surf_ads_list.extend(ads_euler(conf, surf, molec_ctr, site,
                                                   num_pts, norm_vect,
                                                   coll_bttm,
                                                   molec_neigh_coords[i]))
                elif algo == 'chemcat':
                    surf_ads_list.extend(ads_chemcat(site, molec_ctr, num_pts))
    return surf_ads_list


def run_screening(inp_vars):
    """Carry out the screening of adsorbate coordinates on a surface

    @param inp_vars: Calculation parameters from input file.
    """
    import os
    from modules.formats import read_coords, read_energies, \
        rdkit_mol_to_ase_atoms, adapt_format
    from modules.clustering import get_rmsd, clustering
    from modules.isolated import get_moments_of_inertia
    from modules.calculation import run_calc

    if not os.path.isdir("isolated"):
        err = "'isolated' directory not found. It is needed in order to carry "
        "out the screening of structures to be adsorbed"
        logger.error(err)
        raise ValueError(err)

    conf_list = read_coords(inp_vars['code'], 'isolated', 'rdkit')
    # TODO Implement neighbors algorithm / align molecule to surface
    # neigh_list = get_neighbors(conf_list[0], inp_vars['molec_ads_ctrs'])
    conf_enrgs = read_energies(inp_vars['code'], 'isolated')
    mois = np.array([get_moments_of_inertia(conf)[0] for conf in conf_list])
    rmsd_mtx = get_rmsd(conf_list)
    exemplars = clustering(rmsd_mtx)  # TODO Kind of Clustering
    clust_conf_list = [conf_list[idx] for idx in exemplars]
    conf_list = [rdkit_mol_to_ase_atoms(conf) for conf in clust_conf_list]
    surf = adapt_format('ase', inp_vars['surf_file'])
    surf_ads_list = adsorb_confs(conf_list, surf, inp_vars['molec_ads_ctrs'],
                                 inp_vars['sites'], inp_vars['ads_algo'],
                                 inp_vars['sample_points_per_angle'],
                                 inp_vars['molec_neigh_ctrs'],
                                 inp_vars['surf_norm_vect'],
                                 inp_vars['collision_bottom'])
    run_calc('screening', inp_vars, surf_ads_list)