"""
Count (an upper bound on) the number of tinyhouse chess positions.

Board: 4x4 = 16 squares. Ranks 2 and 3 are the 8 "middle" squares.

Piece inventory (conserved totals, both colours + banks combined):
  1 white king, 1 black king
  2 "pawn-ish" pieces  : each is a pawn, or a promoted pawn (pawn-wazir /
                         pawn-ferz / pawn-horse)
  2 true wazirs, 2 true ferzes, 2 true horses

Placement domains:
  king, true wazir/ferz/horse ....... any of the 16 squares (or bank, except kings)
  plain pawn ........................ one of the 8 middle squares, or the bank
  promoted pawn (pawn-X) ............ any of the 16 squares, never the bank
  at most one piece per square

Method
------
Every piece is described by a "slot option": (label, place) where
    label = (colour, kind)          -- what the piece is and whose it is
    place = 'M' (middle 8 squares) | 'F' (free: any of 16) | 'bank'

Pieces sharing the same label+place are indistinguishable, so for each group
we enumerate *multisets* of options (combinations_with_replacement).

Given a whole configuration with `a` on-board pieces restricted to the middle
squares and `b` on-board pieces free to use all 16, the number of ways to place
them as *distinguishable* pieces is
        P(8, a) * P(16 - a, b)
(place the middle-restricted ones first; each one eats a square from the 16).
Dividing by k! for every group of k identical on-board pieces turns that into a
count of distinct positions.  Identical pieces sitting in the same bank are
already handled by the multiset enumeration, so they need no divisor.

Finally multiply by 2 for whose turn it is.
"""

from itertools import combinations_with_replacement, product
from collections import Counter
from math import perm, factorial

COLOURS = ("W", "B")

# ---- option lists for each conserved group -------------------------------

# the two pawn-ish pieces
PAWN_OPTIONS = (
    [((c, "pawn"), "M") for c in COLOURS]                       # pawn on board
    + [((c, "pawn"), "bank") for c in COLOURS]                  # pawn in hand
    + [((c, "pawn-" + p), "F")                                  # promoted pawn
       for c in COLOURS for p in ("wazir", "ferz", "horse")]
)


def true_piece_options(kind):
    return ([((c, kind), "F") for c in COLOURS]
            + [((c, kind), "bank") for c in COLOURS])


WAZIR_OPTIONS = true_piece_options("wazir")
FERZ_OPTIONS = true_piece_options("ferz")
HORSE_OPTIONS = true_piece_options("horse")

KINGS = [(("W", "king"), "F"), (("B", "king"), "F")]

# ---- enumeration ---------------------------------------------------------


def count_positions(verbose=False):
    groups = [
        list(combinations_with_replacement(PAWN_OPTIONS, 2)),
        list(combinations_with_replacement(WAZIR_OPTIONS, 2)),
        list(combinations_with_replacement(FERZ_OPTIONS, 2)),
        list(combinations_with_replacement(HORSE_OPTIONS, 2)),
    ]

    total = 0
    by_pieces_on_board = Counter()      # how many pieces stand on the board
    configs = 0

    for combo in product(*groups):
        pieces = KINGS + [item for grp in combo for item in grp]

        a = sum(1 for _, place in pieces if place == "M")
        b = sum(1 for _, place in pieces if place == "F")

        placements = perm(8, a) * perm(16 - a, b)
        if placements == 0:
            continue

        # divide out permutations of identical on-board pieces
        board_counts = Counter(p for p in pieces if p[1] != "bank")
        divisor = 1
        for k in board_counts.values():
            divisor *= factorial(k)

        n = placements // divisor
        total += n
        configs += 1
        by_pieces_on_board[a + b] += n

    total *= 2                                  # side to move

    if verbose:
        print(f"piece configurations considered : {configs}")
        print("positions by number of pieces standing on the board "
              "(before the x2 for side to move):")
        for k in sorted(by_pieces_on_board):
            print(f"   {k:2d} pieces on board : {by_pieces_on_board[k]:>14,}")
        print()
    return total


if __name__ == "__main__":
    t = count_positions(verbose=True)
    print(f"TOTAL tinyhouse positions (upper bound): {t:,}")
    print(f"                                       ~ {t:.4g}")