# 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 int64_t
import numpy as np


cdef extern from "sample.h" namespace "sample":
    void apply_mag_field_cxx(
        int64_t num_hop_sc,
        const int64_t* hop_ind,
        double complex* hop_eng,
        const double* dr,
        const double* orb_pos,
        double intensity,
        int64_t gauge)

#-------------------------------------------------------------------------------
#                   Functions for advanced Sample utilities
#-------------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
def apply_mag_field(
    int64_t [:,::1] hop_ind,
    double complex [::1] hop_eng,
    double [:,::1] dr,
    double [:,::1] orb_pos,
    double intensity,
    int64_t gauge):
    """
    Add magnetic field perpendicular to xOy-plane via Peierls substitution.

    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
        hopping energies in accordance with hop_end in eV
    dr: (num_hop_sc, 3) float64 array
        distances corresponding to hop_ind in NM
    orb_pos: (num_orb_sc, 3) float64 array
        CARTESIAN coordinates of all orbitals in super cell in NM
    intensity: float
        intensity of magnetic field in Tesla
    gauge: int64
        gauge of vector field

    Returns
    -------
    None. Results are saved in hop_eng.
    """
    cdef int64_t num_hop_sc = hop_ind.shape[0]
    apply_mag_field_cxx(
        num_hop_sc,
        &hop_ind[0][0],
        &hop_eng[0],
        &dr[0][0],
        &orb_pos[0][0],
        intensity,
        gauge)
