"""Chain pigeonhole under SR + PDS, no chord constraints.

Under the correct PDS modelling, each tire's T'_{f'} has γ inner-spoke
pendants regardless of O's chord structure (chord edges of O become
G'-edges between inner-spoke vertices, but those edges are NOT in
T'_{f'} since neither endpoint is in V(f')).  So the only inputs are
cycle lengths (m, k) per tire.

For each tire T with B_out length m and B_in length k:
  Π_T ⊆ {1,2,3}^m × {1,2,3}^k
  = { (σ_U, σ_D) : σ comes from a proper edge 3-colouring of C_{m+k} }

Step 2 (pairwise): for adjacent tires T_1, T_2 sharing γ, do
π_U(T_1)|_γ and π_D(T_2)|_γ overlap?  By step-1 saturation, yes —
T_1's γ-side saturates {1,2,3}^γ when m_1 ≥ |γ|, etc.

Step 3 (chain consistency): chain T_1 | T_2 | ... | T_n.  σ at each
shared cycle must be jointly consistent.  Propagate forward via the
joint supports and check if state ever empties.
"""
from __future__ import annotations

import numpy as np
from itertools import product

from tire_fiber_chords import d_positions_for, u_positions_for
from tire_fiber_chords_fast import proper_cycle_colorings, induced_sigma_vec


_pi_cache: dict = {}


def joint_support_SR(m: int, k: int) -> set[tuple[tuple[int, ...], tuple[int, ...]]]:
    """Π_T = {(σ_U, σ_D)} pairs for SR tire (no chord)."""
    key = (m, k)
    if key in _pi_cache:
        return _pi_cache[key]
    n = m + k
    d_pos = d_positions_for(m, k)
    u_pos = u_positions_for(m, k)
    c = proper_cycle_colorings(n)
    sigma = induced_sigma_vec(c)
    sigma_u = sigma[:, u_pos]
    sigma_d = sigma[:, d_pos]
    pairs = set(zip(map(tuple, sigma_u.tolist()),
                    map(tuple, sigma_d.tolist())))
    _pi_cache[key] = pairs
    return pairs


# --- Step 2 ---

def step2_SR(gamma: int, m_1: int, k_2: int) -> dict:
    """T_1 has (m=m_1, k=γ).  T_2 has (m=γ, k=k_2)."""
    pi1 = joint_support_SR(m_1, gamma)
    pi2 = joint_support_SR(gamma, k_2)
    sd1 = {p[1] for p in pi1}   # γ-projection from T_1's side
    su2 = {p[0] for p in pi2}   # γ-projection from T_2's side
    fwd = sd1 & su2
    rev = sd1 & {s[::-1] for s in su2}
    return {
        'γ': gamma, 'm_1': m_1, 'k_2': k_2,
        '|sd1|': len(sd1), '|su2|': len(su2), '3^γ': 3**gamma,
        'sd1_saturates': len(sd1) == 3**gamma,
        'su2_saturates': len(su2) == 3**gamma,
        'fwd': len(fwd), 'rev': len(rev),
        'compatible': bool(fwd or rev),
    }


# --- Step 3 ---

def step3_chain(chain: list[tuple[int, int]]) -> dict:
    """Forward-propagate joint supports along a tire chain.

    chain[i] = (m_i, k_i) for tire T_i.  Adjacency requires k_{i+1} = m_i
    (T_{i+1}'s B_in = T_i's B_out = shared γ).

    State at step i = set of σ at the "current" shared cycle γ_i = L_i,
    which is reachable through T_1, ..., T_i.

    Returns: compatibility result and the state size trajectory.
    """
    # Validate adjacency
    for i in range(len(chain) - 1):
        if chain[i][0] != chain[i + 1][1]:
            return {'error': f'chain adjacency mismatch at i={i}: '
                             f'T_{i+1}.m={chain[i][0]} != T_{i+2}.k={chain[i+1][1]}'}

    pis = [joint_support_SR(m, k) for (m, k) in chain]

    # Initial state: σ at L_1 from T_1's σ_U-projection (any σ_D ok since
    # the innermost boundary L_0 has no further constraint).
    state = {p[0] for p in pis[0]}
    trajectory = [len(state)]

    for i in range(1, len(chain)):
        new_state = set()
        pi = pis[i]
        # T_i has B_in = L_i, B_out = L_{i+1}.  pair = (σ at L_{i+1}, σ at L_i).
        # For each pair, if σ at L_i in current state, add σ at L_{i+1} to new state.
        for sigma_outer, sigma_inner in pi:
            if sigma_inner in state:
                new_state.add(sigma_outer)
        trajectory.append(len(new_state))
        if not new_state:
            return {
                'chain': chain,
                'compatible': False,
                'failed_at_tire_index': i,
                'trajectory': trajectory,
            }
        state = new_state

    return {
        'chain': chain,
        'compatible': True,
        'final_state_size': len(state),
        'trajectory': trajectory,
    }


def step3_chain_with_reflections(chain: list[tuple[int, int]]) -> dict:
    """Like step3_chain but also tries flipping orientation at each
    junction (since adjacent tires may have reversed γ-orientations
    in the actual plane embedding).

    More expensive: 2^(len(chain)-1) orientation choices."""
    n = len(chain)
    pis = [joint_support_SR(m, k) for (m, k) in chain]

    # Try every combination of orientation flips at junctions.
    # flip[i] = True means flip σ at L_{i+1} when bridging T_i → T_{i+1}.
    # Equivalent: reverse the σ_D of T_{i+1} before matching.

    best = None
    for flips in product([False, True], repeat=n - 1):
        state = {p[0] for p in pis[0]}
        traj = [len(state)]
        ok = True
        for i in range(1, n):
            pi = pis[i]
            new_state = set()
            if flips[i - 1]:
                # Reverse σ at L_i (the boundary between T_{i-1} and T_i)
                state_match = state
                for sigma_outer, sigma_inner in pi:
                    if sigma_inner[::-1] in state_match:
                        new_state.add(sigma_outer)
            else:
                for sigma_outer, sigma_inner in pi:
                    if sigma_inner in state:
                        new_state.add(sigma_outer)
            traj.append(len(new_state))
            if not new_state:
                ok = False
                break
            state = new_state
        if ok:
            if best is None or len(state) > best['final_state_size']:
                best = {
                    'chain': chain,
                    'compatible': True,
                    'final_state_size': len(state),
                    'trajectory': traj,
                    'flips': flips,
                }

    if best is not None:
        return best
    return {
        'chain': chain,
        'compatible': False,
        'tried_orientations': 2 ** (n - 1),
    }


# --- Test scenarios ---

PAIRWISE_CASES = [
    # (γ, m_1, k_2): T_1 has B_in=γ; T_2 has B_out=γ.
    (3, 3, 3),
    (4, 4, 3),
    (4, 4, 4),
    (5, 5, 3),
    (5, 6, 5),
    (6, 6, 3),
    (6, 6, 4),
    (6, 6, 5),
    (6, 6, 6),
    (8, 8, 4),
    (8, 8, 6),
    (8, 10, 8),
    (10, 10, 8),
    (12, 12, 6),
]

# 3-tire and longer chains.
# Each tuple (m, k) -- adjacent (m_i, k_i), (m_{i+1}, k_{i+1}) need k_{i+1} = m_i.
CHAIN_CASES = [
    # Strictly outward-growing PDS:
    [(4, 3), (5, 4), (6, 5)],
    [(4, 3), (6, 4), (8, 6)],
    [(5, 4), (6, 5), (7, 6)],
    [(6, 3), (8, 6), (10, 8)],
    # Stalled growth:
    [(4, 3), (5, 4), (5, 5)],
    [(4, 3), (4, 4), (5, 4)],
    # Shrinking (anti-PDS, for stress):
    [(3, 4), (3, 3)],   # would need k_2=m_1=3 -- ok
    # Longer chains:
    [(4, 3), (5, 4), (6, 5), (7, 6)],
    [(4, 3), (5, 4), (6, 5), (7, 6), (8, 7)],
    [(4, 3), (5, 4), (6, 5), (7, 6), (8, 7), (9, 8)],
]


def main():
    print("=" * 80)
    print("Step 2 under SR + PDS (no chord)")
    print("=" * 80)
    print(f"{'γ':>3} {'m_1':>4} {'k_2':>4} {'|sd1|':>6} {'|su2|':>6} {'3^γ':>7} "
          f"{'sat1':>5} {'sat2':>5} {'fwd':>5} {'rev':>5} compat")
    for γ, m_1, k_2 in PAIRWISE_CASES:
        r = step2_SR(γ, m_1, k_2)
        sat1 = "✓" if r['sd1_saturates'] else "·"
        sat2 = "✓" if r['su2_saturates'] else "·"
        ok = "YES" if r['compatible'] else "NO"
        print(f"{γ:>3} {m_1:>4} {k_2:>4} {r['|sd1|']:>6} {r['|su2|']:>6} {r['3^γ']:>7} "
              f"{sat1:>5} {sat2:>5} {r['fwd']:>5} {r['rev']:>5} {ok}")

    print()
    print("=" * 80)
    print("Step 3: chain consistency (3+ tires) under SR + PDS")
    print("=" * 80)
    for chain in CHAIN_CASES:
        adj_ok = all(chain[i][0] == chain[i + 1][1] for i in range(len(chain) - 1))
        adj = "OK" if adj_ok else "(adjacency mismatch)"
        if not adj_ok:
            print(f"  SKIP {chain}: {adj}")
            continue
        # Use the orientation-checking version
        r = step3_chain_with_reflections(chain)
        chain_str = " | ".join(f"({m},{k})" for (m, k) in chain)
        if r.get('compatible'):
            print(f"  {chain_str}: COMPATIBLE, final state size {r['final_state_size']}, "
                  f"trajectory {r['trajectory']}")
        else:
            failed = r.get('failed_at_tire_index', '?')
            print(f"  {chain_str}: **INCOMPATIBLE** (failed at tire index {failed})")


if __name__ == '__main__':
    main()
