"""Decompose S_1 ∩ S_2 into S_3-orbits for each (T_1, T_2) pair in step-2.

Question: for structurally different (T_1, T_2) sharing γ, does the same
canonical orbit show up across pairs, or is each intersection a different
orbit?  Are intersections always unions of complete S_3-orbits?
"""
from itertools import permutations

from tire_fiber_step2 import (
    CASES,
    project_T1_D,
    project_T2_U,
    intersect_with_reflection,
    fmt_cfg,
)


def s3_orbit(sigma):
    """Return the S_3 orbit of sigma under color permutations.

    Acts on each entry of sigma by pi : {1,2,3} -> {1,2,3}, generating
    up to 6 elements (fewer if sigma has fewer than 3 distinct colors).
    """
    orbit = set()
    for pi in permutations([1, 2, 3]):
        # pi is a permutation of (1,2,3); color c maps to pi[c-1]
        mapped = tuple(pi[c - 1] for c in sigma)
        orbit.add(mapped)
    return orbit


def cyclic_rotations(sigma):
    """Cyclic rotations of sigma, as a set."""
    n = len(sigma)
    return {tuple(sigma[(i + k) % n] for i in range(n)) for k in range(n)}


def canonical_orbit_rep(sigma):
    """Return a canonical representative of sigma's combined orbit under
    S_3 color action AND cyclic rotation of positions: the lexicographically
    smallest element of (S_3 ∪ cyclic) closure."""
    closure = set()
    for pi in permutations([1, 2, 3]):
        for k in range(len(sigma)):
            rot = tuple(pi[sigma[(i + k) % len(sigma)] - 1] for i in range(len(sigma)))
            closure.add(rot)
    return min(closure)


def decompose_into_s3_orbits(S):
    """Decompose set S into S_3-orbits (color permutation, NOT cyclic).
    Returns dict {canonical_rep -> list of all elements in that orbit}."""
    orbits = {}
    seen = set()
    for sigma in sorted(S):
        if sigma in seen:
            continue
        orb = s3_orbit(sigma)
        orb_in_S = orb & S
        rep = min(orb_in_S)
        orbits[rep] = sorted(orb_in_S)
        seen |= orb_in_S
    return orbits


def decompose_into_combined_orbits(S):
    """Decompose into orbits under S_3 color action AND cyclic position
    rotation.  Returns dict {canonical_rep -> elements}."""
    orbits = {}
    seen = set()
    for sigma in sorted(S):
        if sigma in seen:
            continue
        closure = set()
        for pi in permutations([1, 2, 3]):
            for k in range(len(sigma)):
                rot = tuple(pi[sigma[(i + k) % len(sigma)] - 1]
                            for i in range(len(sigma)))
                closure.add(rot)
        in_S = closure & S
        rep = min(in_S)
        orbits[rep] = sorted(in_S)
        seen |= in_S
    return orbits


def is_complete_s3_closed(S):
    """Check if S is closed under S_3 color permutations."""
    for sigma in S:
        for tau in s3_orbit(sigma):
            if tau not in S:
                return False
    return True


def main():
    rep_to_pairs = {}      # canonical rep -> list of (T1, T2) pairs in which it appears
    pair_to_orbits = []    # list of (T1, T2, orbits_dict)

    print(f"{'γ':>2}  {'T1':<32s}  {'T2':<32s}  "
          f"{'|S1∩S2|':>7s}  {'S3-closed?':>10s}  "
          f"{'#S3-orbits':>10s}  {'orbit sizes':>15s}")
    print("-" * 120)
    for gamma, t1, t2 in CASES:
        m_1, ch1, model1 = t1
        k_2, ch2, model2 = t2
        S1 = project_T1_D(gamma, m_1, ch1, model1)
        S2 = project_T2_U(gamma, k_2, ch2, model2)
        forward, reverse = intersect_with_reflection(S1, S2)
        # use the larger of forward, reverse for analysis (typically same)
        S = forward if len(forward) >= len(reverse) else reverse

        if not S:
            continue
        closed = is_complete_s3_closed(S)
        orbits = decompose_into_s3_orbits(S)
        orbit_sizes = sorted(len(o) for o in orbits.values())

        # Combined orbits (S_3 + cyclic) -- to identify canonical patterns
        combined = decompose_into_combined_orbits(S)
        for rep in combined:
            rep_to_pairs.setdefault(rep, []).append(
                (gamma, fmt_cfg(m_1, ch1, model1), fmt_cfg(k_2, ch2, model2))
            )

        pair_to_orbits.append((gamma, t1, t2, orbits, combined))

        print(
            f"{gamma:>2}  {fmt_cfg(m_1, ch1, model1):<32s}  "
            f"{fmt_cfg(k_2, ch2, model2):<32s}  "
            f"{len(S):>7d}  {('yes' if closed else 'NO'):>10s}  "
            f"{len(orbits):>10d}  {str(orbit_sizes):>15s}"
        )

    print("-" * 120)
    print()
    print("=== Combined (S_3 + cyclic) orbit canonical representatives ===")
    print("    Showing canonical rep, orbit size, and which (T_1, T_2) pairs contain it.")
    print()
    for rep, pairs in sorted(rep_to_pairs.items(),
                             key=lambda kv: (-len(kv[1]), kv[0])):
        # Compute full combined orbit size
        closure = set()
        for pi in permutations([1, 2, 3]):
            for k in range(len(rep)):
                rot = tuple(pi[rep[(i + k) % len(rep)] - 1]
                            for i in range(len(rep)))
                closure.add(rot)
        print(f"  rep {rep}  |orbit| = {len(closure)}  appears in {len(pairs)} pair(s):")
        for g, t1s, t2s in pairs:
            print(f"      γ={g}  T_1={t1s}  T_2={t2s}")

    # Highlight the rainbow rep
    print()
    print("=== Rainbow pattern check ===")
    rainbow_reps = []
    for rep in rep_to_pairs:
        if len(rep) == 6:
            # Check if pattern is (a,b,c,b,c,a)
            a, b, c = rep[0], rep[1], rep[2]
            if (rep == (a, b, c, b, c, a)
                    and len({a, b, c}) == 3):
                rainbow_reps.append(rep)
    print(f"Found {len(rainbow_reps)} canonical reps matching (a,b,c,b,c,a) pattern:")
    for rep in rainbow_reps:
        n = len(rep_to_pairs[rep])
        print(f"  {rep}  appears in {n} pair(s)")


if __name__ == '__main__':
    main()
