# -*- coding: utf-8 -*-
"""Copyright 2019 DScribe developers

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import numpy as np

import scipy.sparse
from scipy.spatial.distance import cdist

from ase import Atoms


def get_adjacency_matrix(radius, pos1, pos2=None, output_type="coo_matrix"):
    """Calculates a sparse adjacency matrix by only considering distances
    within a certain cutoff. Uses a k-d tree to reach O(n log(N)) time
    complexity.

    Args:
        radius(float): The cutoff radius within which distances are
            calculated. Distances outside this radius are not included.
        pos1(np.ndarray): A list of N-dimensional positions.
        pos2(np.ndarray): A list of N-dimensional positions. If not provided,
            is assumed to be the same as pos1.
        output_type(str): Which container to use for output data. Options:
            "dok_matrix", "coo_matrix", "dict", or "ndarray". Default:
            "dok_matrix".

    Returns:
        dok_matrix | np.array | coo_matrix | dict: Symmetric sparse 2D
        matrix containing the pairwise distances.
    """
    tree1 = scipy.spatial.cKDTree(
        pos1,
        leafsize=16,
        compact_nodes=True,
        copy_data=False,
        balanced_tree=True,
        boxsize=None,
    )
    if pos2 is None:
        pos2 = pos1
    tree2 = scipy.spatial.cKDTree(
        pos2,
        leafsize=16,
        compact_nodes=True,
        copy_data=False,
        balanced_tree=True,
        boxsize=None,
    )
    dmat = tree1.sparse_distance_matrix(tree2, radius, output_type=output_type)

    return dmat


def get_adjacency_list(adjacency_matrix):
    """Used to transform an adjacency matrix into an adjacency list. The
    adjacency list provides much faster access to the neighbours of a node.

    Args:
        adjacency_matrix(scipy.sparse.spmatrix): The adjacency matrix from
            which the adjacency list is constructed from. Any of the scipy
            sparse matrix classes.

    Returns:
        list: A list of neighbouring indices. The list of neighbouring indices
        for atom at index i is given by accessing the ith element of this list.
    """
    # Ensure that we have a coo-matrix
    if type(adjacency_matrix) != scipy.sparse.coo_matrix:
        adjacency_matrix = adjacency_matrix.tocoo()

    # Build adjacency list
    adjacency_list = [[] for i in range(adjacency_matrix.shape[0])]
    for i, j in zip(adjacency_matrix.row, adjacency_matrix.col):
        adjacency_list[i].append(j)

    return adjacency_list


def get_extended_system(system, radial_cutoff, centers=None, return_cell_indices=False):
    """Used to create a periodically extended system. If centers are not
    specified, simply takes returns the original system multiplied by an
    integer amount of times in each direction to cover the radial cutoff. If
    centers are provided, returns the exact atoms that are within the given
    radial cutoff from the given centers.

    Args:
        original_system (ase.Atoms): The original periodic system to duplicate.
        radial_cutoff (float): The radial cutoff to use in constructing the
            extended system.
        centers (np.ndarray): Array of xyz-coordinates from which the distance
            is calculated. If provided, these centers are used to calculate the
            exact distance and only atoms within the radial cutoff from these
            centers are returned.
        return_cell_indices (boolean): Whether to return an array of cell
            indices for each atom in the extended system.

    Returns:
        ase.Atoms | tuple: If return_cell_indices is False, returns the new
        extended system. Else returns a tuple containing the new extended
        system as the first entry and the index of the periodically repeated
        cell for each atom as the second entry.
    """
    # Determine the upper limit of how many copies we need in each cell vector
    # direction. We take as many copies as needed to reach the radial cutoff.
    # Notice that we need to use vectors that are perpendicular to the cell
    # vectors to ensure that the correct atoms are included for non-cubic cells.
    cell = np.asarray(system.get_cell())
    pbc = system.get_pbc()
    a1, a2, a3 = cell[0], cell[1], cell[2]
    b1 = np.cross(a2, a3, axis=0)
    b2 = np.cross(a3, a1, axis=0)
    b3 = np.cross(a1, a2, axis=0)
    p1 = np.dot(a1, b1) / np.dot(b1, b1) * b1  # Projections onto perpendicular vectors
    p2 = np.dot(a2, b2) / np.dot(b2, b2) * b2
    p3 = np.dot(a3, b3) / np.dot(b3, b3) * b3
    xyz_arr = np.linalg.norm(np.array([p1, p2, p3]), axis=1)
    cell_images = np.ceil(radial_cutoff / xyz_arr)
    nx = int(cell_images[0]) if pbc[0] else 0
    ny = int(cell_images[1]) if pbc[1] else 0
    nz = int(cell_images[2]) if pbc[2] else 0
    n_copies_axis = np.array([nx, ny, nz], dtype=int)

    # If no centers are given, and the cell indices are not requested, simply
    # return the multiplied system. This is much faster.
    if centers is None and not return_cell_indices:
        n_atoms = len(system)
        n_rep = np.product(2 * n_copies_axis + 1)  # Number of repeated copies
        ext_pos = np.tile(system.get_positions(), (n_rep, 1))

        # Calculate the extended system positions so that the original cell
        # stays in place: both in space and in index
        i_curr = 0
        for m0 in np.append(np.arange(0, nx + 1), np.arange(-nx, 0)):
            for m1 in np.append(np.arange(0, ny + 1), np.arange(-ny, 0)):
                for m2 in np.append(np.arange(0, nz + 1), np.arange(-nz, 0)):
                    ext_pos[i_curr : i_curr + n_atoms] += np.dot((m0, m1, m2), cell)
                    i_curr += n_atoms

        ext_symbols = np.tile(system.get_atomic_numbers(), n_rep)
        extended_system = Atoms(
            positions=ext_pos,
            symbols=ext_symbols,
        )

        return extended_system

    # If centers are given and/or cell indices are needed, the process is done
    # one cell at a time to keep track of the cell inded and the distances.
    # This is a bit slower.
    else:
        # We need to specify that the relative positions should not be wrapped.
        # Otherwise the repeated systems may overlap with the positions taken
        # with get_positions()
        relative_pos = np.array(system.get_scaled_positions(wrap=False))
        numbers = np.array(system.numbers)
        cartesian_pos = np.array(system.get_positions())

        # Create copies of the cell but keep track of the atoms in the
        # original cell
        num_extended = []
        pos_extended = []
        num_extended.append(numbers)
        pos_extended.append(cartesian_pos)
        a = np.array([1, 0, 0])
        b = np.array([0, 1, 0])
        c = np.array([0, 0, 1])
        cell_indices = [np.zeros((len(system), 3), dtype=int)]

        for i in range(-n_copies_axis[0], n_copies_axis[0] + 1):
            for j in range(-n_copies_axis[1], n_copies_axis[1] + 1):
                for k in range(-n_copies_axis[2], n_copies_axis[2] + 1):
                    if i == 0 and j == 0 and k == 0:
                        continue

                    # Calculate the positions of the copied atoms and filter
                    # out the atoms that are farther away than the given
                    # cutoff.
                    num_copy = np.array(numbers)
                    pos_copy = np.array(relative_pos)

                    pos_shifted = pos_copy - i * a - j * b - k * c
                    pos_copy_cartesian = np.dot(pos_shifted, cell)

                    # Only distances to the atoms within the interaction limit
                    # are considered.
                    distances = cdist(pos_copy_cartesian, centers)
                    weight_mask = distances < radial_cutoff

                    # Create a boolean mask that says if the atom is within the
                    # range from at least one atom in the original cell
                    valids_mask = np.any(weight_mask, axis=1)

                    if np.any(valids_mask):
                        valid_pos = pos_copy_cartesian[valids_mask]
                        valid_num = num_copy[valids_mask]
                        valid_ind = np.tile(
                            np.array([i, j, k], dtype=int), (len(valid_num), 1)
                        )

                        pos_extended.append(valid_pos)
                        num_extended.append(valid_num)
                        cell_indices.append(valid_ind)

        pos_extended = np.concatenate(pos_extended)
        num_extended = np.concatenate(num_extended)
        cell_indices = np.vstack(cell_indices)

        extended_system = Atoms(
            positions=pos_extended, numbers=num_extended, cell=cell, pbc=False
        )

        return extended_system, cell_indices
