#!/usr/bin/env python3
"""Generate all paired multiplicities and private-row choices for (k,h)."""
from __future__ import annotations

import argparse
from collections import Counter
from pathlib import Path
from typing import Iterator


def partitions(total: int, length: int, lower: int = 1) -> Iterator[tuple[int, ...]]:
    """Positive integer partitions in increasing order."""
    if length == 0:
        if total == 0:
            yield ()
        return
    if length == 1:
        if total >= lower:
            yield (total,)
        return
    for first in range(lower, total // length + 1):
        for rest in partitions(total - first, length - 1, first):
            yield (first,) + rest


def align(complement: tuple[int, ...], remainder: tuple[int, ...]):
    """Align the remainder in every distinct way, sorting within equal complements."""
    counts = Counter(remainder)
    prefix = []

    def extend(index: int):
        if index == len(complement):
            yield tuple(zip(complement, prefix))
            return
        lower = prefix[-1] if index and complement[index] == complement[index - 1] else 1
        for value in sorted(counts):
            if counts[value] and value >= lower:
                counts[value] -= 1
                prefix.append(value)
                yield from extend(index + 1)
                prefix.pop()
                counts[value] += 1

    yield from extend(0)


def cases(k: int, h: int):
    for t in range(k + 1, min(3*h, 3*(k-h)) + 1):
        for complement in partitions(3*(k-h), t):
            for remainder in partitions(3*h, t):
                for paired_rows in align(complement, remainder):
                    counts = Counter(paired_rows)
                    for p1 in range(min(1, counts[1, 2]) + 1):
                        for p2 in range(min(1-p1, counts[2, 1]) + 1):
                            for doubles in range(min(k-p1-p2, counts[1, 1]) + 1):
                                private = doubles + p1 + p2
                                # Private rows and distinct attachment rows are disjoint.
                                if private + doubles > t:
                                    continue
                                remaining = counts.copy()
                                remaining[1, 1] -= doubles
                                remaining[1, 2] -= p1
                                remaining[2, 1] -= p2
                                rows = sorted(remaining.elements())
                                assert len(rows) == t - private
                                yield t, doubles, p1, p2, rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('k', type=int)
    parser.add_argument('h', type=int)
    parser.add_argument('output', type=Path, nargs='?')
    args = parser.parse_args()
    if not 5 <= args.h < args.k:
        parser.error('Expected 5 <= h < k')
    destination = args.output or (
        Path(__file__).resolve().parents[1] / 'data' / 'enumeration'
        / f'k{args.k}_h{args.h}_cases.txt'
    )
    destination.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with destination.open('w') as stream:
        for index, (t, doubles, p1, p2, rows) in enumerate(cases(args.k, args.h)):
            entry = [index, t, doubles, p1, p2] + [x for row in rows for x in row]
            stream.write(' '.join(map(str, entry)) + '\n')
            count += 1
    print(f'Initial multiplicity assignments and private rows: {count}')


if __name__ == '__main__':
    main()