"""
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}")