#include "mask.h"

#include <set>

namespace tbplas::tests {

Mask::Mask(const Box& starting_box, int num_grid, int num_iter)
    : num_grid_(num_grid)
{
    boxes_.push_back(starting_box);
    for (int i = 0; i < num_iter; ++i) {
        std::list<Box> new_boxes;
        for (const auto& b : boxes_) {
            new_boxes.splice(new_boxes.end(), partition_box(b));
        }
        boxes_ = new_boxes;
    }
}

std::list<Box> Mask::partition_box(const Box& box) const
{
    std::list<Box> sub_boxes;
    if (box.empty) {
        sub_boxes.push_back(box);
    } else {
        int di = static_cast<int>((box.i1 - box.i0 + 1) / num_grid_);
        int dj = static_cast<int>((box.j1 - box.j0 + 1) / num_grid_);
        for (int ii = 0; ii < num_grid_; ++ii) {
            int i0 = box.i0 + ii * di;
            int i1 = i0 + di;
            for (int jj = 0; jj < num_grid_; ++jj) {
                int j0 = box.j0 + jj * dj;
                int j1 = j0 + dj;
                bool empty = (1 <= ii) && (ii < num_grid_ - 1) && (1 <= jj) && (jj < num_grid_ - 1);
                sub_boxes.emplace_back(i0, j0, i1, j1, empty);
            }
        }
    }
    return sub_boxes;
}

void Mask::etch_prim_cell(model_t& prim_cell, const int& width) const
{
    std::set<size_t> masked_id_pc;
    for (const auto& b : boxes_) {
        if (b.empty) {
            for (int ia = b.i0; ia < b.i1; ++ia) {
                for (int ib = b.j0; ib < b.j1; ++ib) {
                    masked_id_pc.insert(ia * width + ib);
                }
            }
        }
    }
    prim_cell.remove_orbitals(masked_id_pc);
}

void Mask::etch_orbitals(
    std::vector<Orbital>& orbital_list,
    const int& width) const
{
    std::set<size_t> masked_id_pc;
    std::vector<Orbital> new_list;
    for (const auto& b : boxes_) {
        if (b.empty) {
            for (int ia = b.i0; ia < b.i1; ++ia) {
                for (int ib = b.j0; ib < b.j1; ++ib) {
                    masked_id_pc.insert(ia * width + ib);
                }
            }
        }
    }
    for (int i = 0; i < orbital_list.size(); ++i) {
        if (masked_id_pc.find(i) == masked_id_pc.end()) {
            new_list.push_back(orbital_list[i]);
        }
    }
    orbital_list = new_list;
}

} // namespace tbplas::tests