/**
 * @file Space.cpp
 * @author marisa <marisa@fwmari.net>
 * @brief Space definitions
 * @version 0.1
 * @date 2023-08-06
 *
 * Copyright (c) 2023 marisa <https://git.fwmari.net/marisa/MolecularDynamics2>
 *
 */

#include <iostream>
#include <stdexcept>

#include "Space.hpp"

Space::Space(double cutoff) {
  if (cutoff <= 0.0)
    throw std::invalid_argument("cutoff must be positive");

  this->m_cutoff = cutoff;
}

void Space::addAtom(Atom atom) {
  if (this->m_locked)
    throw std::runtime_error("cannot add atom to locked space");

  this->m_atoms.push_back(atom);
}

void Space::addAtom(std::vector<Atom> atoms) {
  if (this->m_locked)
    throw std::runtime_error("cannot add atoms to locked space");

  this->m_atoms.insert(this->m_atoms.end(), atoms.begin(), atoms.end());
}

void Space::setBox(double x, double y, double z) {
  if (this->m_locked)
    throw std::runtime_error("cannot set box of locked space");

  if (x <= 0.0 || y <= 0.0 || z <= 0.0)
    throw std::invalid_argument("box dimensions must be positive");

  this->m_box.x = x;
  this->m_box.y = y;
  this->m_box.z = z;
}

void Space::prepareSpace() {
  if (this->m_locked)
    throw std::runtime_error("cannot prepare locked space");

  if (this->m_atoms.empty())
    throw std::runtime_error("cannot prepare space with no atoms");

  if (this->m_box.x <= 0.0 || this->m_box.y <= 0.0 || this->m_box.z <= 0.0)
    throw std::runtime_error("cannot prepare space with invalid box");

  if (this->m_cutoff <= 0.0)
    throw std::runtime_error("cannot prepare space with invalid cutoff");

  // Calculate the number of cells in each dimension
  this->m_cellsDim.x = static_cast<int>(this->m_box.x / this->m_cutoff);
  this->m_cellsDim.y = static_cast<int>(this->m_box.y / this->m_cutoff);
  this->m_cellsDim.z = static_cast<int>(this->m_box.z / this->m_cutoff);

  this->m_locked = true;
}

void Space::buildCells() {
  if (!this->m_locked)
    throw std::runtime_error("cannot build cells of unlocked space");

  this->m_cells.clear();

  for (int i = 0; i < this->m_cellsDim.x; i++)
    for (int j = 0; j < this->m_cellsDim.y; j++)
      for (int k = 0; k < this->m_cellsDim.z; k++) {
        Cell cell;
        cell.m_x   = i * this->m_cutoff;
        cell.m_y   = j * this->m_cutoff;
        cell.m_z   = k * this->m_cutoff;
        cell.m_idx = i * this->m_cellsDim.y * this->m_cellsDim.z + j * this->m_cellsDim.z + k;
        this->m_cells.push_back(cell);
      }

  std::cout << this->m_cells.size() << " cells built" << std::endl;

  for (std::size_t a = 0; a < this->m_atoms.size(); a++) {
    Atom &atom = this->m_atoms[a];

    int i = static_cast<int>(atom.x / this->m_cutoff);
    int j = static_cast<int>(atom.y / this->m_cutoff);
    int k = static_cast<int>(atom.z / this->m_cutoff);

    int idx = i * this->m_cellsDim.y * this->m_cellsDim.z + j * this->m_cellsDim.z + k;

    Cell &cell = this->m_cells[idx];

    cell.x.push_back(atom.x);
    cell.y.push_back(atom.y);
    cell.z.push_back(atom.z);
    cell.idx.push_back(a);
  }
}

void Space::printCells() const {
  for (const Cell &cell : this->m_cells)
    cell.printCell();
}

Cell Space::getNeighboringCells(int idx) const {
  Cell neighboringCells;

  const Cell &cell = this->m_cells[idx];

  // For each cell adjacent to the cell of the given index, add its atoms to neighboringCells
  for (int i = -1; i <= 1; i++)
    for (int j = -1; j <= 1; j++)
      for (int k = -1; k <= 1; k++) {
        int x = cell.m_x + i * this->m_cutoff;
        int y = cell.m_y + j * this->m_cutoff;
        int z = cell.m_z + k * this->m_cutoff;

        if (x < 0 || x >= this->m_box.x || y < 0 || y >= this->m_box.y || z < 0 ||
            z >= this->m_box.z)
          continue;

        std::cout << "Adding cell (" << x << ", " << y << ", " << z << ") to neighboringCells"
                  << std::endl;

        int idx = i * this->m_cellsDim.y * this->m_cellsDim.z + j * this->m_cellsDim.z + k;

        const Cell &neighboringCell = this->m_cells[idx];

        neighboringCells.x.insert(
            neighboringCells.x.end(), neighboringCell.x.begin(), neighboringCell.x.end());
        neighboringCells.y.insert(
            neighboringCells.y.end(), neighboringCell.y.begin(), neighboringCell.y.end());
        neighboringCells.z.insert(
            neighboringCells.z.end(), neighboringCell.z.begin(), neighboringCell.z.end());
        neighboringCells.idx.insert(
            neighboringCells.idx.end(), neighboringCell.idx.begin(), neighboringCell.idx.end());
      }

  return neighboringCells;
}

void Cell::printCell() const {
  std::cout << "Cell " << this->m_idx << " (" << this->m_x << ", " << this->m_y << ", " << this->m_z
            << "):" << std::endl;

  for (std::size_t i = 0; i < this->idx.size(); i++)
    std::cout << "  " << this->idx[i] << ": (" << this->x[i] << ", " << this->y[i] << ", "
              << this->z[i] << ")" << std::endl;
}