"""Verify the candidate induced γ-partition conjecture for T_2 (the inner
tire whose chord is on B_in^(2), not on γ).

Setup:
  T_2 is an SP tire with B_out^(2) = γ (length k) and B_in^(2) = C_{k_2}
  (length k_2 divisible by 3) plus chord(s) creating O^(2)-faces of
  size 3 each.  Balanced annular triangulation interleaves D-triangles
  (one per B_in^(2)-edge) and U-triangles (one per γ-edge) on the dual
  cycle T'_ann of length k + k_2.

The conjecture is that π_U(C(T_2)) ⊇ L(γ, ~F_2), the Latin subset for
some induced γ-face partition ~F_2.  Candidate construction (from
worst_case_proof_sketch.tex): for each O^(2)-face F^(2) (a triple of
B_in^(2)-edges = a triple of D-triangles), the associated γ-edges are
the U-triangles that lie in the cyclic span of F^(2)'s D-positions on
T'_ann, with boundary U-triangles assigned by some rule (e.g. "next
F^(2) cyclically").

This script:
  - Constructs the candidate induced partition.
  - Computes π_U(C(T_2)) by brute enumeration.
  - Checks whether L(γ, induced partition) ⊆ π_U.
  - Reports findings.
"""
from itertools import permutations, product
from collections import defaultdict

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


def induced_gamma_partition_candidate_1(k, k_2, chords, d_positions, u_positions):
    """Candidate 1: assign each U-triangle to the O^(2)-face whose
    D-triangle interval on T'_ann contains it.  Boundary U-triangles
    are assigned to the cyclically-next F^(2)."""
    n = k + k_2
    # First find face structure of O^(2)
    from tire_fiber_chords import compute_faces_from_chords
    o_faces = compute_faces_from_chords(k_2, chords)  # list of lists of B_in^(2)-edge indices
    # For each B_in^(2)-edge a, find which face it's in
    edge_to_face = {}
    for face_idx, edges in enumerate(o_faces):
        for a in edges:
            edge_to_face[a] = face_idx
    # d_positions[i] = T'_ann position of B_in^(2)-edge i.
    # Walk T'_ann cyclically.  At each U-position, find the next D-position
    # cyclically, and assign U to that D's face.
    d_pos_set = set(d_positions)
    d_pos_to_edge = {p: i for i, p in enumerate(d_positions)}

    u_to_face = {}
    for u_pos in u_positions:
        # Find the next D-position cyclically after u_pos
        for offset in range(1, n):
            check_pos = (u_pos + offset) % n
            if check_pos in d_pos_set:
                next_b_in_edge = d_pos_to_edge[check_pos]
                u_to_face[u_pos] = edge_to_face[next_b_in_edge]
                break

    # Each U-position corresponds to a γ-edge.  γ-edges are indexed by
    # u_positions in cyclic order.
    # u_positions[i] ↔ γ-edge i in the cyclic order on T'_ann
    u_pos_to_gamma_edge = {p: i for i, p in enumerate(u_positions)}

    # Build the partition: face_idx → list of γ-edge indices
    partition = defaultdict(list)
    for u_pos, face_idx in u_to_face.items():
        partition[face_idx].append(u_pos_to_gamma_edge[u_pos])
    return [sorted(v) for v in partition.values()]


def induced_gamma_partition_candidate_2(k, k_2, chords, d_positions, u_positions):
    """Candidate 2: assign each U-triangle to the O^(2)-face whose
    D-triangle interval on T'_ann contains it (using PREVIOUS D for
    boundary assignment instead of next)."""
    n = k + k_2
    from tire_fiber_chords import compute_faces_from_chords
    o_faces = compute_faces_from_chords(k_2, chords)
    edge_to_face = {}
    for face_idx, edges in enumerate(o_faces):
        for a in edges:
            edge_to_face[a] = face_idx
    d_pos_set = set(d_positions)
    d_pos_to_edge = {p: i for i, p in enumerate(d_positions)}

    u_to_face = {}
    for u_pos in u_positions:
        for offset in range(1, n):
            check_pos = (u_pos - offset) % n
            if check_pos in d_pos_set:
                prev_b_in_edge = d_pos_to_edge[check_pos]
                u_to_face[u_pos] = edge_to_face[prev_b_in_edge]
                break

    u_pos_to_gamma_edge = {p: i for i, p in enumerate(u_positions)}
    partition = defaultdict(list)
    for u_pos, face_idx in u_to_face.items():
        partition[face_idx].append(u_pos_to_gamma_edge[u_pos])
    return [sorted(v) for v in partition.values()]


def latin_set(partition, k):
    """Set of σ ∈ {1,2,3}^k where σ restricted to each partition block
    is a permutation of {1,2,3} (= each block has exactly 3 elements
    using all 3 colors)."""
    L = set()
    if any(len(b) != 3 for b in partition):
        return None  # not all triples
    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):
    """Yield all ways to partition a set of 3n elements into triples."""
    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 check_all_triple_partitions(pi_U, k):
    """Return all triple-partitions of {0,...,k-1} whose Latin set ⊆ π_U."""
    if k % 3 != 0:
        return []
    good = []
    for partition in all_partitions_into_triples(range(k)):
        L = latin_set([list(b) for b in partition], k)
        if L is None:
            continue
        if L <= pi_U:
            good.append([sorted(b) for b in partition])
    return good


def main():
    cases = [
        (6, 6, [(0, 3)], "k=k_2=6, antipodal chord (faces 3+3)"),
        (6, 6, [(0, 2), (3, 5)], "k=k_2=6, two chords (faces 2+2+2)"),
        (6, 3, [], "k=6, k_2=3 (T_2 inner = C_3, no chord, SP)"),
        (9, 9, [(0, 3), (3, 6)], "k=k_2=9, two chords (faces 3+3+3)"),
    ]
    for k, k_2, chords, desc in cases:
        print(f'\n=== {desc} ===')
        print(f'  k = γ length = {k}, k_2 = {k_2}, chords = {chords}')

        # Get T_2's projections.  T_2's outer boundary = γ, so we use π_U.
        from tire_fiber_chords import d_positions_for
        d_pos = d_positions_for(k, k_2)
        u_pos = u_positions_for(k, k_2)
        try:
            if k + k_2 <= 14:
                fibers, _, _ = fiber_distribution(k, k_2, chords)
                pi_U = projection_support(fibers, u_pos)
            else:
                # chunked streaming for larger cases
                pi_U = projection_support_streaming(k, k_2, chords, u_pos)
        except Exception as e:
            print(f'  computation failed: {e}')
            continue
        print(f'  |π_U| = {len(pi_U)}, 3^k = {3**k}')

        # Try candidate induced partitions
        for cand_name, cand_fn in [
            ('Candidate 1 (next-D)', induced_gamma_partition_candidate_1),
            ('Candidate 2 (prev-D)', induced_gamma_partition_candidate_2),
        ]:
            partition = cand_fn(k, k_2, chords, d_pos, u_pos)
            partition_str = ", ".join(str(sorted(b)) for b in partition)
            print(f'  {cand_name}: partition = {partition_str}')

            L = latin_set(partition, k)
            if L is None:
                print(f'    partition not all-3-triples; skipping Latin test')
                continue
            print(f'    |L| = {len(L)}, L ⊆ π_U? {L <= pi_U}')

        # Enumerate all triple-partitions and find which give Latin ⊆ π_U
        if k % 3 == 0:
            good_partitions = check_all_triple_partitions(pi_U, k)
            print(f'  # triple-partitions with Latin ⊆ π_U: {len(good_partitions)}')
            for p in good_partitions[:5]:
                print(f'    Latin ⊆ π_U: {p}')


if __name__ == '__main__':
    main()
