/** \file main.cc
 *  \author Matthew Gretton-Dann
 *  \brief Solves the partirige 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>
#include <set>

namespace {
  /** (x, y) pair storing a position. */
  using Pos = std::pair<int, int>;

  /** Get x co-ordinate from position */
  auto x(Pos const &p) noexcept -> int { return p.first; }

  /** Get y co-ordinate from position */
  auto y(Pos const &p) noexcept -> int { return p.second; }

  /** 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 const &pos, int 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. */
    auto x() const noexcept -> int { return ::x(pos_); }

    /** Get y co-ordinate of closest corner to origin. */
    auto y() const noexcept -> int { return ::y(pos_); }

    /** Get side length. */
    auto length() const noexcept -> int { return length_; }

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

  /** Structure holding the results.
   */
  struct Results {
    Results(int 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;

    auto length() const noexcept -> int { return length_; }

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

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

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

          int i = sq.x() + n - 1;
          while (n != 0) {
            set(s, --i, sq.y() + 1, '0' + (n % 10));
            n /= 10;
          }
        }
      }
    }

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

  /** An N * N grid of characters. */
  struct Grid {
    /** Construct a grid of given side-length. */
    Grid(int 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 */
    auto length() const noexcept -> int { return length_; }

    /** Add a square to the grid. */
    auto add(Square const &sq) noexcept -> void {
      for (auto i = sq.x(); i < sq.x() + sq.length(); ++i) {
        for (auto j = sq.y(); j < sq.y() + sq.length(); ++j) {
          set(i, j, sq.length());
        }
      }
    }

    /** Clear a square from the grid. */
    auto clear(Square const &sq) noexcept -> void {
      for (auto i = sq.x(); i < sq.x() + sq.length(); ++i) {
        for (auto j = sq.y(); j < sq.y() + sq.length(); ++j) {
          set(i, j, empty);
        }
      }
    }

    /** \brief Get length of the largest square that fits at \a pos in the grid.
     */
    auto largest_square(Pos const &pos, int n) const noexcept -> int {
      assert(x(pos) < length_);
      assert(y(pos) < length_);

      /* 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 b = x(pos);
      // Make sure we don't go looking in the next row.
      auto e = std::min(x(pos) + n, length_);
      while (b < e) {
        if (grid_[b + y(pos) * length_] != '.') { break; }
        ++b;
      }
      // Check that this length fits vertically as well.
      auto len = b - x(pos);
      auto ye = std::min(y(pos) + len, length_);
      return ye - y(pos);
    }

    /** Get the next position to check.  n is the size of the square we just
     *  added.
     */
    auto next_pos(Pos const &pos, int n) const noexcept -> Pos {
      auto const b = grid_.begin() + x(pos) + n + y(pos) * length_;
      auto const p = std::find(b, grid_.end(), empty);
      auto const v = p - grid_.begin();
      return std::make_pair(v % length_, v / length_);
    }

  private:

    /** Set the grid position (x, y) to the character c.
     *
     * It is an error if (x, y) is already set to c - as that means we have overlapping
     * squares.
     */
    auto set(int x, int y, char c) noexcept -> void {
      assert(x < length_);
      assert(y < length_);
      assert(grid_[x + y * length_] != c);
      grid_[x + y * length_] = c;
    }

    std::vector<char> grid_; ///< The grid
    int length_; ///< Side length

    static constexpr char empty = '.'; ///< Character used for an empty cell.
  };

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

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

  /** Find a solution to the \a n th Partridge problem.
   *
   * Returns the grid of the solution.
   */
  auto find_solution(int 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 inprogress 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, 0};
    int 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 one 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 = std::make_pair(sq.x(), sq.y());
        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 (x(pos) == 0 && y(pos) == grid.length()) { break; }
    }

    return Results(length, sqs);
  }
} // anon namespace

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