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

    /** Get the square placements in this result. */
    [[nodiscard]] auto squares() const noexcept -> std::vector<Square> const & { return squares_; }

    /** 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_);
      // Size labels may replace interior spaces, but must not duplicate a write.
      assert(s[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_;
  };

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

  /** Optional search instrumentation.
   *
   * Counters for search features which are not implemented by the current
   * single-threaded solver remain zero.  Keeping them in the stable output
   * schema lets later solver implementations remain comparable.
   */
  struct SearchCounters {
    size_t search_nodes = 0;
    size_t loop_iterations = 0;
    size_t attempted_placements = 0;
    size_t backtracks = 0;
    size_t prune_checks = 0;
    size_t prune_hits = 0;
    size_t generated_tasks = 0;
    size_t completed_tasks = 0;
  };

  enum class CandidateOrder {
    ascending,
    descending,
  };

  /** A maximal level skyline segment which is lower than its neighbours. */
  struct Valley {
    size_t x;
    size_t height;
    size_t width;
  };

  /** Return the narrowest local valley, breaking ties by height then x. */
  [[nodiscard]] auto smallest_valley(std::vector<size_t> const &skyline) noexcept
      -> Valley {
    Valley best{0, 0, skyline.size()};
    bool found = false;
    for (size_t begin = 0; begin < skyline.size();) {
      auto end = begin + 1;
      while (end < skyline.size() && skyline[end] == skyline[begin]) {
        ++end;
      }
      auto const left_height =
          begin == 0 ? skyline.size() : skyline[begin - 1];
      auto const right_height =
          end == skyline.size() ? skyline.size() : skyline[end];
      auto const lower_than_left = skyline[begin] < left_height;
      auto const lower_than_right = skyline[begin] < right_height;
      auto const width = end - begin;
      if (lower_than_left && lower_than_right &&
          (!found || width < best.width ||
           (width == best.width && skyline[begin] < best.height) ||
           (width == best.width && skyline[begin] == best.height &&
            begin < best.x))) {
        best = {begin, skyline[begin], width};
        found = true;
      }
      begin = end;
    }
    assert(found);
    return best;
  }

  template<bool Instrument>
  auto search_skyline(size_t const n, size_t const length,
                      CandidateOrder const candidate_order,
                      std::vector<size_t> &skyline, Avail &available,
                      std::vector<Square> &squares,
                      SearchCounters *const counters) noexcept -> bool {
    if constexpr (Instrument) {
      assert(counters != nullptr);
      ++counters->search_nodes;
    }

    if (squares.size() == length) {
      return true;
    }

    auto const valley = smallest_valley(skyline);
    auto const largest =
        std::min({n, valley.width, length - valley.height});
    auto try_side = [&](size_t const side) {
      if constexpr (Instrument) {
        ++counters->loop_iterations;
      }
      if (available[side] == 0) {
        return false;
      }

      if constexpr (Instrument) {
        ++counters->attempted_placements;
      }
      --available[side];
      std::fill_n(skyline.begin() + static_cast<std::ptrdiff_t>(valley.x),
                  side, valley.height + side);
      squares.emplace_back(valley.x + valley.height * length, side);

      if (search_skyline<Instrument>(n, length, candidate_order, skyline,
                                     available, squares, counters)) {
        return true;
      }

      squares.pop_back();
      std::fill_n(skyline.begin() + static_cast<std::ptrdiff_t>(valley.x),
                  side, valley.height);
      ++available[side];
      if constexpr (Instrument) {
        ++counters->backtracks;
      }
      return false;
    };

    if (candidate_order == CandidateOrder::ascending) {
      for (size_t side = 1; side <= largest; ++side) {
        if (try_side(side)) {
          return true;
        }
      }
    } else {
      for (auto side = largest; side != 0; --side) {
        if (try_side(side)) {
          return true;
        }
      }
    }
    return false;
  }

  /** Search directly for a solution to the \a n th Partridge problem.
   *
   * The state is one filled height per board column.  At each node the
   * narrowest local valley is found by a linear scan and candidates are tried
   * at its far-left edge.  A placement or undo touches one entry per square
   * column.  Scanning the profile costs O(board width); trying up to n
   * candidates and updating up to n columns for each costs O(n^2), for
   * O(board width + n^2) local work per node and the same total state.
   */
  template<bool Instrument>
  auto search_solution_impl(size_t const n, CandidateOrder const candidate_order,
                            SearchCounters *const counters) noexcept
      -> Results {
    auto const length = triangle_num(n);
    std::vector<size_t> skyline(length);
    Avail available(n + 1);
    for (size_t side = 0; side <= n; ++side) {
      available[side] = side;
    }
    std::vector<Square> squares;
    squares.reserve(length);
    static_cast<void>(search_skyline<Instrument>(
        n, length, candidate_order, skyline, available, squares, counters));

    return {length, std::move(squares)};
  }

  auto search_solution(
      size_t const n,
      CandidateOrder const candidate_order = CandidateOrder::ascending) noexcept
      -> Results {
    return search_solution_impl<false>(n, candidate_order, nullptr);
  }

  auto search_solution_instrumented(size_t const n,
                                    SearchCounters &counters,
                                    CandidateOrder const candidate_order =
                                        CandidateOrder::ascending) noexcept
      -> Results {
    counters = {};
    return search_solution_impl<true>(n, candidate_order, &counters);
  }

  /** Construct an odd-order solution from its even-order predecessor. */
  auto construct_odd_solution(size_t const odd_order, Results predecessor)
      -> Results {
    assert(odd_order >= 9);
    assert(odd_order % 2 == 1);
    assert(predecessor.length() == triangle_num(odd_order - 1));

    auto const old_length = predecessor.length();
    auto const new_length = triangle_num(odd_order);
    std::vector<Square> squares;
    squares.reserve(predecessor.squares().size() + odd_order);
    for (auto const &square: predecessor.squares()) {
      auto const x = square.pos() % old_length;
      auto const y = square.pos() / old_length;
      squares.emplace_back(x + y * new_length, square.length());
    }

    for (size_t y = 0; y < old_length; y += odd_order) {
      squares.emplace_back(old_length + y * new_length, odd_order);
    }
    for (size_t x = 0; x <= old_length; x += odd_order) {
      squares.emplace_back(x + old_length * new_length, odd_order);
    }
    return {new_length, std::move(squares)};
  }

  [[nodiscard]] auto uses_odd_construction(size_t const n) noexcept -> bool {
    return n >= 9 && n % 2 == 1;
  }

  auto find_solution(size_t const n) noexcept -> Results {
    if (uses_odd_construction(n)) {
      return construct_odd_solution(n, search_solution(n - 1));
    }
    return search_solution(n);
  }

  auto find_solution_instrumented(size_t const n,
                                  SearchCounters &counters) noexcept -> Results {
    if (uses_odd_construction(n)) {
      return construct_odd_solution(
          n, search_solution_instrumented(n - 1, counters));
    }
    return search_solution_instrumented(n, counters);
  }
} // anon namespace

#ifndef PARTRIDGE_TESTING
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;
}
#endif