# cython: language_level=3
# cython: warn.undeclared=True
# cython: warn.unreachable=True
# cython: warn.maybe_uninitialized=True
# cython: warn.unused=True
# cython: warn.unused_arg=True
# cython: warn.unused_result=True
# cython: warn.multiple_declarators=True

"""
NOTE: Although the functions defined in this module has a LONG and UGLY list of
parameters, they are transparent, loosely coupled and easy to debug. Refactoring
them in OOP flavor does not make much profit, but makes them hard to understand,
strongly coupled and bug-prone. We got this lesson after wasting a weekend trying
to rewrite it in OOP.

Also, DO NOT replace the list of parameters with a struct or object with public
attributes, since it will hide the actual parameters in use.
"""
import cython
from libc.stdint cimport int32_t, int64_t, uint64_t
import numpy as np


#------------------------------------------------------------------------------#
#                         C++ stuff exposed to Cython                          #
#------------------------------------------------------------------------------#
cdef extern from "orbital_map.h" namespace "orbital_map":
    cdef cppclass OrbitalMap[cell_index_t]:
        OrbitalMap() except +
        OrbitalMap(
            const cell_index_t* dim,
            const cell_index_t& num_orb_pc) except +
        OrbitalMap(
            const cell_index_t* dim,
            const cell_index_t& num_orb_pc,
            const int64_t* vac_id_sc,
            const int64_t& num_vac) except +
        int64_t id_pc2sc(const cell_index_t* id_pc) const
        int64_t id_pc2sc_vac(const cell_index_t* id_pc) const
        int64_t id_pc2sc_vac2(const cell_index_t* id_pc) const

    # Functions
    void check_id_sc_array_cxx(
        const int64_t& num_orb_sc,
        const int64_t& num_id,
        const int64_t* id_sc_array,
        int32_t* status)

    void check_id_pc_array_cxx[cell_index_t](
        const cell_index_t* dim,
        const cell_index_t& num_orb_pc,
        const int64_t& num_vac,
        const int64_t* vac_id_sc,
        const int64_t& num_id,
        const cell_index_t* id_pc_array,
        int32_t* status)

    void id_sc2pc_array_cxx[cell_index_t](
        const cell_index_t* orb_id_pc,
        const int64_t& num_id,
        const int64_t* id_sc_array,
        cell_index_t* id_pc_array)

    void id_pc2sc_array_cxx[cell_index_t](
        const cell_index_t* dim,
        const cell_index_t& num_orb_pc,
        const int64_t& num_vac,
        const int64_t* vac_id_sc,
        const int64_t& num_id,
        const cell_index_t* id_pc_array,
        int64_t* id_sc_array)


#-------------------------------------------------------------------------------
#      Functions for converting index between pc and sc representations
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def id_pc2sc(
    int32_t [::1] dim,
    int32_t num_orb_pc,
    int32_t [::1] id_pc,
    int64_t [::1] vac_id_sc):
    """
    Convert orbital index from pc to sc representation.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    id_pc: (4,) int32 array
        orbital index in primitive cell representation
    vac_id_sc: (num_vac,) int64 array
        indices of vacancies in super cell representation

    Returns
    ------
    id_sc: int64_t
        orbital index in super representation
    """
    cdef int64_t num_vac = vac_id_sc.shape[0]
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc, &vac_id_sc[0], num_vac)
    if num_vac == 0:
        return orb_map.id_pc2sc(&id_pc[0])
    else:
        return orb_map.id_pc2sc_vac2(&id_pc[0])


@cython.boundscheck(False)
@cython.wraparound(False)
def check_id_sc_array(
    int64_t num_orb_sc,
    int64_t [::1] id_sc_array):
    """
    Check for errors in id_sc_array and return information of the 0th
    encountered error.

    Parameters
    ----------
    num_orb_sc: int64
        number of orbitals in the super cell
    id_sc_array (num_id,) int64 array
        orbital indices in sc representation

    Returns
    -------
    status: (2,) int32 array
        0th element is the error type:
            -1: IDSCIndexError
             0: NO ERROR
        1th element is the index of 0th illegal id_sc
    """
    cdef int32_t [::1] status = np.zeros(2, dtype=np.int32)
    cdef int64_t num_id = id_sc_array.shape[0]
    check_id_sc_array_cxx(
        num_orb_sc,
        num_id,
        &id_sc_array[0],
        &status[0])
    return np.asarray(status)


@cython.boundscheck(False)
@cython.wraparound(False)
def check_id_pc_array(
    int32_t [::1] dim,
    int32_t num_orb_pc,
    int32_t [:,::1] id_pc_array,
    int64_t [::1] vac_id_sc):
    """
    Check for errors in id_pc_array and return information of the 0th
    encountered error.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    id_pc_array: (num_id, 4) int32 array
        orbital indices in primitive cell representation
    vac_id_sc: (num_vac,) int64 array
        indices of vacancies in super cell representation

    Returns
    -------
    status: (3,) int32 array
        0th element is the error type:
            -2: IDPCIndexError
            -1: IDPCVacError
             0: NO ERROR
        1st element is the index of 0th illegal id_pc
        2nd element is the index of wrong element in 0th illegal id_pc
    """
    cdef int32_t [::1] status= np.zeros(3, dtype=np.int32)
    cdef int64_t num_id = id_pc_array.shape[0]
    cdef int64_t num_vac = vac_id_sc.shape[0]
    check_id_pc_array_cxx(
        &dim[0],
        num_orb_pc,
        num_vac,
        &vac_id_sc[0],
        num_id,
        &id_pc_array[0][0],
        &status[0])
    return np.asarray(status)


@cython.boundscheck(False)
@cython.wraparound(False)
def id_sc2pc_array(
    int32_t [:,::1] orb_id_pc,
    int64_t [::1] id_sc_array):
    """
    Convert an array of orbital indices from sc representation to pc
    representation.

    Parameters
    ----------
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of all the orbitals in super cell in pc representation
    id_sc_array (num_id,) int64 array
        orbital indices in sc representation

    Returns
    -------
    id_pc_array: (num_id, 4) int32 array
        orbital indices in primitive cell representation
    """
    # Allocate arrays
    cdef int64_t num_id= id_sc_array.shape[0]
    cdef int32_t [:,::1] id_pc_array = np.zeros((num_id, 4), dtype=np.int32)

    # Convert orbital indices from sc to pc representation
    id_sc2pc_array_cxx(
        &orb_id_pc[0][0],
        num_id,
        &id_sc_array[0],
        &id_pc_array[0][0])
    return np.asarray(id_pc_array)


@cython.boundscheck(False)
@cython.wraparound(False)
def id_pc2sc_array(
    int32_t [::1] dim,
    int32_t num_orb_pc,
    int32_t [:,::1] id_pc_array,
    int64_t [::1] vac_id_sc):
    """
    Convert an array of orbital indices from pc representation to sc
    representation in presence of vacancies.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    id_pc_array: (num_id, 4) int32 array
        orbital indices in primitive cell representation
    vac_id_sc: (num_vac,) int64 array
        indices of vacancies in super cell representation

    Returns
    -------
    id_sc_array: (num_id,) int64 array
        orbital indices in super cell representation
    """
    # Allocate arrays
    cdef int64_t num_vac = vac_id_sc.shape[0]
    cdef int64_t num_id = id_pc_array.shape[0]
    cdef int64_t [::1] id_sc_array = np.zeros(num_id, dtype=np.int64)

    # Convert orbital indices from pc to sc representation
    id_pc2sc_array_cxx(
        &dim[0],
        num_orb_pc,
        num_vac,
        &vac_id_sc[0],
        num_id,
        &id_pc_array[0][0],
        &id_sc_array[0])
    return np.asarray(id_sc_array)


#-------------------------------------------------------------------------------
#                        Functions for testing purposes
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def test_acc_sps(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc):
    """
    In this test, we convert the orbital index as sc->pc->sc, to check if the
    input is restored.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    result: int64
        total absolute error
    """
    cdef int64_t id_sc, id_sc2
    cdef int64_t result
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc)

    result = 0
    for id_sc in range(orb_id_pc.shape[0]):
        id_sc2 = orb_map.id_pc2sc(&orb_id_pc[id_sc][0])
        result += abs(id_sc - id_sc2)
    return result


@cython.boundscheck(False)
@cython.wraparound(False)
def test_acc_psp(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc):
    """
    In this test, we convert the orbital index as pc->sc->pc, to check if the
    input is restored.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    result: int64
        total absolute error
    """
    cdef int32_t [::1] id_pc, id_pc2
    cdef int64_t id_sc
    cdef int64_t result
    cdef int32_t i
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc)

    result = 0
    for id_sc in range(orb_id_pc.shape[0]):
        id_pc = orb_id_pc[id_sc]
        id_pc2 = orb_id_pc[orb_map.id_pc2sc(&id_pc[0])]
        for i in range(4):
            result += abs(id_pc[i] - id_pc2[i])
    return result


@cython.boundscheck(False)
@cython.wraparound(False)
def test_acc_sps_vac(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc,
                     int64_t [::1] vac_id_sc):
    """
    In this test, we convert the orbital index as sc->pc->sc, to check if the
    input is restored, when there are vacancies.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation
    vac_id_sc: (num_orb_sc,) int64 array
        indices of vacancies of super cell in sc representation

    Returns
    -------
    result: int64
        total absolute error
    """
    cdef int64_t id_sc, id_sc2
    cdef int64_t result
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc, &vac_id_sc[0], vac_id_sc.shape[0])

    result = 0
    for id_sc in range(orb_id_pc.shape[0]):
        id_sc2 = orb_map.id_pc2sc_vac(&orb_id_pc[id_sc][0])
        result += abs(id_sc - id_sc2)
        id_sc2 = orb_map.id_pc2sc_vac2(&orb_id_pc[id_sc][0])
        result += abs(id_sc - id_sc2)
    return result


@cython.boundscheck(False)
@cython.wraparound(False)
def test_acc_psp_vac(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc,
                     int64_t [::1] vac_id_sc):
    """
    In this test, we convert the orbital index as pc->sc->pc, to check if the
    input is restored, when there are vacancies.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation
    vac_id_sc: (num_orb_sc,) int64 array
        indices of vacancies of super cell in sc representation

    Returns
    -------
    result: int64
        total absolute error
    """
    cdef int32_t [::1] id_pc, id_pc2
    cdef int64_t id_sc
    cdef int64_t result
    cdef int32_t i
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc, &vac_id_sc[0], vac_id_sc.shape[0])

    result = 0
    for id_sc in range(orb_id_pc.shape[0]):
        id_pc = orb_id_pc[id_sc]
        id_pc2 = orb_id_pc[orb_map.id_pc2sc_vac(&id_pc[0])]
        for i in range(4):
            result += abs(id_pc[i] - id_pc2[i])
        id_pc2 = orb_id_pc[orb_map.id_pc2sc_vac2(&id_pc[0])]
    return result


@cython.boundscheck(False)
@cython.wraparound(False)
def test_speed_pc2sc(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc):
    """
    Test the speed of OrbitalMap.id_pc2sc.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    None.
    """
    cdef int64_t i, id_sc
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc)
    for i in range(orb_id_pc.shape[0]):
        id_sc = orb_map.id_pc2sc(&orb_id_pc[i][0])


@cython.boundscheck(False)
@cython.wraparound(False)
def test_speed_pc2sc_vac(int32_t [::1] dim, int32_t num_orb_pc, int32_t [:,::1] orb_id_pc,
                   int64_t [::1] vac_id_sc):
    """
    Test the speed of OrbitalMap.id_pc2sc_vac.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation
    vac_id_sc: (num_orb_sc,) int64 array
        indices of vacancies of super cell in sc representation

    Returns
    -------
    None.
    """
    cdef int64_t i, id_sc
    cdef OrbitalMap[int32_t] orb_map = OrbitalMap[int32_t](&dim[0], num_orb_pc, &vac_id_sc[0], vac_id_sc.shape[0])
    for i in range(orb_id_pc.shape[0]):
        id_sc = orb_map.id_pc2sc_vac(&orb_id_pc[i][0])


@cython.boundscheck(False)
@cython.wraparound(False)
def test_speed_sc2pc(int32_t [:,::1] orb_id_pc):
    """
    Test the speed of OrbitalMap.id_sc2pc.

    Parameters
    ----------
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    None.
    """
    cdef int64_t id_sc
    cdef int32_t [::1] id_pc
    for id_sc in range(orb_id_pc.shape[0]):
        id_pc = orb_id_pc[id_sc]
