#ifndef TBPLAS_ORBITAL_MAP_H
#define TBPLAS_ORBITAL_MAP_H

#include <cstdint>

#include "memory_view.h"

namespace orbital_map {
using memory_view::ConstMemoryView;
using memory_view::MemoryView;

/**
 * @brief Class for converting indices between primitive cell (pc) and
 * supercell (sc) representations
 *
 * @tparam cell_index_t datatype of index in pc representation
 *
 * @note
 *   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.
 */
template <typename cell_index_t>
class OrbitalMap {
private:
    /// @brief (3,) array, dimension of the super cell
    const cell_index_t* dim_ = nullptr;

    /// @brief number of orbitals in primitive cell
    cell_index_t num_orb_pc_ = 0;

    /// @brief vac_id_sc (num_vac,) array, indices of vacancies in sc representation
    const int64_t* vac_id_sc_ = nullptr;

    /// @brief number of vacancies
    int64_t num_vac_ = 0;

    /// @brief shortcut for dimensions
    cell_index_t dim_c_ = 0;
    int64_t dim_bc_ = 0;

    // If id_sc corresponds to a vacancy, return -1.
    [[nodiscard]] int64_t get_offset(const int64_t& id_sc) const
    {
        int64_t offset = 0;
        if (id_sc < vac_id_sc_[0]) {
            offset = 0;
        } else if (id_sc > vac_id_sc_[num_vac_ - 1]) {
            offset = num_vac_;
        } else {
            offset = 0;
            for (int64_t iv = 0; iv < num_vac_; ++iv) {
                int64_t vac = vac_id_sc_[iv];
                if (vac < id_sc) {
                    ++offset;
                } else if (vac == id_sc) {
                    return -1;
                } else {
                    break;
                }
            }
        }
        return offset;
    }

    // If id_sc corresponds to a vacancy, return -1.
    [[nodiscard]] int64_t get_offset2(const int64_t& id_sc) const
    {
        // Initialize variables
        int64_t i0 = 0;
        int64_t i1 = num_vac_ - 1;
        int64_t x0 = vac_id_sc_[i0];
        int64_t x1 = vac_id_sc_[i1];

        // Check for end points
        int64_t offset = 0;
        if (id_sc < x0) {
            offset = 0;
        } else if (id_sc == x0 || id_sc == x1) {
            offset = -1;
        } else if (id_sc > x1) {
            offset = num_vac_;
        } else {
            // Bi-section algorithm for remaining parts
            offset = i0 + 1;
            while (i1 - i0 > 1) {
                int64_t im = i0 + (i1 - i0) / 2;
                int64_t xm = vac_id_sc_[im];
                if (xm < id_sc) {
                    i0 = im;
                } else if (xm > id_sc) {
                    i1 = im;
                } else {
                    offset = -1;
                    break;
                }
            }
            if (offset != -1) {
                offset = i0 + 1;
            }
        }
        return offset;
    }

public:
    // Default constructor
    OrbitalMap() = default;

    // Constructor with no vacancies
    explicit OrbitalMap(
        const cell_index_t* dim,
        const cell_index_t& num_orb_pc)
        : dim_(dim)
        , num_orb_pc_(num_orb_pc)
        , dim_c_(dim[2])
        , dim_bc_(dim[1] * dim[2])
    {
    }

    // Constructor with vacancies
    explicit OrbitalMap(
        const cell_index_t* dim,
        const cell_index_t& num_orb_pc,
        const int64_t* vac_id_sc,
        const int64_t& num_vac)
        : dim_(dim)
        , num_orb_pc_(num_orb_pc)
        , vac_id_sc_(vac_id_sc)
        , num_vac_(num_vac)
        , dim_c_(dim[2])
        , dim_bc_(dim[1] * dim[2])
    {
        if (num_vac == 0) {
            vac_id_sc_ = nullptr;
        }
    }

    /**
     * @brief Convert orbital or vacancy index from primitive cell (pc) representation
     *   to super cell (sc) representation.
     *
     * @param[in] id_pc (4,) array, index in pc representation
     * @return index in sc representation
     */
    int64_t id_pc2sc(const cell_index_t* id_pc) const
    {
        return (id_pc[0] * dim_bc_ + id_pc[1] * dim_c_ + id_pc[2]) * num_orb_pc_ + id_pc[3];
    }

    /**
     * @brief Convert orbital or vacancy index from pc to sc representation
     *
     * @param[in] ia cell index along a direction
     * @param[in] ib cell index along b direction
     * @param[in] ic cell index along c direction
     * @param[in] io orbital index in pc representation
     * @return index in sc representation
     *
     * @note This function does not require a buffer and is safe in OMP sections.
     */
    int64_t id_pc2sc(
        const cell_index_t& ia,
        const cell_index_t& ib,
        const cell_index_t& ic,
        const cell_index_t& io) const
    {
        return (ia * dim_bc_ + ib * dim_c_ + ic) * num_orb_pc_ + io;
    }

    /**
     * @brief Convert orbital index from pc representation to sc representation
     *   in presence of vacancies.
     *
     * @param[in] id_pc (4,) array, index in primitive cell representation
     * @return orbital index in sc representation, -1 if id_pc is a vacancy
     *
     * @note
     *   vac_id_sc should be sorted in increasing order. Otherwise the 'else' clause
     *   will be activated too early, and the results will be WRONG!
     */
    int64_t id_pc2sc_vac(const cell_index_t* id_pc) const
    {
        // Get id_sc without considering vacancies.
        int64_t id_sc = id_pc2sc(id_pc);

        // Apply the offset
        int64_t offset = get_offset(id_sc);
        if (offset == -1) {
            id_sc = -1;
        } else {
            id_sc -= offset;
        }
        return id_sc;
    }

    // Another version of pc2sc_vac based on bi-section method.
    int64_t id_pc2sc_vac2(const cell_index_t* id_pc) const
    {
        // Get id_sc without considering vacancies.
        int64_t id_sc = id_pc2sc(id_pc);

        // Apply the offset
        int64_t offset = get_offset2(id_sc);
        if (offset == -1) {
            id_sc = -1;
        } else {
            id_sc -= offset;
        }
        return id_sc;
    }

    // Another version of pc2sc_vac2 for enumerated indices.
    // This function does not require a buffer and is safe in OMP sections.
    int64_t id_pc2sc_vac2(
        const cell_index_t& ia,
        const cell_index_t& ib,
        const cell_index_t& ic,
        const cell_index_t& io) const
    {
        // Get id_sc without considering vacancies.
        int64_t id_sc = id_pc2sc(ia, ib, ic, io);

        // Apply the offset
        int64_t offset = get_offset2(id_sc);
        if (offset == -1) {
            id_sc = -1;
        } else {
            id_sc -= offset;
        }
        return id_sc;
    }
};

// -----------------------------------------------------------------------------
//       Functions for converting index between pc and sc representations
// -----------------------------------------------------------------------------
inline void check_id_sc_array_cxx(
    const int64_t& num_orb_sc,
    //
    const int64_t& num_id,
    const int64_t* id_sc_array,
    //
    int32_t* status)
{
    for (int64_t i_orb = 0; i_orb < num_id; ++i_orb) {
        int64_t id_sc = id_sc_array[i_orb];
        if (0 > id_sc || id_sc >= num_orb_sc) {
            status[0] = -1;
            status[1] = i_orb;
            break;
        }
    }
}

template <typename cell_index_t>
void check_id_pc_array_cxx(
    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,
    //
    int32_t* status)
{
    ConstMemoryView<cell_index_t> id_pc_map(id_pc_array, 4);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc, vac_id_sc, num_vac);
    cell_index_t dim_ext[4] = { dim[0], dim[1], dim[2], num_orb_pc };

    for (int64_t i_orb = 0; i_orb < num_id; ++i_orb) {
        const cell_index_t* id_pc = id_pc_map.addr(i_orb, 0);

        // Check for IDPCIndexError
        bool exit_loop = false;
        for (int64_t i_dim = 0; i_dim < 4; ++i_dim) {
            if (0 > id_pc[i_dim] || id_pc[i_dim] >= dim_ext[i_dim]) {
                exit_loop = true;
                status[0] = -2;
                status[1] = i_orb;
                status[2] = i_dim;
                break;
            }
        }
        if (exit_loop) {
            break;
        }

        // Check for IDPCVacError
        if (num_vac > 0 && orb_map.id_pc2sc_vac2(id_pc) == -1) {
            status[0] = -1;
            status[1] = i_orb;
            status[2] = 0;
            break;
        }
    }
}

template <typename cell_index_t>
void id_sc2pc_array_cxx(
    const cell_index_t* orb_id_pc,
    //
    const int64_t& num_id,
    const int64_t* id_sc_array,
    //
    cell_index_t* id_pc_array)
{
    ConstMemoryView<cell_index_t> source(orb_id_pc, 4);
    MemoryView<cell_index_t> target(id_pc_array, 4);
    constexpr int64_t buffer_size = 4 * sizeof(cell_index_t);

#pragma omp parallel for
    for (int64_t i_orb = 0; i_orb < num_id; ++i_orb) {
        memcpy(target.addr(i_orb, 0), source.addr(id_sc_array[i_orb], 0), buffer_size);
    }
}

template <typename cell_index_t>
void id_pc2sc_array_cxx(
    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)
{
    ConstMemoryView<cell_index_t> source(id_pc_array, 4);
    OrbitalMap<cell_index_t> orb_map(dim, num_orb_pc, vac_id_sc, num_vac);

    if (num_vac == 0) {
#pragma omp parallel for
        for (int64_t i_orb = 0; i_orb < num_id; ++i_orb) {
            id_sc_array[i_orb] = orb_map.id_pc2sc(source.addr(i_orb, 0));
        }
    } else {
#pragma omp parallel for
        for (int64_t i_orb = 0; i_orb < num_id; ++i_orb) {
            id_sc_array[i_orb] = orb_map.id_pc2sc_vac2(source.addr(i_orb, 0));
        }
    }
}
} // namespace orbital_map
#endif