"""Numpy-optimized version of tire_fiber_chords.fiber_distribution.

The slow version iterates over all 3^n candidates in {1,2,3}^n; this
version directly constructs the 2^n + 2(-1)^n proper edge 3-colorings
of C_n via a vectorized binary-branch construction, then computes σ
and applies chord constraints in numpy.

Speedup: ~3^(n-1)/2^(n-1) = (3/2)^(n-1) ≈ 1500x at n=18, ~4000x at n=20.
"""
from __future__ import annotations

from collections import Counter
import numpy as np

from tire_fiber_chords import (
    d_positions_for,
    compute_faces_from_chords,
)


# OTHER[c, b] = the b-th color in {1,2,3} \ {c}, b ∈ {0,1}, c ∈ {1,2,3}
# Row 0 is a placeholder (color 0 unused).
OTHER_NP = np.array([
    [0, 0],
    [2, 3],
    [1, 3],
    [1, 2],
], dtype=np.int8)


def proper_cycle_colorings(n: int) -> np.ndarray:
    """Return an (N, n) int8 array of all proper edge 3-colorings of C_n
    where N = 2^n + 2(-1)^n.

    Encoding: a path c_0, c_1, ..., c_{n-1} is built by picking c_0 ∈
    {1,2,3} (3 choices) and then at each step picking one of the two
    colors ≠ previous (2 choices each).  We then drop those paths that
    fail the cyclic constraint c_{n-1} ≠ c_0.
    """
    if n < 2:
        return np.empty((0, n), dtype=np.int8)

    block_size = 1 << (n - 1)
    n_raw = 3 * block_size
    arr = np.empty((n_raw, n), dtype=np.int8)

    # Column 0: blocks of size block_size for c_0 = 1, 2, 3.
    arr[0 * block_size : 1 * block_size, 0] = 1
    arr[1 * block_size : 2 * block_size, 0] = 2
    arr[2 * block_size : 3 * block_size, 0] = 3

    row_idx = np.arange(block_size, dtype=np.int64)
    for i in range(1, n):
        bits_block = ((row_idx >> (i - 1)) & 1).astype(np.int8)
        bits_full = np.tile(bits_block, 3)
        prev = arr[:, i - 1]
        arr[:, i] = OTHER_NP[prev, bits_full]

    mask = arr[:, -1] != arr[:, 0]
    return arr[mask]


def induced_sigma_vec(c: np.ndarray) -> np.ndarray:
    """Given an (N, n) proper-cycle-coloring array, return the (N, n)
    σ array where σ_i = 6 - c_{i-1} - c_i (the third color at vertex i)."""
    c_prev = np.roll(c, 1, axis=1)
    return (6 - c_prev - c).astype(np.int8)


def fiber_distribution_fast(m: int, k: int, chords) -> tuple[dict, list, list]:
    """Drop-in replacement for tire_fiber_chords.fiber_distribution."""
    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
    ]

    c = proper_cycle_colorings(n)
    if c.size == 0:
        return {}, d_positions, d_positions_by_face

    sigma = induced_sigma_vec(c)

    # Chord constraint: σ values at each face's positions are pairwise distinct.
    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]])

    sigma_ok = sigma[mask]

    # Count multiplicities.  np.unique with return_counts is faster than Counter for big arrays.
    if sigma_ok.shape[0] == 0:
        return {}, d_positions, d_positions_by_face
    uniq, counts = np.unique(sigma_ok, axis=0, return_counts=True)
    fibers = {tuple(row): int(c) for row, c in zip(uniq.tolist(), counts.tolist())}
    return fibers, d_positions, d_positions_by_face


def spoke_only_fiber_distribution_fast(n: int) -> dict:
    """Steiner-rich baseline (no chord constraints)."""
    if n < 2:
        return {}
    c = proper_cycle_colorings(n)
    sigma = induced_sigma_vec(c)
    uniq, counts = np.unique(sigma, axis=0, return_counts=True)
    return {tuple(row): int(c) for row, c in zip(uniq.tolist(), counts.tolist())}


def projection_support_fast(fibers: dict, positions: list[int]) -> set:
    return {tuple(sigma[p] for p in positions) for sigma in fibers}


if __name__ == '__main__':
    import time

    print("Benchmark vs. slow version")
    print("-" * 60)
    from tire_fiber_chords import fiber_distribution as fiber_slow

    cases = [
        (4, 4, [(0, 2)]),
        (6, 6, [(0, 3)]),
        (6, 9, [(0, 2), (3, 5), (6, 8)]),
        (9, 9, [(0, 2), (3, 5), (6, 8)]),
        (12, 12, [(0, 3), (4, 7), (8, 11)]),
    ]
    for m, k, ch in cases:
        n = m + k

        t0 = time.time()
        fast, _, _ = fiber_distribution_fast(m, k, ch)
        dt_fast = time.time() - t0

        # Sanity-check vs slow for small n
        if n <= 15:
            t0 = time.time()
            slow, _, _ = fiber_slow(m, k, ch)
            dt_slow = time.time() - t0
            assert fast == slow, f"mismatch at (m={m}, k={k}): fast {len(fast)} != slow {len(slow)}"
            print(f"(m={m}, k={k}, n={n}): fast {dt_fast:.3f}s | slow {dt_slow:.3f}s  | "
                  f"speedup {dt_slow/max(dt_fast, 1e-6):.1f}×  | |C| = {len(fast)}")
        else:
            print(f"(m={m}, k={k}, n={n}): fast {dt_fast:.3f}s  | |C| = {len(fast)}")
