#ifndef TBPLAS_SUPER_H
#define TBPLAS_SUPER_H

#include <complex>
#include <cstdint>
#include <cstring>

#include "common.h"
#include "memory_view.h"
#include "orbital_map.h"

namespace super {
using complex_t = std::complex<double>;
using common::calc_dr;
using memory_view::ConstMemoryView;
using memory_view::MemoryView;
using orbital_map::OrbitalMap;

/** @note
 * 1. cell_index_t stands for the datatype of id_pc, namely cell indices and
 *    orbital indices of primitive cell. In most cases, int32_t is enough.
 *    However, it may become int64_t in the future. So we make it template
 *    parameters. This convention applies for all functions in this file.
 *
 * 2. Number of hopping terms in the primitive cell is always int64_t, since
 *    an orbital can have multiple hopping terms and they can easily exceed
 *    2^31-1.
 *
 * 3. 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.
 *
 * 4. 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.
 */

// -----------------------------------------------------------------------------
//      Functions for constructing attributes of OrbitalSet class
// -----------------------------------------------------------------------------
template <typename cell_index_t>
void build_orb_id_pc_cxx(
    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)
{
    // Memory views
    MemoryView<cell_index_t> id_pc_map(orb_id_pc, 4);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc, vac_id_sc, num_vac);

    // OMP private variables
    cell_index_t ia = 0;
    cell_index_t ib = 0;
    cell_index_t ic = 0;
    cell_index_t i_orb = 0;
    int64_t id_sc = 0;
    cell_index_t* id_pc = nullptr;

#pragma omp parallel for private(ib, ic, i_orb, id_sc, id_pc)
    for (ia = 0; ia < dim[0]; ++ia) {
        for (ib = 0; ib < dim[1]; ++ib) {
            for (ic = 0; ic < dim[2]; ++ic) {
                for (i_orb = 0; i_orb < num_orb_pc; ++i_orb) {
                    id_sc = orb_map.id_pc2sc(ia, ib, ic, i_orb);
                    id_pc = id_pc_map.addr(id_sc, 0);
                    id_pc[0] = ia;
                    id_pc[1] = ib;
                    id_pc[2] = ic;
                    id_pc[3] = i_orb;
                }
            }
        }
    }
}

template <typename cell_index_t>
void build_orb_id_pc_vac_cxx(
    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)
{
    // Memory views
    MemoryView<cell_index_t> id_pc_map(orb_id_pc, 4);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc, vac_id_sc, num_vac);

    // OMP private variables
    cell_index_t ia = 0;
    cell_index_t ib = 0;
    cell_index_t ic = 0;
    cell_index_t i_orb = 0;
    int64_t id_sc = 0;
    cell_index_t* id_pc = nullptr;

#pragma omp parallel for private(ib, ic, i_orb, id_sc, id_pc)
    for (ia = 0; ia < dim[0]; ++ia) {
        for (ib = 0; ib < dim[1]; ++ib) {
            for (ic = 0; ic < dim[2]; ++ic) {
                for (i_orb = 0; i_orb < num_orb_pc; ++i_orb) {
                    id_sc = orb_map.id_pc2sc_vac2(ia, ib, ic, i_orb);
                    if (id_sc != -1) {
                        id_pc = id_pc_map.addr(id_sc, 0);
                        id_pc[0] = ia;
                        id_pc[1] = ib;
                        id_pc[2] = ic;
                        id_pc[3] = i_orb;
                    }
                }
            }
        }
    }
}

// -----------------------------------------------------------------------------
//             Functions for building arrays for SuperCell class
// -----------------------------------------------------------------------------
/**
 * @brief Wrap primitive cell index back into the 0th period of super cell
 *   based on super cell dimension and boundary condition.
 *
 * @tparam cell_index_t datatype of cell index
 * @param[in] j0 index of primitive cell before wrapping
 * @param[in] n0 dimension of super cell along given direction
 * @param[in] pbc whether to enforce periodic condition along given direction
 * @param[out] j1 index of primitive cell after wrapping, j0 % n0
 * @param[out] n1 period index after wrapping, j0 // n0
 *
 * @note
 *  Mathematically, this is done by j1 = j0 % n0 and n1 = j0 // n0, following
 *  Python flavor of division (not C).
 *
 *  For open boundary condition (pbc == 0), j1 will be dropped by setting it to
 *  -1 if n1 != 0.
 */
template <typename cell_index_t>
void wrap_pbc(
    const cell_index_t& j0,
    const cell_index_t& n0,
    const cell_index_t& pbc,
    //
    cell_index_t& j1,
    cell_index_t& n1)
{
    // Branches are sorted according to their matching probabilities.
    if (0 <= j0 && j0 < n0) {
        j1 = j0;
        n1 = 0;
    } else if (n0 <= j0 && j0 < 2 * n0) {
        j1 = j0 - n0;
        n1 = 1;
    } else if (-n0 <= j0 && j0 < 0) {
        j1 = j0 + n0;
        n1 = -1;
    } else {
        cell_index_t j1_new = j0;
        while (j1_new < 0) {
            j1_new += n0;
        }
        while (j1_new >= n0) {
            j1_new -= n0;
        }
        j1 = j1_new;
        n1 = (j0 - j1_new) / n0;
    }

    // Apply pbc
    if (pbc == 0 && n1 != 0) {
        j1 = -1;
    }
}

template <typename cell_index_t>
void build_orb_pos_cxx(
    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)
{
    // Memory views
    ConstMemoryView<cell_index_t> orb_id_pc_map(orb_id_pc, 4);
    ConstMemoryView<double> pc_orb_pos_map(pc_orb_pos, 3);
    MemoryView<double> sc_orb_pos_map(sc_orb_pos, 3);

    // OMP private variables
    const cell_index_t* id_pc_i = nullptr;
    const double* pc_pos_i = nullptr;
    double* sc_pos_i = nullptr;
    double ra = 0.0;
    double rb = 0.0;
    double rc = 0.0;

#pragma omp parallel for private(id_pc_i, pc_pos_i, sc_pos_i, ra, rb, rc)
    for (int64_t io_sc = 0; io_sc < num_orb_sc; ++io_sc) {
        id_pc_i = orb_id_pc_map.addr(io_sc, 0);
        pc_pos_i = pc_orb_pos_map.addr(orb_id_pc_map(io_sc, 3), 0);
        sc_pos_i = sc_orb_pos_map.addr(io_sc, 0);
        // pos_frac + rn
        ra = pc_pos_i[0] + id_pc_i[0];
        rb = pc_pos_i[1] + id_pc_i[1];
        rc = pc_pos_i[2] + id_pc_i[2];
        sc_pos_i[0] = ra * pc_lattice[0] + rb * pc_lattice[3] + rc * pc_lattice[6];
        sc_pos_i[1] = ra * pc_lattice[1] + rb * pc_lattice[4] + rc * pc_lattice[7];
        sc_pos_i[2] = ra * pc_lattice[2] + rb * pc_lattice[5] + rc * pc_lattice[8];
    }
}

template <typename cell_index_t>
void build_orb_eng_cxx(
    const double* pc_orb_eng,
    //
    const int64_t& num_orb_sc,
    const cell_index_t* orb_id_pc,
    //
    double* sc_orb_eng)
{
    ConstMemoryView<cell_index_t> orb_id_pc_map(orb_id_pc, 4);
#pragma omp parallel for
    for (int64_t io_sc = 0; io_sc < num_orb_sc; ++io_sc) {
        sc_orb_eng[io_sc] = pc_orb_eng[orb_id_pc_map(io_sc, 3)];
    }
}

template <typename cell_index_t>
int64_t build_hop_gen_cxx(
    const int64_t& num_hop_pc, // use 64 bit to avoid overflow
    const cell_index_t* pc_hop_ind,
    const complex_t* 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,
    complex_t* hop_eng,
    double* dr)
{
    // MemoryViews
    ConstMemoryView<cell_index_t> orb_id_pc_map(orb_id_pc, 4);
    ConstMemoryView<cell_index_t> pc_hop_ind_map(pc_hop_ind, 5);
    MemoryView<int64_t> hop_ind_map(hop_ind, 2);
    MemoryView<double> dr_map(dr, 3);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc, vac_id_sc, num_vac);

    // Buffers and intermediate results
    int64_t counter = 0;

    // Get id_sc for bra
    for (int64_t io_sc = 0; io_sc < num_orb_sc; ++io_sc) {
        const cell_index_t* id_pc_i = orb_id_pc_map.addr(io_sc, 0);
        int64_t id_sc_i = io_sc;

        // Loop over hopping terms in primitive cell
        for (int64_t ih = 0; ih < num_hop_pc; ++ih) {
            const cell_index_t* hop_term = pc_hop_ind_map.addr(ih, 0);
            if (id_pc_i[3] == hop_term[3]) {

                // Apply boundary condition
                cell_index_t ja0 = id_pc_i[0] + hop_term[0];
                cell_index_t jb0 = id_pc_i[1] + hop_term[1];
                cell_index_t jc0 = id_pc_i[2] + hop_term[2];
                cell_index_t ja = 0;
                cell_index_t jb = 0;
                cell_index_t jc = 0;
                cell_index_t na = 0;
                cell_index_t nb = 0;
                cell_index_t nc = 0;
                wrap_pbc(ja0, dim[0], pbc[0], ja, na);
                wrap_pbc(jb0, dim[1], pbc[1], jb, nb);
                wrap_pbc(jc0, dim[2], pbc[2], jc, nc);

                // Drop the hopping term if it is out of boundary
                bool in_bound = (ja != -1) && (jb != -1) && (jc != -1);
                if (in_bound) {
                    // Evaluate id_sc_j
                    int64_t id_sc_j = 0;
                    if (num_vac == 0) {
                        id_sc_j = orb_map.id_pc2sc(ja, jb, jc, hop_term[4]);
                    } else {
                        id_sc_j = orb_map.id_pc2sc_vac2(ja, jb, jc, hop_term[4]);
                    }

                    // Check if id_sc_j corresponds to a vacancy
                    if (num_vac == 0 || id_sc_j != -1) {
                        hop_ind_map(counter, 0) = id_sc_i;
                        hop_ind_map(counter, 1) = id_sc_j;
                        hop_eng[counter] = pc_hop_eng[ih];
                        calc_dr(
                            sc_orb_pos, sc_orb_pos, id_sc_i, id_sc_j,
                            sc_lattice, na, nb, nc,
                            dr_map.addr(counter, 0));
                        ++counter;
                    }
                }
            }
        }
    }
    return counter;
}

template <typename cell_index_t>
void stat_hop_cxx(
    const int64_t& num_hop, // use 64 bit to avoid overflow
    const cell_index_t* pc_hop_ind,
    const cell_index_t* pbc,
    //
    int32_t* status)
{
    // Memory views
    ConstMemoryView<cell_index_t> hop_ind_map(pc_hop_ind, 5);

    // OMP private variables
    const cell_index_t* hop_ind = nullptr;

    // Determine pbc and free hopping terms
    // status: 0-free, 1-pbc
#pragma omp parallel for private(hop_ind)
    for (int64_t ih = 0; ih < num_hop; ++ih) {
        hop_ind = hop_ind_map.addr(ih, 0);
        if ((pbc[0] == 0 && hop_ind[0] != 0) || (pbc[1] == 0 && hop_ind[1] != 0) || (pbc[2] == 0 && hop_ind[2] != 0)) {
            status[ih] = 0;
        } else {
            status[ih] = 1;
        }
    }
}

template <typename cell_index_t>
void split_hop_cxx(
    const int64_t& num_hop, // use 64 bit to avoid overflow
    const cell_index_t* pc_hop_ind,
    const complex_t* pc_hop_eng,
    //
    const int32_t* status,
    //
    cell_index_t* ind_pbc,
    cell_index_t* ind_free,
    complex_t* eng_pbc,
    complex_t* eng_free)
{
    // Memory views
    ConstMemoryView<cell_index_t> ind_all_map(pc_hop_ind, 5);
    MemoryView<cell_index_t> ind_pbc_map(ind_pbc, 5);
    MemoryView<cell_index_t> ind_free_map(ind_free, 5);

    // Intermediate results
    int64_t counter_pbc = 0;
    int64_t counter_free = 0;
    constexpr int64_t buffer_size = 5 * sizeof(cell_index_t);

    for (int64_t ih = 0; ih < num_hop; ++ih) {
        const cell_index_t* p0 = ind_all_map.addr(ih, 0);
        if (status[ih] == 1) {
            cell_index_t* p1 = ind_pbc_map.addr(counter_pbc, 0);
            memcpy(p1, p0, buffer_size);
            eng_pbc[counter_pbc] = pc_hop_eng[ih];
            ++counter_pbc;
        } else {
            cell_index_t* p1 = ind_free_map.addr(counter_free, 0);
            memcpy(p1, p0, buffer_size);
            eng_free[counter_free] = pc_hop_eng[ih];
            ++counter_free;
        }
    }
}

template <typename cell_index_t>
void build_hop_pbc_cxx(
    const int64_t& num_hop_pc, // use 64 bit to avoid overflow
    const cell_index_t* pc_hop_ind,
    const complex_t* 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,
    complex_t* hop_eng,
    double* dr)
{
    // MemoryViews
    ConstMemoryView<cell_index_t> pc_hop_ind_map(pc_hop_ind, 5);
    MemoryView<int64_t> hop_ind_map(hop_ind, 2);
    MemoryView<double> dr_map(dr, 3);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc);

    // Buffers and intermediate results
    int64_t counter = 0;

    for (cell_index_t ia = 0; ia < dim[0]; ++ia) {
        for (cell_index_t ib = 0; ib < dim[1]; ++ib) {
            for (cell_index_t ic = 0; ic < dim[2]; ++ic) {

                // Evaluate the offset for orbital i
                int64_t offset_sc_i = orb_map.id_pc2sc(ia, ib, ic, 0);

                for (int64_t ih = 0; ih < num_hop_pc; ++ih) {
                    // Default offset_j and new hop_rn assuming rn == (0, 0, 0)
                    int64_t offset_sc_j = offset_sc_i;
                    cell_index_t na = 0;
                    cell_index_t nb = 0;
                    cell_index_t nc = 0;

                    // Update offset_j and new hop_rn if rn != (0, 0, 0)
                    const cell_index_t* hop_term = pc_hop_ind_map.addr(ih, 0);
                    bool rn_non_zero = (hop_term[0] != 0) || (hop_term[1] != 0) || (hop_term[2] != 0);
                    if (rn_non_zero) {
                        cell_index_t ja0 = ia + hop_term[0];
                        cell_index_t jb0 = ib + hop_term[1];
                        cell_index_t jc0 = ic + hop_term[2];
                        cell_index_t ja = 0;
                        cell_index_t jb = 0;
                        cell_index_t jc = 0;
                        wrap_pbc(ja0, dim[0], pbc[0], ja, na);
                        wrap_pbc(jb0, dim[1], pbc[1], jb, nb);
                        wrap_pbc(jc0, dim[2], pbc[2], jc, nc);
                        offset_sc_j = orb_map.id_pc2sc(ja, jb, jc, 0);
                    }

                    // Fill hopping terms
                    int64_t id_sc_i = offset_sc_i + hop_term[3];
                    int64_t id_sc_j = offset_sc_j + hop_term[4];
                    hop_ind_map(counter, 0) = id_sc_i;
                    hop_ind_map(counter, 1) = id_sc_j;
                    hop_eng[counter] = pc_hop_eng[ih];
                    calc_dr(
                        sc_orb_pos, sc_orb_pos, id_sc_i, id_sc_j,
                        sc_lattice, na, nb, nc,
                        dr_map.addr(counter, 0));
                    ++counter;
                }
            }
        }
    }
}

inline 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)
{
    // Count the number or hopping terms for each orbital
    ConstMemoryView<int64_t> hop_ind_map(hop_ind_sc, 2);
    for (int64_t ih = 0; ih < num_hop_sc; ++ih) {
        const int64_t* ind = hop_ind_map.addr(ih, 0);
        ++hop_count[ind[0]];
        ++hop_count[ind[1]];
    }

    // Determine the number of orbitals to trim
    // Since we have reduced hop_ind, so dangling orbitals have
    // hopping terms <= 1.
    int64_t num_orb_trim = 0;
    for (int64_t io_sc = 0; io_sc < num_orb_sc; ++io_sc) {
        if (hop_count[io_sc] <= 1) {
            ++num_orb_trim;
        }
    }
    return num_orb_trim;
}

template <typename cell_index_t>
void get_orb_id_trim_cxx(
    const int64_t& num_orb_sc,
    const cell_index_t* orb_id_pc,
    //
    const int64_t* hop_count,
    //
    cell_index_t* orb_id_trim)
{
    // Memory views
    ConstMemoryView<cell_index_t> full_map(orb_id_pc, 4);
    MemoryView<cell_index_t> trim_map(orb_id_trim, 4);

    // Buffers and intermediate results
    constexpr int64_t buffer_size = 4 * sizeof(cell_index_t);
    int64_t counter = 0;

    for (int64_t io_sc = 0; io_sc < num_orb_sc; ++io_sc) {
        if (hop_count[io_sc] <= 1) {
            const cell_index_t* p0 = full_map.addr(io_sc, 0);
            cell_index_t* p1 = trim_map.addr(counter, 0);
            memcpy(p1, p0, buffer_size);
            ++counter;
        }
    }
}

// -----------------------------------------------------------------------------
//            Functions for building arrays for SCInterHopping class
// -----------------------------------------------------------------------------
inline 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)
{
    // Memory views
    ConstMemoryView<int64_t> hop_ind_map(hop_ind, 5);
    MemoryView<double> dr_map(dr, 3);

    // OMP private variables
    const int64_t* ind = nullptr;
    int64_t na = 0;
    int64_t nb = 0;
    int64_t nc = 0;
    int64_t id_bra = 0;
    int64_t id_ket = 0;

#pragma omp parallel for private(ind, na, nb, nc, id_bra, id_ket)
    for (int64_t ih = 0; ih < num_hop; ++ih) {
        ind = hop_ind_map.addr(ih, 0);
        na = ind[0];
        nb = ind[1];
        nc = ind[2];
        id_bra = ind[3];
        id_ket = ind[4];
        calc_dr(
            pos_bra, pos_ket, id_bra, id_ket,
            sc_lat_ket, na, nb, nc,
            dr_map.addr(ih, 0));
    }
}
} // namespace super
#endif