# 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
# distutils: language = c++

import cython
from libc.stdint cimport int32_t, int64_t, uint64_t
from libcpp.string cimport string


# C++ model class
cdef extern from "pymodel.h" namespace "tbplas::python":
    cdef cppclass PyPrimitiveCell[hop_t]:
        PyPrimitiveCell(
            uint64_t num_orb,
            uint64_t num_hop, 
            const double* orb_eng,
            const double* orb_pos,
            const int32_t* hop_ind,
            const hop_t* hop_eng,
            const double* lattice,
            const double* origin) except +
        void build_ham_dense(double* k_point, double complex* ham, int convention) 


# C++ auxiliary functions for settings calculation parameters
cdef extern from "interface_diag.h" namespace "tbplas::python":
    void set_config[solver_t](
        solver_t* solver,
        const string& config_name)

    void set_k_points[solver_t](
        solver_t* solver,
        const double* k_points,
        uint64_t num_kpt)

    void unset_k_points[solver_t](solver_t* solver)

    void set_k_array[solver_t](
        solver_t* solver,
        const double* ka_array,
        uint64_t num_ka,
        const double* kb_array,
        uint64_t num_kb,
        double kc)


# C++ data types and solvers
cdef extern from "diag/base.h" namespace "tbplas::diag":
    cdef cppclass BandData:
        pass

    cdef cppclass DosData:
        pass

    cdef cppclass DiagSolver[model_t]:
        DiagSolver(const model_t* model) except +
        DiagSolver(const model_t* model, const model_t* overlap) except +
        BandData calc_bands()
        DosData calc_dos()
        void calc_states()


cdef extern from "diag/lindhard.h" namespace "tbplas::diag":
    cdef cppclass DynPolData:
        pass

    cdef cppclass EpsilonData:
        pass

    cdef cppclass ACCondData:
        pass

    cdef cppclass Lindhard[model_t]:
        Lindhard(const model_t* model) except +
        Lindhard(const model_t* model, const model_t* overlap) except +
        DynPolData calc_dyn_pol()
        ACCondData calc_ac_cond()


cdef extern from "diag/magnetic.h" namespace "tbplas::diag":
    cdef cppclass SpinData:
        pass

    cdef cppclass SpinTexture[model_t]:
        SpinTexture(const model_t* model) except +
        SpinTexture(const model_t* model, const model_t* overlap) except +
        SpinData calc_spin_texture()


cdef extern from "diag/z2.h" namespace "tbplas::diag":
    cdef cppclass Z2Data:
        pass

    cdef cppclass Z2[model_t]:
        Z2(const model_t* model) except +
        Z2(const model_t* model, const model_t* overlap) except +
        Z2Data calc_phases()


cdef extern from "diag/berry.h" namespace "tbplas::diag":
    cdef cppclass BerryData:
        pass

    cdef cppclass Berry[model_t]:
        Berry(const model_t* model) except +
        Berry(const model_t* model, const model_t* overlap) except +
        BerryData calc_berry_curvature_kubo()
        BerryData calc_berry_curvature_wilson()


# Python model class
cdef class CPPPrimitiveCell:
    cdef PyPrimitiveCell[double complex]* c_prim_cell

    def __cinit__(self,
        double [::1] orb_eng,
        double [:, ::1] orb_pos,
        int32_t [:, ::1] hop_ind,
        double complex [::1] hop_eng,
        double [:, ::1] lattice,
        double [:] origin):
        cdef uint64_t num_orb = orb_eng.shape[0]
        cdef uint64_t num_hop = hop_ind.shape[0]
        self.c_prim_cell = new PyPrimitiveCell[double complex](
            num_orb,
            num_hop,
            &orb_eng[0],
            &orb_pos[0][0],
            &hop_ind[0][0],
            &hop_eng[0],
            &lattice[0][0],
            &origin[0])

    def __dealloc__(self):
        del self.c_prim_cell

    def build_ham_dense(self,
        double [::1] k_point,
        double complex [::1, :] ham,
        int convention):
        """
        Set Hamiltonian for given k-point.

        NOTE: This function assumes ham to be pre-allocated in column-major
        order. So the 'order' attribute of numpy array must be 'F' (Fortran).
        """
        self.c_prim_cell.build_ham_dense(&k_point[0], &ham[0][0], convention)


# Python solvers
cdef class CPPDiagSolver:
    cdef DiagSolver[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        CPPPrimitiveCell overlap,
        const string& config_name):
        if overlap is not None:
            self.c_solver = new DiagSolver[PyPrimitiveCell[double complex]](model.c_prim_cell, overlap.c_prim_cell)
        else:
            self.c_solver = new DiagSolver[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

    def calc_bands(self, double [:, ::1] k_points):
        set_config(self.c_solver, self.config_name)
        set_k_points(self.c_solver, &k_points[0][0], k_points.shape[0])
        self.c_solver.calc_bands()
        unset_k_points(self.c_solver)

    def calc_dos(self, double [:, ::1] k_points):
        set_config(self.c_solver, self.config_name)
        set_k_points(self.c_solver, &k_points[0][0], k_points.shape[0])
        self.c_solver.calc_dos()
        unset_k_points(self.c_solver)

    def calc_states(self, double [:, ::1] k_points):
        set_config(self.c_solver, self.config_name)
        set_k_points(self.c_solver, &k_points[0][0], k_points.shape[0])
        self.c_solver.calc_states()
        unset_k_points(self.c_solver)


cdef class CPPLindhard:
    cdef Lindhard[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        CPPPrimitiveCell overlap,
        const string& config_name):
        if overlap is not None:
            self.c_solver = new Lindhard[PyPrimitiveCell[double complex]](model.c_prim_cell, overlap.c_prim_cell)
        else:
            self.c_solver = new Lindhard[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

    def calc_dyn_pol(self):
        set_config(self.c_solver, self.config_name)
        self.c_solver.calc_dyn_pol()

    def calc_ac_cond(self):
        set_config(self.c_solver, self.config_name)
        self.c_solver.calc_ac_cond()


cdef class CPPSpinTexture:
    cdef SpinTexture[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        CPPPrimitiveCell overlap,
        const string& config_name):
        if overlap is not None:
            self.c_solver = new SpinTexture[PyPrimitiveCell[double complex]](model.c_prim_cell, overlap.c_prim_cell)
        else:
            self.c_solver = new SpinTexture[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

    def calc_spin_texture(self, double [:, ::1] k_points):
        set_config(self.c_solver, self.config_name)
        set_k_points(self.c_solver, &k_points[0][0], k_points.shape[0])
        self.c_solver.calc_spin_texture()
        unset_k_points(self.c_solver)


cdef class CPPZ2:
    cdef Z2[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        CPPPrimitiveCell overlap,
        const string& config_name):
        if overlap is not None:
            self.c_solver = new Z2[PyPrimitiveCell[double complex]](model.c_prim_cell, overlap.c_prim_cell)
        else:
            self.c_solver = new Z2[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

    def calc_phases(self,
        double [::1] ka_array,
        double [::1] kb_array,
        double kc):
        set_config(self.c_solver, self.config_name)
        set_k_array(
            self.c_solver,
            &ka_array[0],
            ka_array.shape[0],
            &kb_array[0],
            kb_array.shape[0],
            kc)
        self.c_solver.calc_phases()


cdef class CPPBerry:
    cdef Berry[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        CPPPrimitiveCell overlap,
        const string& config_name):
        if overlap is not None:
            self.c_solver = new Berry[PyPrimitiveCell[double complex]](model.c_prim_cell, overlap.c_prim_cell)
        else:
            self.c_solver = new Berry[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

    def calc_berry_curvature_kubo(self):
        set_config(self.c_solver, self.config_name)
        self.c_solver.calc_berry_curvature_kubo()

    def calc_berry_curvature_wilson(self):
        set_config(self.c_solver, self.config_name)
        self.c_solver.calc_berry_curvature_wilson()
