/** \file main.cc
 *  \author Matthew Gretton-Dann
 *  \brief Solves the Partridge problem for user specified size.
 *
 * Copyright 2025, Matthew-Gretton-Dann
 * SPDX: Apache-2.0
 */

#include <cassert>
#include <utility>
#include <string_view>
#include <vector>
#include <iostream>

namespace {
  using size_t = std::uint64_t;

  /** (x, y) pair storing a position. */
  using Pos = size_t;

  /** A square - consisting of position of closest corner to origin, and side-length.
   */
  struct Square {
    /** Construct a square.
     *  \param pos    Position of closest corner to origin
     *  \param length Side length.
     */
    Square(Pos pos, size_t const length) noexcept : pos_(pos), length_(length) {
    }

    Square(Square const &other) noexcept = default;

    Square(Square &&other) noexcept = default;

    Square &operator=(Square const &other) noexcept = default;

    Square &operator=(Square &&other) noexcept = default;

    ~Square() noexcept = default;

    /** Get x co-ordinate of closest corner to origin. */
    [[nodiscard]] auto pos() const noexcept -> Pos { return pos_; }

    /** Get side length. */
    [[nodiscard]] auto length() const noexcept -> size_t { return length_; }

  private:
    Pos pos_; ///< Position of corner closest to origin
    size_t length_; ///< Side length
  };

  /** Structure holding the results.
   */
  struct Results {
    Results(size_t length, std::vector<Square> squares) : length_(length), squares_(std::move(squares)) {
    }

    Results(Results const &other) noexcept = delete;

    Results &operator=(Results const &other) noexcept = delete;

    Results &operator=(Results &&other) noexcept = default;

    Results(Results &&other) noexcept = default;

    ~Results() noexcept = default;

    [[nodiscard]] auto length() const noexcept -> size_t { return length_; }

    /** Output the grid. */
    auto output() const -> void {
      std::string out(length_ * length_, '.');
      for (auto const &sq: squares_) {
        prettify_sq(out, sq);
      }
      for (size_t idx = 0; idx < length_ * length_; idx += length_) {
        std::cout << std::string_view(out.data() + idx, length_) << '\n';
      }
    }

  private:
    auto set(std::string &s, size_t x, size_t y, char c) const noexcept -> void {
      assert(x < length_);
      assert(y < length_);
      assert(grid_[x + y * length_] != c);
      s[x + y * length_] = c;
    }

    [[nodiscard]] auto sq_x(Square const &sq) const noexcept -> size_t { return sq.pos() % length_; }
    [[nodiscard]] auto sq_y(Square const &sq) const noexcept -> size_t { return sq.pos() / length_; }

    auto prettify_sq(std::string &s, Square const &sq) const noexcept -> void {
      switch (sq.length()) {
        case 1: set(s, sq_x(sq), sq_y(sq), '*');
          break;
        case 2: set(s, sq_x(sq), sq_y(sq), '+');
          set(s, sq_x(sq) + 1, sq_y(sq), '+');
          set(s, sq_x(sq), sq_y(sq) + 1, '+');
          set(s, sq_x(sq) + 1, sq_y(sq) + 1, '+');
          break;
        default: {
          auto n = sq.length();
          set(s, sq_x(sq), sq_y(sq), '+');
          set(s, sq_x(sq) + n - 1, sq_y(sq), '+');
          set(s, sq_x(sq), sq_y(sq) + n - 1, '+');
          set(s, sq_x(sq) + n - 1, sq_y(sq) + n - 1, '+');
          for (size_t i = 1; i < n - 1; ++i) {
            set(s, sq_x(sq) + i, sq_y(sq), '-');
            set(s, sq_x(sq) + i, sq_y(sq) + n - 1, '-');
            set(s, sq_x(sq), sq_y(sq) + i, '|');
            for (size_t j = 1; j < n - 1; ++j) {
              set(s, sq_x(sq) + j, sq_y(sq) + i, ' ');
            }
            set(s, sq_x(sq) + n - 1, sq_y(sq) + i, '|');
          }

          size_t i = sq_x(sq) + n - 1;
          while (n != 0) {
            set(s, --i, sq_y(sq) + 1, static_cast<char>('0' + static_cast<char>(n % 10)));
            n /= 10;
          }
        }
      }
    }

    size_t length_;
    std::vector<Square> squares_;
  };

  /** An N * N grid of characters. */
  struct Grid {
    // Type to use for the grid contents
    using T = std::int_fast64_t;

    /** Construct a grid of given side-length. */
    explicit Grid(size_t length) : grid_(length * length, empty), length_(length) {
    }

    Grid(Grid const &other) = delete;

    Grid(Grid &&other) noexcept = default;

    Grid &operator=(Grid const &other) = delete;

    Grid &operator=(Grid &&other) noexcept = default;

    ~Grid() noexcept = default;

    /** Get grid length */
    [[nodiscard]] auto end() const noexcept -> size_t { return static_cast<size_t>(grid_.size()); }

    /** Add a square to the grid. */
    auto add(Square const &sq) noexcept -> void {
      /* One would expect the fastest way to do this would be to have x be the
       * fastest increasing index so we store [pos, pos + 1,..., pos+length, ...]
       * But experimentation tells us this isn't so, and storing
       * [pos, pos + length, ..., pos + 1, ...] is faster!
       */
      for (auto x = 0; x < sq.length(); ++x) {
        for (auto y = sq.pos(); y < sq.pos() + sq.length() * length_; y += length_) {
          grid_[x + y] = filled;
        }
      }
    }

    /** Clear a square from the grid. */
    auto clear(Square const &sq) noexcept -> void {
      for (auto x = 0; x < sq.length(); ++x) {
        for (auto y = sq.pos(); y < sq.pos() + sq.length() * length_; y += length_) {
          grid_[x + y] = empty;
        }
      }
    }

    /** \brief Get length of the largest square that fits at \a pos in the grid.
     */
    [[nodiscard]] auto largest_square(Pos pos, size_t n) const noexcept -> size_t {
      assert(pos < end());

      /* Because of how we walk through the grid (starting at 0,0 then increasing
       * x followed by y) we can assume that if the position (b, y) is clear
       * (i.e. a '.') then (b, y + i) is clear for all i > 0.
       *
       * This means we only need to look for the first non-clear position along the
       * current row.
       */
      auto const pos_x = pos % length_;
      auto const pos_y0 = pos - pos_x;
      auto b = pos;
      // Make sure we don't go looking in the next row.
      auto const e = std::min(pos + n, pos_y0 + length_);
      while (b < e) {
        if (grid_[b] != empty) { break; }
        ++b;
      }
      // Check that this length fits vertically as well.
      auto const len = b - pos;
      auto const pos_y = pos / length_;
      auto const ye = std::min(pos_y + len, length_);
      return ye - pos_y;
    }

    /** Get the next position to check starting at pos.
     *
     * Returns grid_.length() if no more positions available.
     */
    [[nodiscard]] auto next_pos(Pos pos) const noexcept -> Pos {
      auto const b = grid_.begin() + static_cast<std::ptrdiff_t>(pos);
      auto const p = std::find(b, grid_.end(), empty);
      return p - grid_.begin();
    }

  private:
    std::vector<T> grid_; ///< The grid
    size_t length_; ///< Side length

    static constexpr char empty = 0; ///< Character used for an empty cell.
    static constexpr char filled = 1; ///< Character used for a filled cell,
  };

  /** Get the n-th triangular number. */
  auto triangle_num(size_t n) noexcept -> size_t { return (n * (n + 1)) / 2; }

  /** Vector used to identify the available squares. */
  using Avail = std::vector<size_t>;

  /** Find a solution to the \a n th Partridge problem.
   *
   * Returns the grid of the solution.
   */
  auto find_solution(size_t const n) noexcept -> Results {
    /* Implementation is iterative, as opposed to recursive.
     *
     * The recursive implementation is easier to understand - but is
     * slightly slower because of the repeated function calls (and
     * entry/exit).
     *
     * The basic algorithm is to start at the origin of the grid we
     * want to place squares on and iterate over the permutations of
     * available squares until we find one that fits.
     */

    // grid is our in-progress grid of square positions.
    auto const length = triangle_num(n);
    Grid grid(length);

    /* avail_sqs is a vector indexed by square length indicating how many
     * squares are available.  Initially set up so that avail_sqs[i] = i.
     */
    Avail avail_sqs;
    for (auto i = 0; i <= n; ++i) { avail_sqs.push_back(i); }

    /* sqs is a vector used as a stack of the squares currently placed.
     * We reserve the length we need so as not to have too many allocations.
     */
    std::vector<Square> sqs;
    sqs.reserve(length);

    // Start at the origin with a square of longest side length.
    Pos pos = 0;
    size_t idx = n;

    while (true) {
      /* If the idx is 0 we've looked at all possible square lengths for this
       * position, and they've failed.  Pop the last square of the stack, remove
       * it from the grid and try the next smaller size in the same position.
       */
      if (idx == 0) {
        // No squares on the stack -> failed to find a solution.
        if (sqs.empty()) { break; }

        auto sq = sqs.back();
        sqs.pop_back();
        grid.clear(sq);
        ++avail_sqs[sq.length()];
        pos = sq.pos();
        idx = sq.length() - 1;
        continue;
      }

      // If there are no squares available of the current size try the next one.
      if (avail_sqs[idx] == 0) {
        --idx;
        continue;
      }

      /* Place a square of side length idx at pos, push this onto the stack and
       * set up to look at the next position.
       */
      auto const sq = Square(pos, idx);
      --avail_sqs[idx];
      grid.add(sq);
      sqs.push_back(sq);

      pos = grid.next_pos(pos + idx);
      idx = grid.largest_square(pos, n);

      // Have we reached the end?  If so success!
      if (pos == grid.end()) { break; }
    }

    return {length, sqs};
  }
} // anon namespace

int main(int argc, char **argv) {
  auto n = (argc == 1) ? 8 : std::atol(argv[1]);
  auto const grid = find_solution(n);
  std::cout << "Partridge problem " << n << " side length " << grid.length() << '\n';
  grid.output();
  return 0;
}