/** \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>

/** (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
};

/** 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);
      }
    }
  }

  auto prettify() noexcept -> void {
    for (auto idx = 0; idx < grid_.size(); ++idx) {
      if (grid_[idx] < 32) {
        auto const pos = std::make_pair(idx % length_, idx / length_);
        prettify_sq(Square(pos, grid_[idx]));
      }
    }
  }

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

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

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

  /** 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 const int 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>;

/** Recursive part of solution finder.
 *
 * Fills \a grid with squares from \a avail_sqs, starting at the current
 * position \a pos.  \a n is the side-length of the largest square.
 *
 * Returns \c true if a solution is found, \a c false if not.
 */
auto find_solution_impl(Grid &grid, int n, Pos const &pos, Avail &avail_sqs) noexcept -> bool {
  if (x(pos) == 0 && y(pos) == grid.length()) { return true; }

  /* Walk through the possible squares from the largest that will fit down to 0.
   * We don't want to walk up from 1 to the largest because we can show that a
   * 1x1 piece will never fit along an edge (or one in from the edge) - as
   * there is no other 1x1 piece to go alongside it).  This means we know the
   * first guess is definitely wrong if we start small.
   *
   * If we know a square of side length N will fit then we know a square of side length
   * N - 1 will fit.
   */
  for (auto idx = grid.largest_square(pos, n); idx != 0; --idx) {
    if (avail_sqs[idx] == 0) { continue; }

    auto const sq = Square(pos, idx);
    --avail_sqs[idx];
    grid.add(sq);
    if (find_solution_impl(grid, n, grid.next_pos(pos, idx), avail_sqs)) { return true; }
    ++avail_sqs[idx];
    grid.clear(sq);
  }

  return false;
}

/** Find a solution to the partridge problem of size n.
 *
 * Returns the solution grid.
 */
auto find_solution(int n) -> Grid {
  auto length = triangle_num(n);
  Grid grid(length);
  Avail avail_sqs;
  for (auto i = 0; i <= n; ++i) { avail_sqs.push_back(i); }
  find_solution_impl(grid, n, std::make_pair(0, 0), avail_sqs);
  grid.prettify();
  return grid;
}

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;
}