/** \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 valley_capacity_checks = 0;
    size_t valley_capacity_prunes = 0;
    size_t large_square_checks = 0;
    size_t large_square_prunes = 0;
    size_t generated_tasks = 0;
    size_t completed_tasks = 0;
  };

  /** Deterministic order in which a skyline node tries fitting squares. */
  enum class SearchPolicy {
    ascending,
    descending,
    /** Close the selected valley when possible, then try smaller sizes first. */
    best_fit,
  };

  /** Whether to remove equivalent board orientations from the search. */
  enum class SymmetryBreaking {
    disabled,
    d4_unit_square,
  };

  /** Cheap necessary conditions applied before branching at a search node. */
  enum class Pruning {
    disabled = 0,
    valley_capacity = 1,
    large_square = 2,
    all = 3,
  };

  /** Return whether a cell is the canonical representative of its D4 orbit.
   *
   * Reflect a cell into the left half of the board, rotate so its distance
   * from the left edge is no greater than its distance from the top edge, and
   * finally reflect it into the top half.  The resulting fundamental region is
   * the closed triangle x <= y <= floor((length - 1) / 2).  Closed boundaries
   * retain representatives whose orbit is smaller than eight.
   */
  [[nodiscard]] auto canonical_unit_position(size_t const x,
                                             size_t const y,
                                             size_t const length) noexcept
      -> bool {
    assert(x < length);
    assert(y < length);
    return x <= y && y <= (length - 1) / 2;
  }

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

  /** Return whether the remaining narrow pieces can fill a valley to its rim.
   *
   * Until the valley reaches the lower of its two neighbouring heights, no
   * square wider than the valley can enter it.  The area of all remaining
   * squares no wider than the valley is therefore an upper bound on the area
   * available to fill this isolated strip.  This includes the familiar
   * width-one and width-two gap checks without special cases.
   */
  [[nodiscard]] auto valley_has_sufficient_capacity(
      Valley const valley, std::vector<size_t> const &skyline,
      Avail const &available) noexcept -> bool {
    auto const end = valley.x + valley.width;
    auto const left_height =
        valley.x == 0 ? skyline.size() : skyline[valley.x - 1];
    auto const right_height =
        end == skyline.size() ? skyline.size() : skyline[end];
    auto const rim_height = std::min(left_height, right_height);
    auto const required_area = valley.width * (rim_height - valley.height);

    size_t available_area = 0;
    auto const largest = std::min(
        valley.width, static_cast<size_t>(available.size() - 1));
    for (size_t side = 1; side <= largest; ++side) {
      available_area += available[side] * side * side;
    }
    return available_area >= required_area;
  }

  /** Return whether the skyline contains an empty box for a square.
   *
   * A side-k square fits exactly when k consecutive columns have heights no
   * greater than board height minus k.  Tracking the current qualifying run
   * checks this in one pass without allocations.
   */
  [[nodiscard]] auto skyline_has_empty_square(
      std::vector<size_t> const &skyline, size_t const side) noexcept -> bool {
    assert(side > 0);
    assert(side <= skyline.size());
    auto const maximum_height = skyline.size() - side;
    size_t run = 0;
    for (auto const height: skyline) {
      run = height <= maximum_height ? run + 1 : 0;
      if (run == side) {
        return true;
      }
    }
    return false;
  }

  /** Return whether the largest remaining square has any feasible position.
   *
   * Feasible empty boxes are monotonic in the side length: a box which fits
   * the largest remaining square also fits every smaller one.  It is therefore
   * sufficient to test only the largest size with non-zero multiplicity.
   */
  [[nodiscard]] auto remaining_large_square_fits(
      std::vector<size_t> const &skyline, Avail const &available) noexcept
      -> bool {
    for (auto side = available.size() - 1; side != 0; --side) {
      if (available[side] != 0) {
        return skyline_has_empty_square(skyline, side);
      }
    }
    return true;
  }

  template<bool Instrument, bool BreakD4Symmetry, bool PruneValleyCapacity,
           bool PruneLargeSquare>
  auto search_skyline(size_t const n, size_t const length,
                      SearchPolicy const policy,
                      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);
    if constexpr (PruneValleyCapacity) {
      if constexpr (Instrument) {
        ++counters->prune_checks;
        ++counters->valley_capacity_checks;
      }
      if (!valley_has_sufficient_capacity(valley, skyline, available)) {
        if constexpr (Instrument) {
          ++counters->prune_hits;
          ++counters->valley_capacity_prunes;
        }
        return false;
      }
    }
    if constexpr (PruneLargeSquare) {
      if constexpr (Instrument) {
        ++counters->prune_checks;
        ++counters->large_square_checks;
      }
      if (!remaining_large_square_fits(skyline, available)) {
        if constexpr (Instrument) {
          ++counters->prune_hits;
          ++counters->large_square_prunes;
        }
        return false;
      }
    }
    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 (BreakD4Symmetry) {
        if (side == 1) {
          if constexpr (Instrument) {
            ++counters->prune_checks;
          }
          if (!canonical_unit_position(valley.x, valley.height, length)) {
            if constexpr (Instrument) {
              ++counters->prune_hits;
            }
            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, BreakD4Symmetry, PruneValleyCapacity,
                         PruneLargeSquare>(
              n, length, policy, 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 (policy == SearchPolicy::best_fit && largest == valley.width &&
        try_side(largest)) {
      return true;
    }
    if (policy != SearchPolicy::descending) {
      for (size_t side = 1; side <= largest; ++side) {
        if (policy == SearchPolicy::best_fit && side == valley.width) {
          continue;
        }
        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, bool BreakD4Symmetry, bool PruneValleyCapacity,
           bool PruneLargeSquare>
  auto search_solution_impl(size_t const n, SearchPolicy const policy,
                            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, BreakD4Symmetry, PruneValleyCapacity,
                       PruneLargeSquare>(
            n, length, policy, skyline, available, squares, counters));

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

  template<bool Instrument, bool BreakD4Symmetry>
  auto search_solution_dispatch(size_t const n, SearchPolicy const policy,
                                Pruning const pruning,
                                SearchCounters *const counters) noexcept
      -> Results {
    switch (pruning) {
      case Pruning::all:
        return search_solution_impl<Instrument, BreakD4Symmetry, true, true>(
            n, policy, counters);
      case Pruning::valley_capacity:
        return search_solution_impl<Instrument, BreakD4Symmetry, true, false>(
            n, policy, counters);
      case Pruning::large_square:
        return search_solution_impl<Instrument, BreakD4Symmetry, false, true>(
            n, policy, counters);
      case Pruning::disabled:
        return search_solution_impl<Instrument, BreakD4Symmetry, false, false>(
            n, policy, counters);
    }
    assert(false);
    return search_solution_impl<Instrument, BreakD4Symmetry, false, false>(
        n, policy, counters);
  }

  auto search_solution(
      size_t const n,
      SearchPolicy const policy = SearchPolicy::ascending,
      SymmetryBreaking const symmetry =
          SymmetryBreaking::d4_unit_square,
      Pruning const pruning = Pruning::all) noexcept
      -> Results {
    if (symmetry == SymmetryBreaking::d4_unit_square) {
      return search_solution_dispatch<false, true>(
          n, policy, pruning, nullptr);
    }
    return search_solution_dispatch<false, false>(
        n, policy, pruning, nullptr);
  }

  auto search_solution_instrumented(size_t const n,
                                    SearchCounters &counters,
                                    SearchPolicy const policy =
                                        SearchPolicy::ascending,
                                    SymmetryBreaking const symmetry =
                                        SymmetryBreaking::d4_unit_square,
                                    Pruning const pruning =
                                        Pruning::all) noexcept
      -> Results {
    counters = {};
    if (symmetry == SymmetryBreaking::d4_unit_square) {
      return search_solution_dispatch<true, true>(
          n, policy, pruning, &counters);
    }
    return search_solution_dispatch<true, false>(
        n, policy, pruning, &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,
      SearchPolicy const policy = SearchPolicy::ascending,
      SymmetryBreaking const symmetry =
          SymmetryBreaking::d4_unit_square,
      Pruning const pruning = Pruning::all) noexcept -> Results {
    if (uses_odd_construction(n)) {
      return construct_odd_solution(
          n, search_solution(n - 1, policy, symmetry, pruning));
    }
    return search_solution(n, policy, symmetry, pruning);
  }

  auto find_solution_instrumented(size_t const n,
                                  SearchCounters &counters,
                                  SearchPolicy const policy =
                                      SearchPolicy::ascending,
                                  SymmetryBreaking const symmetry =
                                      SymmetryBreaking::d4_unit_square,
                                  Pruning const pruning =
                                      Pruning::all) noexcept
      -> Results {
    if (uses_odd_construction(n)) {
      return construct_odd_solution(
          n, search_solution_instrumented(
                 n - 1, counters, policy, symmetry, pruning));
    }
    return search_solution_instrumented(
        n, counters, policy, symmetry, pruning);
  }
} // 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