"""Step 1 of birkhoff_heesch_reducibility.tex: enumerate the fiber
distribution N(T'_{f'}; σ) for spoke-only tire annular face connectors.

In the spoke-only setting (Remark 1.16.1) the tire annular face
connector is

    G_n  =  C_n  +  one pendant edge at each cycle vertex

where n = m + k = |V(B_out)| + |V(B_in)|.  Each cycle vertex v_i has
degree 3 (two cycle edges + one pendant spoke).  At k = 3 colors, the
proper-edge-coloring constraint forces the spoke color s_i at v_i to
be the unique color not used by the two cycle edges at v_i.

So the map

    proper edge 3-coloring of G_n   -->   σ ∈ {1,2,3}^n

is well-defined (each cycle coloring maps to one σ).  The fiber count
N(σ) := #{cycle colorings inducing σ} partitions the P_e(G_n, 3) =
2^n + 2(-1)^n proper edge colorings into fibers over σ.

This script computes the full fiber distribution and various
summaries.  Output is written to fiber_data.json next to this file.
"""
from __future__ import annotations

import json
import os
from collections import Counter
from itertools import product


COLORS = (1, 2, 3)


def induced_sigma(cycle_coloring: tuple[int, ...]) -> tuple[int, ...]:
    """Given a proper edge 3-coloring (c(e_0), ..., c(e_{n-1})) of C_n,
    return the induced spoke configuration σ where σ_i = third color at v_i.
    v_i is incident to edges e_{i-1} and e_i (indices mod n)."""
    n = len(cycle_coloring)
    return tuple(
        ({1, 2, 3} - {cycle_coloring[(i - 1) % n], cycle_coloring[i]}).pop()
        for i in range(n)
    )


def enumerate_proper_cycle_colorings(n: int):
    """Yield all proper edge 3-colorings of C_n (i.e., n-tuples in {1,2,3}^n
    with consecutive entries distinct, including wraparound)."""
    for ce in product(COLORS, repeat=n):
        if all(ce[i] != ce[(i + 1) % n] for i in range(n)):
            yield ce


def fiber_distribution(n: int) -> dict[tuple[int, ...], int]:
    """Return Counter mapping σ -> N(σ) for G_n = C_n + n pendants."""
    fibers: Counter[tuple[int, ...]] = Counter()
    for ce in enumerate_proper_cycle_colorings(n):
        fibers[induced_sigma(ce)] += 1
    return dict(fibers)


def projection_support(fibers: dict[tuple[int, ...], int], positions: tuple[int, ...]) -> set[tuple[int, ...]]:
    """Return the set of σ-projections onto the given positions
    (e.g., positions=(0,1,2,3) extracts the first four coords of each σ)."""
    return {tuple(sigma[p] for p in positions) for sigma in fibers}


def summary_row(n: int) -> dict:
    """Compute summary statistics for n = m + k."""
    fibers = fiber_distribution(n)
    total = sum(fibers.values())
    support = len(fibers)
    fiber_sizes = list(fibers.values())
    size_dist = Counter(fiber_sizes)

    # P_e(C_n, 3) closed form
    pe_closed = 2**n + 2 * ((-1) ** n)

    # Universe size for σ
    universe = 3**n

    # A-reducibility check: is every σ realizable?
    a_reducible = (support == universe)

    return {
        'n': n,
        'P_e_computed': total,
        'P_e_closed_form': pe_closed,
        'match': total == pe_closed,
        'sigma_support': support,
        'sigma_universe': universe,
        'support_fraction': support / universe if universe else 0.0,
        'A_reducible': a_reducible,
        'max_fiber': max(fiber_sizes) if fiber_sizes else 0,
        'min_fiber': min(fiber_sizes) if fiber_sizes else 0,
        'mean_fiber': total / support if support else 0.0,
        'fiber_size_distribution': dict(size_dist),
    }


def half_cycle_projections(n: int, k_values=(3, 4, 5, 6)) -> dict:
    """For each k <= n in k_values, project σ onto contiguous-block
    positions (0, 1, ..., k-1) and tabulate how many of {1,2,3}^k are
    covered.  This is the support of σ restricted to a 'shared cycle'
    of length k.

    Also try the 'alternating' pattern (take every other position)
    when the math works out."""
    fibers = fiber_distribution(n)
    out = {}
    for k in k_values:
        if k > n:
            continue
        # contiguous block: first k positions
        contiguous = projection_support(fibers, tuple(range(k)))
        # 'spread' block: positions evenly spaced (if possible)
        if k <= n:
            spread_positions = tuple(round(i * n / k) % n for i in range(k))
            # ensure distinct
            if len(set(spread_positions)) == k:
                spread = projection_support(fibers, spread_positions)
            else:
                spread = None
        else:
            spread = None
        out[k] = {
            'k': k,
            'positions_contiguous': tuple(range(k)),
            'support_contiguous': len(contiguous),
            'positions_spread': spread_positions if spread else None,
            'support_spread': len(spread) if spread else None,
            'universe_k': 3**k,
        }
    return out


def format_table(rows: list[dict]) -> str:
    lines = [
        'n   |  P_e  |  |C|  |  3^n   | |C|/3^n | A-red? | max | min |  mean | fiber sizes',
        '----+-------+-------+--------+---------+--------+-----+-----+-------+-------------',
    ]
    for r in rows:
        sz_dist = sorted(r['fiber_size_distribution'].items())
        sz_str = ', '.join(f"{s}×{c}" for s, c in sz_dist)
        lines.append(
            f"{r['n']:2d}  | {r['P_e_computed']:5d} | {r['sigma_support']:5d} | "
            f"{r['sigma_universe']:6d} | {r['support_fraction']:7.3f} | "
            f"{'yes' if r['A_reducible'] else 'no':4s}   | "
            f"{r['max_fiber']:3d} | {r['min_fiber']:3d} | "
            f"{r['mean_fiber']:5.2f} | {sz_str}"
        )
    return '\n'.join(lines)


def format_projection_table(projections_by_n: dict[int, dict]) -> str:
    lines = [
        'n   |  k  |  contiguous-block |/ 3^k  | spread |/ 3^k',
        '----+-----+-------------------+-------+--------+-------',
    ]
    for n in sorted(projections_by_n):
        for k, info in sorted(projections_by_n[n].items()):
            uni = info['universe_k']
            sc = info['support_contiguous']
            ss = info['support_spread']
            ss_str = f"{ss:6d} | {ss/uni:.3f}" if ss is not None else f"{'-':6s} |  -   "
            lines.append(
                f"{n:2d}  | {k:2d}  | {sc:7d} / {uni:5d} | {sc/uni:.3f} | {ss_str}"
            )
    return '\n'.join(lines)


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    out_json = os.path.join(here, 'fiber_data.json')
    out_txt = os.path.join(here, 'fiber_data.txt')

    # Step 1: fiber distribution for n = 3..12
    n_range = range(3, 13)
    rows = []
    full_fibers = {}
    projections = {}
    for n in n_range:
        rows.append(summary_row(n))
        full_fibers[n] = fiber_distribution(n)
        projections[n] = half_cycle_projections(n)

    # Output tables
    summary_text = (
        '=== FIBER DISTRIBUTION SUMMARY (G_n = C_n + n pendants, k=3 colors) ===\n\n'
        + format_table(rows)
        + '\n\n=== PROJECTION ONTO SHARED-CYCLE BLOCKS ===\n\n'
        + format_projection_table(projections)
        + '\n'
    )
    print(summary_text)

    with open(out_txt, 'w') as f:
        f.write(summary_text)
    print(f"\nWrote {out_txt}")

    # Save raw data for follow-up analyses (step 2: adjacent-tire overlap)
    data = {
        'summary_rows': rows,
        'projections': {str(n): projections[n] for n in projections},
        # full_fibers: convert tuple keys to strings for JSON
        'full_fibers': {
            str(n): {','.join(map(str, sigma)): count
                     for sigma, count in full_fibers[n].items()}
            for n in full_fibers
        },
    }
    with open(out_json, 'w') as f:
        json.dump(data, f, indent=2)
    print(f"Wrote {out_json}")


if __name__ == '__main__':
    main()
