"""6-hour counterexample search for the Latin-substructure conjecture.

Logic: pick a chord configuration cs1 for T1 (with B_in = γ of length k)
and a chord configuration cs2 for T2 (with B_in of length k_2), with
various m_1 (T1's B_out length) and k_2.  For each (m_1, k, cs1) and
(k, k_2, cs2) pair (both SP-feasible), compute π_D(T1) and π_U(T2) on
γ; check if their intersection (forward or reverse) is empty.

Append-log: each test writes a JSON-Lines record to LOG_PATH.  On
restart, we read the log to skip already-tested pairs.

Run until time budget exhausted; report counterexamples loudly.
"""
from __future__ import annotations

import json
import os
import sys
import time
from itertools import product
from pathlib import Path

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 (
    fiber_distribution_fast,
    spoke_only_fiber_distribution_fast,
    projection_support_fast,
    proper_cycle_colorings,
)


HERE = Path(__file__).resolve().parent
LOG_PATH = HERE / "counterexample_search.log"


def log_event(event: dict) -> None:
    """Append one JSON line to the log."""
    event = dict(event)  # don't mutate caller
    event.setdefault('wall_time', time.time())
    with open(LOG_PATH, 'a') as f:
        f.write(json.dumps(event, default=str) + "\n")
        f.flush()


def read_existing_keys() -> set[str]:
    """Read the log and return set of pair-keys already tested."""
    keys = set()
    if not LOG_PATH.exists():
        return keys
    with open(LOG_PATH) as f:
        for line in f:
            try:
                ev = json.loads(line)
            except json.JSONDecodeError:
                continue
            if ev.get('type') == 'pair' and 'key' in ev:
                keys.add(ev['key'])
    return keys


def make_key(gamma: int, m_1: int, ch1, model1: str, k_2: int, ch2, model2: str) -> str:
    ch1_str = json.dumps(sorted([sorted(c) for c in ch1])) if ch1 else "[]"
    ch2_str = json.dumps(sorted([sorted(c) for c in ch2])) if ch2 else "[]"
    return f"gamma={gamma}|t1=({m_1},{ch1_str},{model1})|t2=({k_2},{ch2_str},{model2})"


def is_sp_feasible(k: int, chords) -> bool:
    if not chords:
        return k <= 3
    faces = compute_faces_from_chords(k, chords)
    return max(len(f) for f in faces) <= 3


# Hand-curated chord configurations per k.  Each entry should be
# SP-feasible (every face has at most 3 B_in edges).
CHORD_CONFIGS: dict[int, list] = {
    3: [[]],

    # k=4: only one essentially-distinct chord modulo rotation.
    4: [[(0, 2)]],

    # k=5: short or long chord.
    5: [[(0, 2)], [(0, 3)]],

    # k=6: antipodal, or 2-chord matching.
    6: [
        [(0, 3)],
        [(0, 2), (3, 5)],
        [(1, 3), (4, 0)],
    ],

    # k=7: needs ≥ 2 chords for all-faces-≤-3.
    7: [
        [(0, 2), (3, 6)],
        [(0, 3), (4, 6)],
        [(0, 2), (3, 5)],     # outer face: edges 5,6 → size 2
        [(0, 3), (3, 6)],     # uses vertex 3 twice
    ],

    # k=8.
    8: [
        [(0, 2), (3, 5), (6, 0)],  # vertex 0 used twice
        [(0, 3), (4, 7)],
        [(0, 3), (3, 6)],
        [(0, 2), (3, 6)],
        [(0, 2), (4, 6)],
    ],

    # k=9: 2-chord tight (faces 3,3,3) and 3-chord variants.
    9: [
        [(0, 3), (4, 7)],
        [(0, 2), (3, 5), (6, 8)],
        [(0, 2), (4, 6), (3, 7)],
        [(0, 3), (3, 6), (6, 0)],
    ],

    # k=10.
    10: [
        [(0, 3), (4, 7), (7, 0)],
        [(0, 2), (3, 5), (6, 8), (9, 1)],
        [(0, 3), (4, 7), (8, 0)],
        [(0, 3), (3, 6), (6, 9), (9, 0)],
    ],

    # k=11.
    11: [
        [(0, 3), (3, 6), (6, 9), (9, 0)],
        [(0, 2), (3, 5), (6, 8), (9, 0)],
    ],

    # k=12.
    12: [
        [(0, 3), (4, 7), (8, 11)],            # tight 4-faces-of-3
        [(0, 3), (3, 6), (6, 9), (9, 0)],     # tight ring
        [(0, 3), (4, 11), (5, 7), (8, 10)],   # nested
        [(0, 2), (3, 5), (6, 9), (10, 0)],    # mixed
    ],

    # k=13.
    13: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (0, 12)],
        [(0, 3), (4, 7), (8, 11), (12, 1)],
    ],

    # k=14.
    14: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (12, 0)],
        [(0, 2), (3, 5), (6, 8), (9, 11), (12, 0)],
    ],

    # k=15: tight ring.
    15: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (0, 12)],
        [(0, 3), (4, 7), (8, 11), (12, 14), (12, 0)],
    ],

    # k=18: tight ring.
    18: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (12, 15), (0, 15)],
        [(0, 3), (4, 7), (8, 11), (12, 15), (16, 1)],
    ],

    # k=21.
    21: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (12, 15), (15, 18), (0, 18)],
    ],

    # k=24.
    24: [
        [(0, 3), (3, 6), (6, 9), (9, 12), (12, 15), (15, 18), (18, 21), (0, 21)],
    ],
}


# Filter to SP-feasible only.
for k in list(CHORD_CONFIGS):
    feasible = [cs for cs in CHORD_CONFIGS[k] if is_sp_feasible(k, cs)]
    if feasible:
        CHORD_CONFIGS[k] = feasible
    else:
        CHORD_CONFIGS[k] = []


def compute_pi_D(gamma: int, m_1: int, chords_1, model: str) -> set:
    if model == 'SR':
        fibers = spoke_only_fiber_distribution_fast(m_1 + gamma)
    else:
        fibers, _, _ = fiber_distribution_fast(m_1, gamma, chords_1)
    return projection_support_fast(fibers, d_positions_for(m_1, gamma))


def compute_pi_U(gamma: int, k_2: int, chords_2, model: str) -> set:
    if model == 'SR':
        fibers = spoke_only_fiber_distribution_fast(gamma + k_2)
    else:
        fibers, _, _ = fiber_distribution_fast(gamma, k_2, chords_2)
    return projection_support_fast(fibers, u_positions_for(gamma, k_2))


def test_pair(gamma, m_1, ch1, model1, k_2, ch2, model2):
    """Return result dict."""
    t0 = time.time()
    S1 = compute_pi_D(gamma, m_1, ch1, model1)
    S2 = compute_pi_U(gamma, k_2, ch2, model2)
    forward = S1 & S2
    S2_rev = {s[::-1] for s in S2}
    reverse = S1 & S2_rev
    dt = time.time() - t0
    compatible = bool(forward or reverse)
    return {
        'type': 'pair',
        'key': make_key(gamma, m_1, ch1, model1, k_2, ch2, model2),
        'gamma': gamma,
        't1': {'m_1': m_1, 'chords': ch1, 'model': model1},
        't2': {'k_2': k_2, 'chords': ch2, 'model': model2},
        's1_size': len(S1),
        's2_size': len(S2),
        'fwd_size': len(forward),
        'rev_size': len(reverse),
        'compatible': compatible,
        'time_seconds': round(dt, 3),
        # Tag whether this is a "strict-Latin" pair (where the paper conjecture
        # applies) -- both tires must have every O-face of exactly 3 B_in edges.
        't1_all_three': has_all_three_faces(gamma, ch1) if model1 == 'SP' else None,
        't2_all_three': has_all_three_faces(k_2, ch2) if model2 == 'SP' else None,
    }


def estimate_cost(n: int) -> float:
    """Rough estimate of seconds to compute fiber distribution at cycle length n.
    Calibrated from prior benchmarks (n=18: 0.1s; n=24: 6.6s)."""
    if n <= 12:
        return 0.01
    return 0.001 * (2 ** (n - 12))


# Cap on cycle length to avoid memory blowup.
# At n=27: 3 * 2^26 paths × 27 bytes ≈ 5 GB.
# At n=30: would be ~40 GB.
MAX_N = 27

# Time cap on a single tire computation.
MAX_TIRE_TIME_SEC = 600  # 10 minutes max per tire


def has_all_three_faces(k: int, chords) -> bool:
    """The strict precondition of the paper's Latin conjecture:
    every O-face has exactly 3 B_in edges."""
    if not chords:
        return k == 3
    faces = compute_faces_from_chords(k, chords)
    return all(len(f) == 3 for f in faces)


def add_rotations(k: int, chords):
    """Generate all cyclic rotations of a chord set on C_k (yields tuples)."""
    seen = set()
    for shift in range(k):
        rotated = tuple(sorted([tuple(sorted([(a + shift) % k, (b + shift) % k]))
                                 for (a, b) in chords]))
        if rotated not in seen:
            seen.add(rotated)
            yield [list(c) for c in rotated]


def main(budget_hours: float = 6.0):
    deadline = time.time() + budget_hours * 3600
    log_event({'type': 'session_start', 'budget_hours': budget_hours,
               'pid': os.getpid()})

    already_tested = read_existing_keys()
    print(f"[start] already tested: {len(already_tested)} pairs")
    log_event({'type': 'startup', 'already_tested_count': len(already_tested)})

    # Iterate over k values in increasing order.
    k_values = sorted(CHORD_CONFIGS.keys())
    m_values = [3, 6, 9, 12, 15]
    k2_values = [3, 6, 9, 12, 15]

    n_counterexamples = 0
    n_tested = 0
    n_skipped = 0

    for gamma in k_values:
        if time.time() >= deadline:
            break
        cs_list = CHORD_CONFIGS[gamma]
        if not cs_list:
            log_event({'type': 'k_skip_no_configs', 'gamma': gamma})
            continue
        # Prioritize strict-Latin (all-3-faces) configs for the conjecture test.
        cs_list_sorted = sorted(cs_list, key=lambda cs: not has_all_three_faces(gamma, cs))
        log_event({'type': 'k_start', 'gamma': gamma,
                   'n_configs': len(cs_list_sorted),
                   'n_all3_configs': sum(1 for cs in cs_list_sorted if has_all_three_faces(gamma, cs))})

        for m_1 in m_values:
            if time.time() >= deadline: break
            for ch1 in cs_list_sorted:
                if time.time() >= deadline: break
                n_1 = m_1 + gamma
                if n_1 > MAX_N:
                    continue
                if estimate_cost(n_1) > MAX_TIRE_TIME_SEC:
                    continue
                for k_2 in k2_values:
                    if time.time() >= deadline: break
                    cs2_list = CHORD_CONFIGS.get(k_2, [])
                    if not cs2_list:
                        continue
                    cs2_sorted = sorted(cs2_list, key=lambda cs: not has_all_three_faces(k_2, cs))
                    for ch2 in cs2_sorted:
                        if time.time() >= deadline: break
                        n_2 = gamma + k_2
                        if n_2 > MAX_N:
                            continue
                        if estimate_cost(n_2) > MAX_TIRE_TIME_SEC:
                            continue

                        key = make_key(gamma, m_1, ch1, 'SP', k_2, ch2, 'SP')
                        if key in already_tested:
                            n_skipped += 1
                            continue

                        budget_remaining = deadline - time.time()
                        # Conservative: skip if estimated cost > budget remaining.
                        est = estimate_cost(n_1) + estimate_cost(n_2) + 0.5
                        if est > budget_remaining:
                            log_event({'type': 'cost_skip', 'key': key,
                                       'est_seconds': est,
                                       'budget_remaining': budget_remaining})
                            continue

                        try:
                            result = test_pair(gamma, m_1, ch1, 'SP', k_2, ch2, 'SP')
                        except MemoryError as e:
                            log_event({'type': 'memory_error', 'key': key,
                                       'error': str(e)})
                            continue
                        except Exception as e:
                            log_event({'type': 'exception', 'key': key,
                                       'error': type(e).__name__ + ": " + str(e)})
                            continue

                        n_tested += 1
                        already_tested.add(key)
                        log_event(result)
                        if not result['compatible']:
                            n_counterexamples += 1
                            log_event({'type': 'COUNTEREXAMPLE', 'pair': result})
                            print(f"\n*** COUNTEREXAMPLE FOUND at γ={gamma}, "
                                  f"T1={(m_1, ch1)}, T2={(k_2, ch2)} ***\n")

                        # Print periodic progress
                        if n_tested % 25 == 0:
                            elapsed = time.time() - (deadline - budget_hours * 3600)
                            print(f"[{elapsed/60:.1f}m] tested {n_tested}, "
                                  f"skipped {n_skipped}, "
                                  f"counterexamples {n_counterexamples}, "
                                  f"current γ={gamma}")

    log_event({'type': 'session_end',
               'n_tested': n_tested,
               'n_skipped': n_skipped,
               'n_counterexamples': n_counterexamples,
               'elapsed_hours': (time.time() - (deadline - budget_hours * 3600)) / 3600})
    print(f"\n[done] tested {n_tested}, skipped {n_skipped}, "
          f"counterexamples {n_counterexamples}")


if __name__ == '__main__':
    budget = float(sys.argv[1]) if len(sys.argv) > 1 else 6.0
    main(budget_hours=budget)
