# 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) 

    cdef cppclass PySample[hop_t]:
        PySample(
            uint64_t num_orb,
            uint64_t num_hop, 
            const double* orb_eng,
            const double* orb_pos,
            const int64_t* hop_ind,
            const hop_t* hop_eng,
            const double* dr,
            const double* lattice) except +


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

    cdef void set_wf0[solver_t](
        solver_t* solver,
        const double complex* wft_wf0,
        int64_t num_orb)

    cdef void unset_wf0[solver_t](solver_t* solver)


# C++ solvers
cdef extern from "tbpm/solver.h" namespace "tbplas::tbpm":
    cdef cppclass TBPMSolver[model_t]:
        TBPMSolver(const model_t*) except +
        void calc_ldos_haydock()
        void calc_corr_dos()
        void calc_corr_ac_cond()
        void calc_corr_dyn_pol()
        void calc_quasi_eigenstates()
        void calc_corr_dc_cond()
        void calc_hall_cond()
        void calc_wft()
        void calc_corr_bands()


# Python class for importing
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


cdef class CPPSample:
    cdef PySample[double complex]* c_sample

    def __cinit__(self,
        double [::1] orb_eng,
        double [:, ::1] orb_pos,
        int64_t [:, ::1] hop_ind,
        double complex [::1] hop_eng,
        double [:, ::1] dr,
        double [:, ::1] lattice):
        cdef uint64_t num_orb = orb_eng.shape[0]
        cdef uint64_t num_hop = hop_ind.shape[0]
        self.c_sample = new PySample[double complex](
            num_orb,
            num_hop,
            &orb_eng[0],
            &orb_pos[0][0],
            &hop_ind[0][0],
            &hop_eng[0],
            &dr[0][0],
            &lattice[0][0])

    def __dealloc__(self):
        del self.c_sample


cdef class CPPTBPMSolverForPrimitiveCell:
    cdef TBPMSolver[PyPrimitiveCell[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPPrimitiveCell model,
        const string& config_name):
        self.c_solver = new TBPMSolver[PyPrimitiveCell[double complex]](model.c_prim_cell)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

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

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

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

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

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

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

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

    def calc_wft(self, double complex [::1] wft_wf0):
        set_config(self.c_solver, self.config_name)
        set_wf0(self.c_solver, &wft_wf0[0], wft_wf0.shape[0])
        self.c_solver.calc_wft()
        unset_wf0(self.c_solver)

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


cdef class CPPTBPMSolverForSample:
    cdef TBPMSolver[PySample[double complex]]* c_solver
    cdef string config_name

    def __cinit__(self,
        CPPSample model,
        const string& config_name):
        self.c_solver = new TBPMSolver[PySample[double complex]](model.c_sample)
        self.config_name = config_name

    def __dealloc__(self):
        del self.c_solver

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

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

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

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

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

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

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

    def calc_wft(self, double complex [::1] wft_wf0):
        set_config(self.c_solver, self.config_name)
        set_wf0(self.c_solver, &wft_wf0[0], wft_wf0.shape[0])
        self.c_solver.calc_wft()
        unset_wf0(self.c_solver)

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