# cython: language_level=3
"""
Utilities for manipulating coordinates or list of coordinates, under periodic
boundary conditions or otherwise.
"""

__author__ = "Will Richards"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Will Richards"
__email__ = "wmdrichards@gmail.com"
__date__ = "Nov 27, 2011"

import numpy as np

cimport cython
cimport numpy as np
from libc.math cimport fabs, floor, round
from libc.stdlib cimport free, malloc

np.import_array()

# Create images, 2D array of all length 3 combinations of [-1, 0, 1]
rng = np.arange(-1, 2, dtype=np.float64)
arange = rng[:, None] * np.array([1, 0, 0])[None, :]
brange = rng[:, None] * np.array([0, 1, 0])[None, :]
crange = rng[:, None] * np.array([0, 0, 1])[None, :]
images_t = arange[:, None, None] + brange[None, :, None] + \
    crange[None, None, :]
images = images_t.reshape((27, 3))

cdef np.float_t[:, ::1] images_view = images


# ---------------------------------------------------------------------------
# (N, 3) @ (3, 3) specialized matmul.
# The lattice matrix in pymatgen is always 3x3, so dot_2d / dot_2d_mod are
# called only in this shape. Hand-unrolling the inner two dims:
#   - holds the lattice in 9 scalar locals (registers)
#   - makes the i-loop unit-stride on both read and write
#   - lets the compiler auto-vectorize the outer N-loop
# This replaces the previous generic ijk loop, which had a non-unit-stride
# inner loop for row-major C arrays.
# ---------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
cdef void dot_2d_3x3(
    np.float_t[:, ::1] a,
    np.float_t[:, ::1] b,
    np.float_t[:, ::1] o,
) noexcept nogil:
    cdef int i, I = a.shape[0]
    cdef double a0, a1, a2
    cdef double b00 = b[0, 0], b01 = b[0, 1], b02 = b[0, 2]
    cdef double b10 = b[1, 0], b11 = b[1, 1], b12 = b[1, 2]
    cdef double b20 = b[2, 0], b21 = b[2, 1], b22 = b[2, 2]
    for i in range(I):
        a0 = a[i, 0]
        a1 = a[i, 1]
        a2 = a[i, 2]
        o[i, 0] = a0 * b00 + a1 * b10 + a2 * b20
        o[i, 1] = a0 * b01 + a1 * b11 + a2 * b21
        o[i, 2] = a0 * b02 + a1 * b12 + a2 * b22


@cython.boundscheck(False)
@cython.wraparound(False)
cdef void dot_2d_mod_3x3(
    np.float_t[:, ::1] a,
    np.float_t[:, ::1] b,
    np.float_t[:, ::1] o,
) noexcept nogil:
    """Same as dot_2d_3x3 but applies a mod-1 to each input frac coord first.

    Implemented as ``a - floor(a)`` (numpy ``np.mod(a, 1)`` semantics — always
    in [0, 1) for any real input) rather than ``a % 1``: the latter compiles to
    a slow Python-style helper on cdef double and has C-vs-Python sign
    differences for negative inputs.
    """
    cdef int i, I = a.shape[0]
    cdef double a0, a1, a2
    cdef double b00 = b[0, 0], b01 = b[0, 1], b02 = b[0, 2]
    cdef double b10 = b[1, 0], b11 = b[1, 1], b12 = b[1, 2]
    cdef double b20 = b[2, 0], b21 = b[2, 1], b22 = b[2, 2]
    for i in range(I):
        a0 = a[i, 0] - floor(a[i, 0])
        a1 = a[i, 1] - floor(a[i, 1])
        a2 = a[i, 2] - floor(a[i, 2])
        o[i, 0] = a0 * b00 + a1 * b10 + a2 * b20
        o[i, 1] = a0 * b01 + a1 * b11 + a2 * b21
        o[i, 2] = a0 * b02 + a1 * b12 + a2 * b22


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def pbc_shortest_vectors(lattice, fcoords1, fcoords2, mask=None, return_d2=False, lll_frac_tol=None):
    """Get the shortest vectors between two lists of coordinates taking into
    account periodic boundary conditions and the lattice.

    Args:
        lattice: lattice to use
        fcoords1: First set of fractional coordinates. e.g., [0.5, 0.6, 0.7]
            or [[1.1, 1.2, 4.3], [0.5, 0.6, 0.7]]. Must be np.float64
        fcoords2: Second set of fractional coordinates.
        mask (int_ array): Mask of matches that are not allowed.
            i.e. if mask[1,2] == True, then subset[1] cannot be matched
            to superset[2]
        lll_frac_tol (float_ array of length 3): Fractional tolerance (per LLL lattice vector) over which
            the calculation of minimum vectors will be skipped.
            Can speed up calculation considerably for large structures.

    Returns:
        np.array: of displacement vectors from fcoords1 to fcoords2
            first index is fcoords1 index, second is fcoords2 index
    """

    # Ensure correct shape
    fcoords1, fcoords2 = np.atleast_2d(fcoords1, fcoords2)

    pbc = lattice.pbc
    cdef int n_pbc = sum(pbc)
    cdef int n_pbc_im = int(3 ** n_pbc)

    cdef np.float_t[:, ::1] frac_im
    cdef int i, j, k, l, I, J, bestk = 0

    if n_pbc == 3:
        fcoords1 = lattice.get_lll_frac_coords(fcoords1)
        fcoords2 = lattice.get_lll_frac_coords(fcoords2)
        matrix = lattice.lll_matrix
        frac_im = images_view
    else:
        frac_im = <np.float_t[:n_pbc_im, :3]> malloc(3 * n_pbc_im * sizeof(np.float_t))
        matrix = lattice.matrix.copy()
        k = 0
        for i in range(27):
            for j in range(3):
                if not pbc[j] and images_view[i, j] != 0:
                    break
            else:
                frac_im[k] = images_view[i]
                k += 1

    cdef np.float_t[:, ::1] lat = np.asarray(matrix, dtype=np.float64, order="C")

    I = len(fcoords1)
    J = len(fcoords2)

    cdef np.float_t[:, ::1] fc1 = fcoords1
    cdef np.float_t[:, ::1] fc2 = fcoords2

    cdef np.float_t[:, ::1] cart_f1 = <np.float_t[:I, :3]> malloc(3 * I * sizeof(np.float_t))
    cdef np.float_t[:, ::1] cart_f2 = <np.float_t[:J, :3]> malloc(3 * J * sizeof(np.float_t))
    cdef np.float_t[:, ::1] cart_im = <np.float_t[:n_pbc_im, :3]> malloc(3 * n_pbc_im * sizeof(np.float_t))

    cdef bint has_mask = mask is not None
    cdef np.int64_t[:, :] mask_arr
    if has_mask:
        mask_arr = np.asarray(mask, dtype=np.int64, order="C")

    cdef bint has_ftol = (lll_frac_tol is not None)
    cdef np.float_t[:] ftol
    if has_ftol:
        ftol = np.asarray(lll_frac_tol, dtype=np.float64, order="C")


    dot_2d_mod_3x3(fc1, lat, cart_f1)
    dot_2d_mod_3x3(fc2, lat, cart_f2)
    dot_2d_3x3(frac_im, lat, cart_im)

    vectors = np.empty((I, J, 3))
    d2 = np.empty((I, J))
    cdef np.float_t[:, :, ::1] vs = vectors
    cdef np.float_t[:, ::1] ds = d2
    cdef np.float_t best, d, da, db, dc, fdist
    cdef bint within_frac = True
    # 3-element scratch buffer for the per-(i,j) preimage. Heap-alloc just for
    # 24 bytes is wasteful and the memoryview indirection prevents the compiler
    # from holding these in registers; a stack array is faster and simpler.
    cdef double pre_im[3]

    for i in range(I):
        for j in range(J):
            within_frac = False
            if (not has_mask) or (mask_arr[i, j] == 0):
                within_frac = True
                if has_ftol:
                    for l in range(3):
                        fdist = fc2[j, l] - fc1[i, l]
                        if fabs(fdist - round(fdist)) > ftol[l]:
                            within_frac = False
                            break
                if within_frac:
                    pre_im[0] = cart_f2[j, 0] - cart_f1[i, 0]
                    pre_im[1] = cart_f2[j, 1] - cart_f1[i, 1]
                    pre_im[2] = cart_f2[j, 2] - cart_f1[i, 2]
                    best = 1e100
                    bestk = 0
                    for k in range(n_pbc_im):
                        # compilers have a hard time unrolling this
                        da = pre_im[0] + cart_im[k, 0]
                        db = pre_im[1] + cart_im[k, 1]
                        dc = pre_im[2] + cart_im[k, 2]
                        d = da * da + db * db + dc * dc
                        if d < best:
                            best = d
                            bestk = k
                    ds[i, j] = best
                    vs[i, j, 0] = pre_im[0] + cart_im[bestk, 0]
                    vs[i, j, 1] = pre_im[1] + cart_im[bestk, 1]
                    vs[i, j, 2] = pre_im[2] + cart_im[bestk, 2]
            if not within_frac:
                ds[i, j] = 1e20
                vs[i, j, 0] = 1e20
                vs[i, j, 1] = 1e20
                vs[i, j, 2] = 1e20

    if not (n_pbc == 3):
        free(&frac_im[0,0])
    free(&cart_f1[0,0])
    free(&cart_f2[0,0])
    free(&cart_im[0,0])

    if return_d2:
        return vectors, d2
    else:
        return vectors

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def is_coord_subset_pbc(subset, superset, atol, mask, pbc=(True, True, True)):
    """
    Tests if all fractional coords in subset are contained in superset.
    Allows specification of a mask determining pairs that are not
    allowed to match to each other.

    Args:
        subset, superset: List of fractional coords.
        pbc: a tuple defining the periodic boundary conditions along the three
            axis of the lattice.

    Returns:
        True if all of subset is in superset.
    """

    cdef np.float_t[:, :] fc1 = subset
    cdef np.float_t[:, :] fc2 = superset
    cdef np.float_t[:] t = atol
    cdef np.int64_t[:, :] m = np.asarray(mask, dtype=np.int64, order="C")

    cdef int i, j, k, len_fc1, len_fc2
    cdef np.float_t d
    cdef bint ok
    cdef bint pbc_int[3]

    pbc_int = pbc

    len_fc1 = fc1.shape[0]
    len_fc2 = fc2.shape[0]

    for i in range(len_fc1):
        ok = False
        for j in range(len_fc2):
            if m[i, j]:
                continue
            ok = True
            for k in range(3):
                d = fc1[i, k] - fc2[j, k]
                if fabs(d - round(d) * pbc_int[k]) > t[k]:
                    ok = False
                    break
            if ok:
                break
        if not ok:
            break
    return bool(ok)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def coord_list_mapping_pbc(subset, superset, atol=1e-8, pbc=(True, True, True)):
    """
    Gives the index mapping from a subset to a superset.
    Superset cannot contain duplicate matching rows.

    Args:
        subset, superset: List of frac_coords.
        pbc: a tuple defining the periodic boundary conditions along the three
            axis of the lattice.

    Returns:
        list of indices such that superset[indices] = subset
    """
    inds = -np.ones(len(subset), dtype=np.int64)
    subset = np.atleast_2d(subset)
    superset = np.atleast_2d(superset)

    cdef np.float_t[:, :] fc1 = subset
    cdef np.float_t[:, :] fc2 = superset
    cdef np.float_t[:] t = atol
    cdef np.int64_t[:] c_inds = inds
    cdef np.float_t d
    cdef bint ok_inner, ok_outer
    cdef bint pbc_int[3]

    pbc_int = pbc

    len_fc1 = fc1.shape[0]
    len_fc2 = fc2.shape[0]

    for i in range(len_fc1):
        ok_outer = False
        for j in range(len_fc2):
            ok_inner = True
            for k in range(3):
                d = fc1[i, k] - fc2[j, k]
                if fabs(d - round(d) * pbc_int[k]) > t[k]:
                    ok_inner = False
                    break
            if ok_inner:
                if c_inds[i] >= 0:
                    raise ValueError("Something wrong with the inputs, likely duplicates in superset")
                c_inds[i] = j
                ok_outer = True
                # We don't break here so we can check for duplicates in superset
        if not ok_outer:
            break

    if not ok_outer:
        raise ValueError("not a subset of superset")

    return inds
