"""Streaming/chunked variant of tire_fiber_chords_fast.py for n > 27.

The in-memory version materializes an (N, n) int8 array with N = 3·2^(n-1),
which at n = 30 is ~30 GB.  This chunked version processes paths in
chunks of ~2^20 rows, never holding the full array in memory.

For projection support: we only need the set of σ-projected rows,
which is typically much smaller than 3^k.  We accumulate into a Python
set.

For γ up to ~18 with m ≥ γ this brings n up to ~36, with one tire
taking minutes (versus impossible).
"""
from __future__ import annotations

import numpy as np

from tire_fiber_chords import (
    d_positions_for,
    u_positions_for,
    compute_faces_from_chords,
)
from tire_fiber_chords_fast import OTHER_NP, induced_sigma_vec


def proper_cycle_colorings_chunked(n: int, chunk_size: int = 1 << 20):
    """Yield (rows, n) int8 chunks of proper edge 3-colorings of C_n.

    Each chunk is a numpy array; chunks together cover every proper
    edge 3-coloring exactly once.  After filtering for the cycle closure
    constraint, a chunk may be smaller than chunk_size."""
    if n < 2:
        return

    block_size = 1 << (n - 1)
    eff_chunk = min(chunk_size, block_size)

    for c0 in (1, 2, 3):
        for start in range(0, block_size, eff_chunk):
            end = min(start + eff_chunk, block_size)
            n_rows = end - start
            arr = np.empty((n_rows, n), dtype=np.int8)
            arr[:, 0] = c0

            row_idx = np.arange(start, end, dtype=np.int64)
            for i in range(1, n):
                bits = ((row_idx >> (i - 1)) & 1).astype(np.int8)
                prev = arr[:, i - 1]
                arr[:, i] = OTHER_NP[prev, bits]

            mask = arr[:, -1] != arr[:, 0]
            if mask.any():
                yield arr[mask]


def projection_support_streaming(
    m: int,
    k: int,
    chords,
    positions: list[int],
    chunk_size: int = 1 << 20,
) -> set:
    """Compute the projected support {σ[positions] : σ realisable}
    without materializing the full σ array.

    Returns: set of tuples in {1,2,3}^len(positions).
    """
    n = m + k
    d_positions = d_positions_for(m, k)
    o_faces = compute_faces_from_chords(k, chords)
    d_positions_by_face = [
        [d_positions[a] for a in face_edges]
        for face_edges in o_faces
    ]

    support: set[tuple[int, ...]] = set()
    for c_chunk in proper_cycle_colorings_chunked(n, chunk_size):
        sigma = induced_sigma_vec(c_chunk)

        mask = np.ones(sigma.shape[0], dtype=bool)
        for face in d_positions_by_face:
            if len(face) <= 1:
                continue
            for i in range(len(face)):
                for j in range(i + 1, len(face)):
                    mask &= (sigma[:, face[i]] != sigma[:, face[j]])

        if not mask.any():
            continue
        sigma_ok = sigma[mask]
        projected = sigma_ok[:, positions]
        support.update(map(tuple, projected.tolist()))

    return support


def spoke_only_projection_support_streaming(
    n: int, positions: list[int], chunk_size: int = 1 << 20
) -> set:
    """Steiner-rich baseline (no chord constraints)."""
    support: set[tuple[int, ...]] = set()
    for c_chunk in proper_cycle_colorings_chunked(n, chunk_size):
        sigma = induced_sigma_vec(c_chunk)
        projected = sigma[:, positions]
        support.update(map(tuple, projected.tolist()))
    return support


if __name__ == '__main__':
    import time
    from tire_fiber_chords_fast import fiber_distribution_fast, projection_support_fast

    print("Validation: chunked vs in-memory")
    print("-" * 60)
    cases = [
        (4, 4, [(0, 2)]),
        (6, 6, [(0, 3)]),
        (6, 9, [(0, 2), (3, 5), (6, 8)]),
        (9, 9, [(0, 3), (4, 7)]),
        (12, 12, [(0, 3), (4, 7), (8, 11)]),
    ]
    for m, k, ch in cases:
        n = m + k
        d_pos = d_positions_for(m, k)

        t0 = time.time()
        fibers_mem, _, _ = fiber_distribution_fast(m, k, ch)
        S_mem = projection_support_fast(fibers_mem, d_pos)
        dt_mem = time.time() - t0

        t0 = time.time()
        S_chunked = projection_support_streaming(m, k, ch, d_pos)
        dt_chunked = time.time() - t0

        ok = S_mem == S_chunked
        status = "OK" if ok else "MISMATCH"
        print(f"(m={m}, k={k}, n={n}): mem {dt_mem:.2f}s, chunked {dt_chunked:.2f}s, "
              f"|S|={len(S_chunked)}, {status}")

    print()
    print("Stress test (n > 24):")
    print("-" * 60)
    cases_big = [
        (12, 13, [(0, 3), (4, 7), (8, 11), (12, 1)]),    # n=25
        (15, 15, [(0, 3), (4, 7), (8, 11), (12, 14), (12, 0)]),  # n=30
    ]
    for m, k, ch in cases_big:
        n = m + k
        d_pos = d_positions_for(m, k)
        t0 = time.time()
        S = projection_support_streaming(m, k, ch, d_pos)
        dt = time.time() - t0
        print(f"(m={m}, k={k}, n={n}): {dt:.1f}s, |S|={len(S)}")
