# 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 "hop_dict.h" namespace "hop_dict":
    cdef cppclass HopDict:
        HopDict() except +
        void add_hop_origin(
            const int64_t& num_hop,
            const int64_t* hop_ind,
            const double complex* hop_eng,
            const double* dr)
        void add_hop_modifier(
            const int64_t& num_hop,
            const int64_t* hop_ind,
            const double complex* hop_eng,
            const double* sc_lattice,
            const double* sc_orb_pos)
        uint64_t get_num_hop() const
        void to_array(int64_t* hop_ind, double complex* hop_eng, double* dr) const


cdef extern from "orbital_map.h" namespace "orbital_map":
    # Functions
    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)


cdef extern from "super.h" namespace "super":
    # Functions
    void build_orb_id_pc_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,
        cell_index_t* orb_id_pc)

    void build_orb_id_pc_vac_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,
        cell_index_t* orb_id_pc)

    void build_orb_pos_cxx[cell_index_t](
        const double* pc_lattice,
        const double* pc_orb_pos,
        const int64_t& num_orb_sc,
        const cell_index_t* orb_id_pc,
        double* sc_orb_pos)

    void build_orb_eng_cxx[cell_index_t](
        const double* pc_orb_eng,
        const int64_t& num_orb_sc,
        const cell_index_t* orb_id_pc,
        double* sc_orb_eng)

    int64_t build_hop_gen_cxx[cell_index_t](
        const int64_t& num_hop_pc,
        const cell_index_t* pc_hop_ind,
        const double complex* pc_hop_eng,
        const cell_index_t* dim,
        const cell_index_t* pbc,
        const cell_index_t& num_orb_pc,
        const int64_t& num_orb_sc,
        const cell_index_t* orb_id_pc,
        const int64_t& num_vac,
        const int64_t* vac_id_sc,
        const double* sc_lattice,
        const double* sc_orb_pos,
        int64_t* hop_ind,
        double complex* hop_eng,
        double* dr)

    void stat_hop_cxx[cell_index_t](
        const int64_t& num_hop,
        const cell_index_t* pc_hop_ind,
        const cell_index_t* pbc,
        int32_t* status)

    void split_hop_cxx[cell_index_t](
        const int64_t& num_hop,
        const cell_index_t* pc_hop_ind,
        const double complex* pc_hop_eng,
        const int32_t* status,
        cell_index_t* ind_pbc,
        cell_index_t* ind_free,
        double complex* eng_pbc,
        double complex* eng_free)

    void build_hop_pbc_cxx[cell_index_t](
        const int64_t& num_hop_pc,
        const cell_index_t* pc_hop_ind,
        const double complex* pc_hop_eng,
        const cell_index_t* dim,
        const cell_index_t* pbc,
        const cell_index_t& num_orb_pc,
        const double* sc_lattice,
        const double* sc_orb_pos,
        int64_t* hop_ind,
        double complex* hop_eng,
        double* dr)

    int64_t count_hop_cxx(
        const int64_t& num_orb_sc,
        const int64_t& num_hop_sc,
        const int64_t* hop_ind_sc,
        int64_t* hop_count)

    void get_orb_id_trim_cxx[cell_index_t](
        const int64_t& num_orb_sc,
        const cell_index_t* orb_id_pc,
        const int64_t* hop_count,
        cell_index_t* orb_id_trim)

    void build_inter_dr_cxx(
        const int64_t& num_hop,
        const int64_t* hop_ind,
        const double* pos_bra,
        const double* pos_ket,
        const double* sc_lat_ket,
        double* dr)


#------------------------------------------------------------------------------#
#                       Class for applying hop_modifier                        #
#------------------------------------------------------------------------------#
cdef class CPPHopDict:
    """Container of hopping terms in the supercell."""
    cdef HopDict backend

    def __cinit__(self):
        self.backend = HopDict()

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def add_hop_origin(self,
        int64_t [:,::1] hop_ind,
        double complex [::1] hop_eng,
        double [:,::1] dr):
        """
        Add the original hopping terms in the supercell.

        Parameters
        ----------
        hop_ind: (num_hop_sc, 2) int64 array
            row and column indices of hopping terms in sc representation
            reduced by a half using the conjugate relation
        hop_eng: (num_hop_sc,) complex128 array
            reduced hopping energy of <i|H|j> in eV
        dr: (num_hop_sc, 3) float64 array
            reduced distances corresponding to hop_ind in NM

        Returns
        -------
        None.

        NOTES
        -----
        Lower triangular terms will be transformed to upper triangular.
        """
        cdef int64_t num_hop = hop_ind.shape[0]
        self.backend.add_hop_origin(
            num_hop,
            &hop_ind[0][0],
            &hop_eng[0],
            &dr[0][0])

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def add_hop_modifier(self,
        int64_t [:,::1] hop_ind,
        double complex [::1] hop_eng,
        double [:,::1] sc_lattice,
        double [:,::1] sc_orb_pos):
        """
        Add hopping terms in hop_modifier.

        Parameters
        ----------
        hop_ind: (num_hop, 5) int64 array
            indices of hopping terms in the hop modifier
            reduced by a half using the conjugate relation
        hop_eng: (num_hop,) complex128 array
            reduced energies of hopping terms in the hop modifier
        sc_lattice: (3, 3) float64 array
            CARTESIAN lattice vectors of supercell in NM
        sc_orb_pos: (num_orb_sc, 3) float64 array
            CARTESIAN coordinates of orbitals of super cell in NM

        Returns
        -------
        None.

        NOTES
        -----
        Lower triangular terms will be transformed to upper triangular.
        """
        cdef int64_t num_hop = hop_ind.shape[0]
        self.backend.add_hop_modifier(
            num_hop,
            &hop_ind[0][0],
            &hop_eng[0],
            &sc_lattice[0][0],
            &sc_orb_pos[0][0])

    @cython.boundscheck(False)
    @cython.wraparound(False)
    def to_array(self):
        """
        Build the arrays of hopping indices, energies and distances.

        Parameters
        ----------
        None.

        Returns
        -------
        hop_ind: (num_hop_sc, 2) int64 array
            row and column indices of hopping terms in sc representation
            reduced by a half using the conjugate relation
        hop_eng: (num_hop_sc,) complex128 array
            reduced hopping energy of <i|H|j> in eV
        dr: (num_hop_sc, 3) float64 array
            reduced distances corresponding to hop_ind in NM

        NOTES
        -----
        Dimension of the super cell along each direction should be no smaller
        than a minimum value. Otherwise the result will be wrong. See the
        documentation of 'OrbitalSet' class for more details.

        The hopping terms have been reduced by the conjugate relation. So only
        half of the terms are stored in hop_* and dr.
        """
        cdef int64_t num_hop = self.backend.get_num_hop()
        cdef int64_t [:,::1] hop_ind = np.zeros((num_hop, 2), dtype=np.int64)
        cdef double complex [::1] hop_eng = np.zeros(num_hop, dtype=np.complex128)
        cdef double [:,::1] dr = np.zeros((num_hop, 3), dtype=np.float64)
        self.backend.to_array(&hop_ind[0][0], &hop_eng[0], &dr[0][0])
        return np.asarray(hop_ind), np.asarray(hop_eng), np.asarray(dr)


#-------------------------------------------------------------------------------
#       Functions for constructing attributes of OrbitalSet class
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def build_orb_id_pc(
    int32_t [::1] dim,
    int32_t num_orb_pc,
    int64_t [::1] vac_id_sc):
    """
    Build the indices of orbitals in primitive cell representation.

    Parameters
    ----------
    dim: (3,) int32 array
        dimension of the super cell
    num_orb_pc: int32
        number of orbitals in primitive cell
    vac_id_sc: (num_vac,) int64 arrray
        indices of vacancies in sc representation
    
    Returns
    -------
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals in pc representation
    """
    cdef int64_t num_vac = vac_id_sc.shape[0]
    cdef int64_t num_orb_sc = np.prod(dim) * num_orb_pc - num_vac
    cdef int32_t [:,::1] orb_id_pc = np.zeros((num_orb_sc, 4), dtype=np.int32)
    if num_vac == 0:
        build_orb_id_pc_cxx(
            &dim[0],
            num_orb_pc,
            num_vac,
            &vac_id_sc[0],
            &orb_id_pc[0][0])
    else:
        build_orb_id_pc_vac_cxx(
            &dim[0],
            num_orb_pc,
            num_vac,
            &vac_id_sc[0],
            &orb_id_pc[0][0])
    return np.asarray(orb_id_pc)


#-------------------------------------------------------------------------------
#              Functions for building arrays for SuperCell class
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def build_orb_pos(
    double [:,::1] pc_lattice,
    double [:,::1] pc_orb_pos,
    int32_t [:,::1] orb_id_pc):
    """
    Build the array of Cartesian coordinates for all the orbitals in
    super cell.

    Parameters
    ----------
    pc_lattice: (3, 3) float64 array
        lattice vectors of primitive cell in NM
    pc_orb_pos: (num_orb_pc, 3) float64 array
        FRACTIONAL coordinates of orbitals in primitive cell
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    sc_orb_pos: (num_orb_sc, 3) float64 array
        CARTESIAN coordinates of all orbitals in super cell in NM
    """
    cdef int64_t num_orb_sc = orb_id_pc.shape[0]
    cdef double [:,::1] sc_orb_pos = np.zeros((num_orb_sc, 3), dtype=np.float64)
    build_orb_pos_cxx(
        &pc_lattice[0][0],
        &pc_orb_pos[0][0],
        num_orb_sc,
        &orb_id_pc[0][0],
        &sc_orb_pos[0][0])
    return np.asarray(sc_orb_pos)


@cython.boundscheck(False)
@cython.wraparound(False)
def build_orb_eng(
    double [::1] pc_orb_eng,
    int32_t [:,::1] orb_id_pc):
    """
    Build the array of energies for all the orbitals in super cell.

    Parameters
    ----------
    pc_orb_eng: (num_orb_pc,) float64 array
        energies of orbitals in primitive cell in eV
    orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals of super cell in pc representation

    Returns
    -------
    sc_orb_eng: (num_orb_sc,) float64 array
        energies of all orbitals in super cell in eV
    """
    cdef int64_t num_orb_sc = orb_id_pc.shape[0]
    cdef double [::1] sc_orb_eng = np.zeros(num_orb_sc, dtype=np.float64)
    build_orb_eng_cxx(
        &pc_orb_eng[0],
        num_orb_sc,
        &orb_id_pc[0][0],
        &sc_orb_eng[0])
    return np.asarray(sc_orb_eng)


@cython.boundscheck(False)
@cython.wraparound(False)
def build_hop_gen(
    int32_t [:,::1] pc_hop_ind,
    double complex [::1] pc_hop_eng,
    int32_t [::1] dim,
    int32_t [::1] pbc,
    int32_t num_orb_pc,
    int32_t [:,::1] orb_id_pc,
    int64_t [::1] vac_id_sc,
    double [:,::1] sc_lattice,
    double [:,::1] sc_orb_pos):
    """
    Build the arrays of hopping terms for constructing sparse Hamiltonian
    and dr in CSR format in general cases.

    Parameters
    ----------
    pc_hop_ind: (num_hop_pc, 5) int32 array
        reduced hopping indices of primitive cell
    pc_hop_eng: (num_orb_pc,) complex128 array
        reduced hopping energies of primitive cell in eV
    dim: (3,) int32 array
        dimension of the super cell
    pbc: (3,) int32 array
        periodic conditions
    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_vac,) int64 array
        indices of vacancies in sc representation
    sc_lattice: (3, 3) float64 array
        CARTESIAN lattice vectors of supercell in NM
    sc_orb_pos: (num_orb_sc, 3) float64 array
        CARTESIAN coordinates of orbitals of super cell in NM

    Returns
    -------
    hop_ind: (num_hop_sc, 2) int64 array
        row and column indices of hopping terms in sc representation
        reduced by a half using the conjugate relation
    hop_eng: (num_hop_sc,) complex128 array
        reduced hopping energy of <i|H|j> in eV
    dr: (num_hop_sc, 3) float64 array
        reduced distances corresponding to hop_ind in NM

    NOTES
    -----
    Dimension of the super cell along each direction should be no smaller
    than a minimum value. Otherwise the result will be wrong. See the
    documentation of 'OrbitalSet' class for more details.

    The hopping terms have been reduced by the conjugate relation. So only
    half of the terms are stored in hop_* and dr.
    """
    # Get the number of hopping terms and allocate arrays
    cdef int64_t num_hop_pc = pc_hop_ind.shape[0]
    cdef int64_t num_hop_sc = num_hop_pc * np.prod(dim)
    cdef int64_t [:,::1] hop_ind = np.zeros((num_hop_sc, 2), dtype=np.int64)
    cdef double complex [::1] hop_eng = np.zeros(num_hop_sc, dtype=np.complex128)
    cdef double [:,::1] dr = np.zeros((num_hop_sc, 3), dtype=np.float64)

    # Fill arrays
    cdef int64_t num_vac = vac_id_sc.shape[0]
    cdef int64_t num_orb_sc = orb_id_pc.shape[0]
    cdef int64_t counter = build_hop_gen_cxx(
        num_hop_pc,
        &pc_hop_ind[0][0],
        &pc_hop_eng[0],
        &dim[0],
        &pbc[0],
        num_orb_pc,
        num_orb_sc,
        &orb_id_pc[0][0],
        num_vac,
        &vac_id_sc[0],
        &sc_lattice[0][0],
        &sc_orb_pos[0][0],
        &hop_ind[0][0],
        &hop_eng[0],
        &dr[0][0])

    # Truncate arrays
    hop_ind = hop_ind[:counter]
    hop_eng = hop_eng[:counter]
    dr = dr[:counter]
    return np.asarray(hop_ind), np.asarray(hop_eng), np.asarray(dr)


@cython.boundscheck(False)
@cython.wraparound(False)
def split_pc_hop(
    int32_t [:,::1] pc_hop_ind,
    double complex [::1] pc_hop_eng,
    int32_t [::1] pbc):
    """
    Split the hopping terms of primitive cell into periodic and free
    categories.

    Parameters
    ----------
    pc_hop_ind: (num_hop, 5) int32 array
        indices of hopping terms in the primitive cell
    pc_hop_eng: (num_hop,) complex128 array
        energies of hopping terms in the primitive cell
    pbc: (3,) int32 array
        periodic boundary conditions

    Returns
    -------
    ind_pbc: (num_hop_pbc, 5) int32 array
        indices of periodic hopping terms
    eng_pbc: (num_hop_pbc,) complex128 array
        energies of periodic hopping terms
    ind_free: (num_hop_free, 5) int32 array
        indices of free hopping terms
    eng_free: (num_hop_free,) complex128 array
        energies of free hopping terms
    """
    # Determine pbc and free hopping terms
    # status: 0-free, 1-pbc
    cdef int64_t num_hop = pc_hop_ind.shape[0]
    cdef int32_t [::1] status = np.zeros(num_hop, dtype=np.int32)
    stat_hop_cxx(
        num_hop,
        &pc_hop_ind[0][0],
        &pbc[0],
        &status[0])

    # Allocate results
    cdef int64_t num_hop_pbc = np.sum(status)
    cdef int64_t num_hop_free = num_hop - num_hop_pbc
    cdef int32_t [:,::1] ind_pbc = np.zeros((num_hop_pbc, 5), dtype=np.int32)
    cdef int32_t [:,::1] ind_free = np.zeros((num_hop_free, 5), dtype=np.int32)
    cdef double complex [::1] eng_pbc = np.zeros(num_hop_pbc, dtype=np.complex128)
    cdef double complex [::1] eng_free = np.zeros(num_hop_free, dtype=np.complex128)

    # Splitting hopping terms
    split_hop_cxx(
        num_hop,
        &pc_hop_ind[0][0],
        &pc_hop_eng[0],
        &status[0],
        &ind_pbc[0][0],
        &ind_free[0][0],
        &eng_pbc[0],
        &eng_free[0])
    return np.asarray(ind_pbc), np.asarray(eng_pbc), np.asarray(ind_free), np.asarray(eng_free)


@cython.boundscheck(False)
@cython.wraparound(False)
def build_hop_pbc(
    int32_t [:,::1] pc_hop_ind,
    double complex [::1] pc_hop_eng,
    int32_t [::1] dim,
    int32_t [::1] pbc,
    int32_t num_orb_pc,
    double [:,::1] sc_lattice,
    double [:,::1] sc_orb_pos):
    """
    Build the arrays of hopping terms arising from periodic hopping terms in the
    primitive cell for constructing sparse Hamiltonian and dr in CSR format.

    Parameters
    ----------
    pc_hop_ind: (num_hop_pc, 5) int32 array
        reduced hopping indices of primitive cell
    pc_hop_eng: (num_orb_pc,) complex128 array
        reduced hopping energies of primitive cell in eV
    dim: (3,) int32 array
        dimension of the super cell
    pbc: (3,) int32 array
        periodic conditions
    num_orb_pc: int32
        number of orbitals in primitive cell
    sc_lattice: (3, 3) float64 array
        CARTESIAN lattice vectors of supercell in NM
    sc_orb_pos: (num_orb_sc, 3) float64 array
        CARTESIAN coordinates of orbitals of super cell in NM

    Returns
    -------
    hop_ind: (num_hop_sc, 2) int64 array
        row and column indices of hopping terms in sc representation
        reduced by a half using the conjugate relation
    hop_eng: (num_hop_sc,) complex128 array
        reduced hopping energy of <i|H|j> in eV
    dr: (num_hop_sc, 3) float64 array
        reduced distances corresponding to hop_ind in NM

    NOTES
    -----
    See 'build_hop'.
    """
    # Get the number of hopping terms and allocate arrays
    cdef int64_t num_hop_pc = pc_hop_ind.shape[0]
    cdef int64_t num_hop_sc = num_hop_pc * np.prod(dim)
    cdef int64_t [:,::1] hop_ind = np.zeros((num_hop_sc, 2), dtype=np.int64)
    cdef double complex [::1] hop_eng = np.zeros(num_hop_sc, dtype=np.complex128)
    cdef double [:,::1] dr = np.zeros((num_hop_sc, 3), dtype=np.float64)

    # Fill arrays
    build_hop_pbc_cxx(
        num_hop_pc,
        &pc_hop_ind[0][0],
        &pc_hop_eng[0],
        &dim[0],
        &pbc[0],
        num_orb_pc,
        &sc_lattice[0][0],
        &sc_orb_pos[0][0],
        &hop_ind[0][0],
        &hop_eng[0],
        &dr[0][0])
    return np.asarray(hop_ind), np.asarray(hop_eng), np.asarray(dr)


@cython.boundscheck(False)
@cython.wraparound(False)
def get_orb_id_trim(
    int32_t [::1] dim,
    int32_t num_orb_pc,
    int32_t [:,::1] orb_id_pc,
    int64_t [::1] vac_id_sc,
    int64_t [:,::1] hop_ind):
    """
    Get the indices of orbitals to trim in super cell representation.

    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
        orbital indices in primitive cell representation
    vac_id_sc: (num_vac,) int64 array
        indices of vacancies in sc representation
    hop_ind: (num_hop_sc, 2) int64 array
        row and column indices of hopping terms in sc representation
        reduced by a half using the conjugate relation

    Returns
    -------
    id_sc_trim: (num_trim,) int64 array
        indices of orbitals to trim in super cell representation
    """
    # Count the number or hopping terms for each orbital
    cdef int64_t num_orb_sc = orb_id_pc.shape[0]
    cdef int64_t num_hop_sc = hop_ind.shape[0]
    cdef int64_t [::1] hop_count = np.zeros(num_orb_sc, dtype=np.int64)
    cdef int64_t num_trim = count_hop_cxx(
        num_orb_sc,
        num_hop_sc,
        &hop_ind[0][0],
        &hop_count[0])

    # Collect orbitals to trim
    cdef int32_t [:,::1] id_pc_trim = np.zeros((num_trim, 4), dtype=np.int32)
    get_orb_id_trim_cxx(
        num_orb_sc,
        &orb_id_pc[0][0],
        &hop_count[0],
        &id_pc_trim[0][0])

    # Convert to sc representation
    # Since indices of orbitals to trim should be define w.r.t pristine super
    # cell, we should set num_vac to 0 to mimic it.
    cdef int64_t [::1] id_sc_trim = np.zeros(num_trim, dtype=np.int64)
    cdef int64_t num_vac = 0
    id_pc2sc_array_cxx(
        &dim[0],
        num_orb_pc,
        num_vac,
        &vac_id_sc[0],
        num_trim,
        &id_pc_trim[0][0],
        &id_sc_trim[0])
    return np.asarray(id_sc_trim)


#-------------------------------------------------------------------------------
#             Functions for building arrays for SCInterHopping class
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def build_inter_dr(
    int64_t [:,::1] hop_ind,
    double [:,::1] pos_bra,
    double [:,::1] pos_ket,
    double [:, ::1] sc_lat_ket):
    """
    Build the array of hopping distances for an 'SCInterHopping' instance.

    Parameters
    ----------
    hop_ind: (num_hop, 5) int64 array
        hopping indices
    pos_bra: (num_hop, 3) float64 array
        Cartesian coordinates of orbitals of the 'bra' super cell in nm
    pos_ket: (num_hop, 3) float64 array
        Cartesian coordinates of orbitals of the 'ket' super cell in nm
    sc_lat_ket: (3, 3) float64 array
        Cartesian coordinates of 'ket' super cell lattice vectors in nm

    Returns
    -------
    dr: (num_hop, 3) float64 array
        hopping distances in nm
    """
    cdef int64_t num_hop = hop_ind.shape[0]
    cdef double [:,::1] dr = np.zeros((num_hop, 3), dtype=np.float64)
    build_inter_dr_cxx(
        num_hop,
        &hop_ind[0][0],
        &pos_bra[0][0],
        &pos_ket[0][0],
        &sc_lat_ket[0][0],
        &dr[0][0])
    return np.asarray(dr)
