"""Enumerate and analyze all surviving triple-partitions of γ at k = 9
for T_2 = (m_2=9, k_2=9, chords (0,3),(3,6), SP).

Look for common structural properties:
  - Block sums modulo k = 9 (rotation invariance)
  - Block differences modulo something
  - Mapping under T_2's annular structure
  - Symmetries (cyclic, reflection)
"""
from itertools import product, permutations
from collections import defaultdict, Counter

from tire_fiber_chords import (
    fiber_distribution,
    projection_support,
    u_positions_for,
    d_positions_for,
)
from tire_fiber_chunked import projection_support_streaming


def latin_set(partition, k):
    L = set()
    if any(len(b) != 3 for b in partition):
        return None
    blocks = [sorted(b) for b in partition]
    for assignment in product(permutations((1, 2, 3)), repeat=len(blocks)):
        sigma = [0] * k
        for block, perm in zip(blocks, assignment):
            for pos, color in zip(block, perm):
                sigma[pos] = color
        L.add(tuple(sigma))
    return L


def all_partitions_into_triples(elements):
    elements = list(elements)
    if not elements:
        yield []
        return
    if len(elements) % 3 != 0:
        return
    if len(elements) == 3:
        yield [tuple(elements)]
        return
    first = elements[0]
    rest = elements[1:]
    from itertools import combinations
    for combo in combinations(rest, 2):
        triple = (first,) + combo
        remaining = [x for x in rest if x not in combo]
        for sub_partition in all_partitions_into_triples(remaining):
            yield [triple] + sub_partition


def normalize(partition):
    """Return a canonical tuple of frozensets for the partition."""
    return tuple(sorted(tuple(sorted(b)) for b in partition))


def find_surviving_partitions(k, pi_U):
    survivors = []
    for partition in all_partitions_into_triples(range(k)):
        L = latin_set(partition, k)
        if L is None:
            continue
        if L <= pi_U:
            survivors.append(tuple(tuple(sorted(b)) for b in partition))
    return survivors


def analyze_block(block, k=9):
    """Return structural properties of a 3-element block of γ-edges."""
    s = sum(block) % k
    d = sorted((b - block[0]) % k for b in block)  # differences from first
    diffs_pairwise = sorted(((b - a) % k for a in block for b in block if a != b))
    return {
        'sum_mod_k': s,
        'first_diffs': d,
        'pairwise_diffs': diffs_pairwise,
    }


def main():
    k, k_2 = 9, 9
    chords = [(0, 3), (3, 6)]
    print(f'Analyzing surviving triple-partitions at k = k_2 = {k}, chords = {chords}')
    print()

    # Compute π_U via chunked streaming (faster for n = k + k_2 = 18)
    u_pos = u_positions_for(k, k_2)
    d_pos = d_positions_for(k, k_2)
    pi_U = projection_support_streaming(k, k_2, chords, u_pos)
    print(f'|π_U| = {len(pi_U)}')
    print(f'd_positions on T_ann_prime = {d_pos}')
    print(f'u_positions on T_ann_prime = {u_pos}')
    print()

    survivors = find_surviving_partitions(k, pi_U)
    print(f'Number of surviving triple-partitions: {len(survivors)}')
    print()

    print('=== Surviving partitions ===')
    for i, p in enumerate(survivors):
        print(f'  {i+1}. {list(p)}')
        for block in p:
            props = analyze_block(block, k)
            print(f'     block {block}: sum mod {k} = {props["sum_mod_k"]}, '
                  f'pairwise diffs = {props["pairwise_diffs"]}')
        print()

    # Structural analysis
    print('=== Cyclic rotation analysis ===')
    print('For each partition, check if it is a rotation of any other:')
    canonical_to_partitions = defaultdict(list)
    for p in survivors:
        # Compute canonical form under cyclic rotation of γ
        rotations = []
        for r in range(k):
            rotated = tuple(tuple(sorted((x + r) % k for x in b)) for b in p)
            rotated_canonical = tuple(sorted(rotated))
            rotations.append(rotated_canonical)
        canonical = min(rotations)
        canonical_to_partitions[canonical].append(p)
    print(f'# distinct cyclic-rotation orbits: {len(canonical_to_partitions)}')
    for can, orbits in canonical_to_partitions.items():
        print(f'  Orbit rep: {list(can)}; members: {len(orbits)}')

    # Block sum analysis
    print()
    print('=== Block sum mod 3 / mod 9 analysis ===')
    sums_mod3 = Counter()
    sums_mod9 = Counter()
    for p in survivors:
        s3 = tuple(sorted(sum(b) % 3 for b in p))
        s9 = tuple(sorted(sum(b) % 9 for b in p))
        sums_mod3[s3] += 1
        sums_mod9[s9] += 1
    print(f'Distribution of sorted block-sums mod 3:')
    for k_, v_ in sums_mod3.items():
        print(f'  {k_}: {v_} partitions')
    print(f'Distribution of sorted block-sums mod 9:')
    for k_, v_ in sums_mod9.items():
        print(f'  {k_}: {v_} partitions')

    # Check if any survivors are "arithmetic progressions" (3-APs)
    print()
    print('=== Arithmetic progression check ===')
    ap_count = 0
    for p in survivors:
        is_ap_partition = all(
            (block[1] - block[0]) % k == (block[2] - block[1]) % k
            for block in p
        )
        if is_ap_partition:
            ap_count += 1
            print(f'  AP partition: {list(p)}')
    print(f'Total AP partitions: {ap_count}')

    # Check if block 1 and block 2 are "shifts" of each other
    print()
    print('=== Shift relationships between blocks ===')
    for p in survivors:
        if len(p) != 3:
            continue
        # Check if block[1] is a cyclic shift of block[0], etc.
        shift_dets = []
        for i in range(3):
            for j in range(3):
                if i == j: continue
                # Is block j a shift of block i?
                for r in range(k):
                    shifted = tuple(sorted((x + r) % k for x in p[i]))
                    if shifted == tuple(sorted(p[j])):
                        shift_dets.append((i, j, r))
                        break
        if shift_dets:
            print(f'  Partition {list(p)}: shift structure {shift_dets}')


if __name__ == '__main__':
    main()
