#!/usr/bin/env python3
"""0/1 TSP-style repair experiment for recovered HCP block graphs.

This focuses on the HCP-491 cubic-ish graph.  The workflow is:

1. Build randomized 2-factors with worker_hc_search.greedy_two_factor until a
   two-cycle cover is found.
2. Convert the two disjoint real cycles into one cyclic permutation by cutting
   one edge on each cycle and inserting one real cross-edge.  The other new
   boundary edge is allowed to be a non-edge.
3. Minimize the 0/1 objective: number of adjacent pairs in the permutation that
   are not graph edges.  A Hamiltonian cycle has objective 0; the best observed
   run has objective 1.
4. Try 2-opt reversals:
      remove (p[i], p[i+1]) and (p[j], p[j+1])
      add    (p[i], p[j])   and (p[i+1], p[j+1])
   The delta is computed under the 0/1 non-edge objective.
"""

import argparse
import glob
import hashlib
import importlib.util
import os
import random
import time


def load_local(name):
    here = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(here, f"{name}.py")
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


hc = load_local("worker_hc_search")


def default_491_path():
    hits = glob.glob("/data/train/*HCP-491-420.cnf.xz")
    if not hits:
        raise FileNotFoundError("could not find /data/train/*HCP-491-420.cnf.xz")
    return hits[0]


def build_graph(cnf):
    hcp_extract = hc.load_extractor()
    _, _, _, layout, edges_by_var, _, n, undirected = hcp_extract.extract(cnf)
    all_edges = [hc.edge_key(a, b) for a, b in undirected]
    edge_set = set(all_edges)
    edge_index = {e: i for i, e in enumerate(all_edges)}
    adj_edges = [[] for _ in range(n)]
    graph_adj = [set() for _ in range(n)]
    for a, b in all_edges:
        ei = edge_index[(a, b)]
        adj_edges[a].append(ei)
        adj_edges[b].append(ei)
        graph_adj[a].add(b)
        graph_adj[b].add(a)

    fixed_block = layout.get("fixed_block")
    preferred_next = None
    if layout.get("low_units"):
        e = edges_by_var.get(abs(layout["low_units"][0]))
        if e is not None and fixed_block in e:
            preferred_next = e[0] if e[1] == fixed_block else e[1]

    return {
        "n": n,
        "layout": layout,
        "all_edges": all_edges,
        "edge_set": edge_set,
        "adj_edges": adj_edges,
        "graph_adj": graph_adj,
        "fixed_block": fixed_block,
        "preferred_next": preferred_next,
    }


def find_two_cycle_cover(g, seed, max_attempts):
    rng = random.Random(seed)
    best_count = None
    best_lengths = None
    for attempt in range(max_attempts):
        selected = hc.greedy_two_factor(g["n"], g["all_edges"], g["adj_edges"], rng)
        if selected is None:
            continue
        cycles, _ = hc.cycle_cover_from_selected(g["n"], selected)
        if cycles is None:
            continue
        lengths = sorted((len(c) for c in cycles), reverse=True)
        if best_count is None or len(cycles) < best_count:
            best_count = len(cycles)
            best_lengths = lengths
        if len(cycles) == 2:
            return selected, cycles, attempt, best_count, best_lengths
    return None, None, None, best_count, best_lengths


def missing(edge_set, a, b):
    return 0 if hc.edge_key(a, b) in edge_set else 1


def objective(edge_set, tour):
    n = len(tour)
    return sum(missing(edge_set, tour[i], tour[(i + 1) % n]) for i in range(n))


def bad_positions(edge_set, tour):
    n = len(tour)
    return [i for i in range(n) if missing(edge_set, tour[i], tour[(i + 1) % n])]


def rotate_to_root(tour, fixed_block, preferred_next):
    if fixed_block is None or fixed_block not in tour:
        return tour
    i = tour.index(fixed_block)
    fwd = tour[i:] + tour[:i]
    rev = [tour[i]] + list(reversed(tour[:i] + tour[i + 1 :]))
    if preferred_next is not None:
        if len(fwd) > 1 and fwd[1] == preferred_next:
            return fwd
        if len(rev) > 1 and rev[1] == preferred_next:
            return rev
    return fwd


def cycle_join_candidates(cycles, graph_adj):
    c0, c1 = cycles
    pos1 = {v: i for i, v in enumerate(c1)}
    out = []
    for i, a in enumerate(c0):
        for c in sorted(graph_adj[a]):
            j = pos1.get(c)
            if j is not None:
                out.append((i, j, a, c))
    return out


def initial_tour_from_cross(cycles, i, j, reverse_second):
    c0, c1 = cycles
    a = c0[i]
    b = c0[(i + 1) % len(c0)]
    first = c0[i + 1 :] + c0[: i + 1]

    if reverse_second:
        c1r = list(reversed(c1))
        jr = c1r.index(c1[j])
        second = c1r[jr:] + c1r[:jr]
    else:
        second = c1[j:] + c1[:j]

    return first + second, {
        "cut_cycle0": (a, b),
        "cross_edge": (a, c1[j]),
        "second_reversed": reverse_second,
    }


def best_initial_tour(g, cycles):
    candidates = cycle_join_candidates(cycles, g["graph_adj"])
    if not candidates:
        return None, None, None

    best = None
    best_meta = None
    best_cost = None
    for i, j, a, c in candidates:
        for reverse_second in (False, True):
            tour, meta = initial_tour_from_cross(cycles, i, j, reverse_second)
            cost = objective(g["edge_set"], tour)
            key = (cost, i, j, reverse_second)
            if best is None or key < best[0]:
                best = (key, tour)
                best_cost = cost
                best_meta = {
                    **meta,
                    "cross_indices": (i, j),
                    "cross_edge": (a, c),
                    "closing_edge": (tour[-1], tour[0]),
                    "initial_cost": cost,
                }
    return best[1], best_meta, best_cost


def two_opt_delta(edge_set, tour, i, j):
    n = len(tour)
    if i == j or (i + 1) % n == j or (j + 1) % n == i:
        return None
    if i > j:
        i, j = j, i
    if i == 0 and j == n - 1:
        return None
    a, b = tour[i], tour[(i + 1) % n]
    c, d = tour[j], tour[(j + 1) % n]
    old = missing(edge_set, a, b) + missing(edge_set, c, d)
    new = missing(edge_set, a, c) + missing(edge_set, b, d)
    return new - old, i, j


def apply_two_opt(tour, i, j):
    if i > j:
        i, j = j, i
    tour[i + 1 : j + 1] = reversed(tour[i + 1 : j + 1])


def repair_tour(g, tour, seed, iterations, plateau_prob, uphill_prob):
    rng = random.Random(seed)
    edge_set = g["edge_set"]
    graph_adj = g["graph_adj"]
    cur = objective(edge_set, tour)
    best_tour = tour[:]
    best_cost = cur
    moves = []

    for it in range(iterations):
        if cur == 0:
            break
        pos = {v: i for i, v in enumerate(tour)}
        candidates = []

        # Main neighborhood: touch endpoints of currently bad adjacencies and
        # pair them with real graph neighbors, so at least one proposed new edge
        # is guaranteed real.
        for bi in bad_positions(edge_set, tour):
            for x in (tour[bi], tour[(bi + 1) % len(tour)]):
                for y in graph_adj[x]:
                    i, j = pos[x], pos[y]
                    d = two_opt_delta(edge_set, tour, i, j)
                    if d is not None:
                        candidates.append(d)

        # Small diversification set.  These are ordinary 2-opt moves; most are
        # rejected unless plateau/uphill probabilities allow them.
        n = len(tour)
        for _ in range(64):
            i, j = sorted(rng.sample(range(n), 2))
            d = two_opt_delta(edge_set, tour, i, j)
            if d is not None:
                candidates.append(d)

        rng.shuffle(candidates)
        accepted = None
        for delta, i, j in candidates:
            if delta < 0:
                accepted = (delta, i, j, "improve")
                break
            if delta == 0 and rng.random() < plateau_prob:
                accepted = (delta, i, j, "plateau")
                break
            if delta > 0 and rng.random() < uphill_prob:
                accepted = (delta, i, j, "uphill")
                break

        if accepted is None:
            break

        delta, i, j, kind = accepted
        before = cur
        apply_two_opt(tour, i, j)
        cur += delta
        moves.append((it, kind, i, j, before, cur))
        if cur < best_cost:
            best_cost = cur
            best_tour = tour[:]

    return best_tour, best_cost, moves


def tour_digest(tour):
    text = " ".join(map(str, tour)) + "\n"
    return hashlib.sha256(text.encode("ascii")).hexdigest()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cnf", nargs="?", default=default_491_path())
    ap.add_argument("--cover-seed", type=int, default=491)
    ap.add_argument("--repair-seed", type=int, default=7)
    ap.add_argument("--max-cover-attempts", type=int, default=10000)
    ap.add_argument("--iterations", type=int, default=20000)
    ap.add_argument("--plateau-prob", type=float, default=0.20)
    ap.add_argument("--uphill-prob", type=float, default=0.002)
    ap.add_argument("--output-tour")
    ap.add_argument("--print-tour", action="store_true")
    args = ap.parse_args()

    start = time.time()
    g = build_graph(args.cnf)
    selected, cycles, attempt, best_count, best_lengths = find_two_cycle_cover(
        g, args.cover_seed, args.max_cover_attempts
    )
    if cycles is None:
        print(f"result FAIL no_two_cycle_cover best_count {best_count} best_lengths {best_lengths}")
        return 1

    initial, meta, initial_cost = best_initial_tour(g, cycles)
    if initial is None:
        print("result FAIL no_cross_edge_between_cover_cycles")
        return 1

    repaired, repaired_cost, moves = repair_tour(
        g,
        initial[:],
        seed=args.repair_seed,
        iterations=args.iterations,
        plateau_prob=args.plateau_prob,
        uphill_prob=args.uphill_prob,
    )
    repaired = rotate_to_root(repaired, g["fixed_block"], g["preferred_next"])
    final_cost = objective(g["edge_set"], repaired)
    bad = [(repaired[i], repaired[(i + 1) % len(repaired)]) for i in bad_positions(g["edge_set"], repaired)]

    print(f"instance {os.path.basename(args.cnf)}")
    print(f"blocks {g['n']} edges {len(g['all_edges'])}")
    print(f"fixed_block {g['fixed_block']} preferred_next {g['preferred_next']}")
    print(f"cover_seed {args.cover_seed} two_cycle_attempt {attempt}")
    print(f"cover_cycle_lengths {sorted((len(c) for c in cycles), reverse=True)}")
    print(f"cross_indices {meta['cross_indices']} second_reversed {meta['second_reversed']}")
    print(f"cut_cycle0 {meta['cut_cycle0']} cross_edge {meta['cross_edge']} closing_edge {meta['closing_edge']}")
    print(f"initial_objective {initial_cost}")
    print(f"repair_seed {args.repair_seed} moves {len(moves)} best_objective {repaired_cost} final_objective {final_cost}")
    print(f"bad_adjacencies {bad}")
    print(f"tour_sha256 {tour_digest(repaired)}")
    print("tour_prefix", " ".join(map(str, repaired[:40])))
    print(f"elapsed {time.time() - start:.3f}s")
    if moves:
        print("move_prefix", moves[:10])
    if args.output_tour:
        with open(args.output_tour, "w") as f:
            f.write(" ".join(map(str, repaired)))
            f.write("\n")
        print(f"wrote_tour {args.output_tour}")
    if args.print_tour:
        print("tour", " ".join(map(str, repaired)))
    return 0 if final_cost == 0 else 2


if __name__ == "__main__":
    raise SystemExit(main())
