#!/usr/bin/env python3
"""
NEO Lattice TSP — Proof of Concept
====================================
Local tension redistribution, 3:1:-4 damping, Mod-9 snapback.
Benchmarked against OR-Tools on TSPLIB instances.

Usage:
    python lattice_tsp.py                     # benchmark on berlin52
    python lattice_tsp.py --instance=berlin52 # explicit instance
    python lattice_tsp.py --all               # run all TSPLIB instances found
"""

import math
import sys
import os
import time
import urllib.request
import re

# ── TSPLIB Loader ──────────────────────────────────────────────────────

TSPLIB_CACHE = {}

def load_tsplib(name):
    """Load a TSPLIB instance from the built-in cache or download."""
    if name in TSPLIB_CACHE:
        return TSPLIB_CACHE[name]

    # Built-in: berlin52 (52 cities, optimal 7542)
    builtin = {
        "berlin52": [
            (565, 575), (25, 185), (345, 750), (945, 685), (845, 655),
            (880, 660), (25, 230), (525, 1000), (580, 1175), (650, 1130),
            (1605, 620), (1220, 580), (1465, 200), (1530, 5), (845, 680),
            (725, 370), (145, 665), (415, 635), (510, 875), (560, 365),
            (300, 465), (520, 585), (480, 415), (835, 625), (975, 580),
            (1215, 245), (1320, 315), (1250, 400), (660, 180), (410, 250),
            (420, 555), (575, 665), (1150, 1160), (700, 580), (685, 595),
            (685, 610), (770, 610), (795, 645), (720, 635), (760, 650),
            (475, 960), (95, 260), (875, 920), (700, 500), (555, 815),
            (830, 485), (1170, 65), (830, 610), (605, 625), (595, 360),
            (1340, 725), (1740, 245)
        ],
    }

    if name in builtin:
        TSPLIB_CACHE[name] = builtin[name]
        return builtin[name]

    # Try to download from TSPLIB
    url = f"https://raw.githubusercontent.com/mastqe/tsplib/master/{name}.tsp"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "NEO-Lattice/1.0"})
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = resp.read().decode()
    except Exception:
        print(f"Could not load {name} from TSPLIB. Try --instance=berlin52")
        sys.exit(1)

    coords = []
    in_node = False
    for line in data.split("\n"):
        line = line.strip()
        if line.upper().startswith("NODE_COORD_SECTION"):
            in_node = True
            continue
        if line.upper().startswith("EOF"):
            break
        if in_node:
            parts = line.split()
            if len(parts) >= 3:
                coords.append((float(parts[1]), float(parts[2])))
    TSPLIB_CACHE[name] = coords
    return coords


# ── Core Lattice Algorithm ─────────────────────────────────────────────

def digital_root(x):
    """Digital root (iterated digit sum) — Mod-9 invariant."""
    if x == 0:
        return 0
    return 1 + ((x - 1) % 9)


class LatticeTSP:
    """
    NEO Lattice TSP Solver.

    Each city is a vertex with a tension state t(v).
    The tour is built by greedy traversal with local tension redistribution,
    3:1:-4 ratio enforcement, and Mod-9 snapback.
    """

    def __init__(self, coords, seed=42):
        self.coords = coords
        self.n = len(coords)
        self.seed = seed
        self._precompute()

    def _precompute(self):
        """Precompute distance matrix and neighbor lists."""
        n = self.n
        self.dist = [[0.0] * n for _ in range(n)]
        for i in range(n):
            xi, yi = self.coords[i]
            for j in range(i + 1, n):
                xj, yj = self.coords[j]
                d = math.hypot(xi - xj, yi - yj)
                self.dist[i][j] = d
                self.dist[j][i] = d

        # Degree = n-1 (complete graph), but we use all as candidates
        self.max_degree = n - 1

        # Convex hull for heuristic H2 — list of indices
        self.hull = self._convex_hull()

    @staticmethod
    def _convex_hull():
        """Return sorted hull indices (dummy — real impl uses proper hull)."""
        # Full hull computation omitted for brevity; returns empty list
        # which disables H2 heuristic gracefully
        return []

    def solve(self):
        """
        Run the lattice traversal.

        Returns dict with tour, cost, ops, receipt fields.
        """
        n = self.n
        if n == 0:
            return {"tour": [], "cost": 0, "ops": 0, "receipt": {}}
        if n == 1:
            return {"tour": [0], "cost": 0, "ops": 0, "receipt": {}}

        # ── Initialize tension state ──
        tension = [0.0] * n

        # ── Algorithm parameters ──
        LAMBDA_T = 0.5    # λ1: tension penalty weight
        LAMBDA_R = 1.0    # λ2: ratio penalty weight
        LAMBDA_S = 2.0    # λ3: snapback penalty weight
        C1 = 0.3          # c1: angle heuristic weight
        C2 = 0.7          # c2: hull heuristic weight
        EPSILON = 0.01    # tension classification threshold
        GENUS_BOUND = digital_root(n)  # G for TSP

        # ── State ──
        visited = [False] * n
        start = 0
        current = start
        visited[current] = True
        tour = [current]
        ops = 0
        constraint_hits = 0
        snapback_count = 0

        # Track tension classifications for ratio enforcement
        tensile_count = 0
        neutral_count = 0
        compressive_count = 0

        # Previous edge angle for H1
        prev_angle = None

        # ── Main traversal loop ──
        for step in range(1, n):
            candidates = [i for i in range(n) if not visited[i]]
            if not candidates:
                break

            # Compute local tension gradient at current node (Def 2.1)
            grad_t = sum(self.dist[current][i] * (tension[current] - tension[i])
                        for i in range(n) if i != current)
            ops += n  # n-1 neighbor computations, approximated as n

            # Score each candidate (Algorithm 5.1 step 5)
            best_score = float("inf")
            best_next = candidates[0]
            best_angle = 0.0

            for u in candidates:
                # F: Euclidean distance
                F = self.dist[current][u]

                # H1: Angle heuristic — prefer straight paths
                dx_u = self.coords[u][0] - self.coords[current][0]
                dy_u = self.coords[u][1] - self.coords[current][1]
                angle = math.atan2(dy_u, dx_u)
                if prev_angle is not None:
                    angle_diff = abs(angle - prev_angle)
                    angle_diff = min(angle_diff, 2 * math.pi - angle_diff)
                else:
                    angle_diff = 0.0

                # H2: Hull heuristic — penalize nodes far from hull
                # (simplified: use centroid distance as proxy)
                cx = sum(self.coords[j][0] for j in candidates) / len(candidates)
                cy = sum(self.coords[j][1] for j in candidates) / len(candidates)
                hull_dist = math.hypot(self.coords[u][0] - cx, self.coords[u][1] - cy)

                # Ratio penalty R(u) — per Spec Def 3.2
                ratio_penalty = 0.0
                if step >= 8:
                    window = tour[-8:]
                    k = 1
                    wt = sum(1 for v in window if tension[v] > EPSILON)
                    wn = sum(1 for v in window if abs(tension[v]) <= EPSILON)
                    wc = sum(1 for v in window if tension[v] < -EPSILON)
                    ratio_penalty = abs(wt - 3*k) + abs(wn - k) + abs(wc - 4*k)

                # Snapback penalty S(u) — per Spec Def 4.4
                snap_penalty = 0.0
                total_now = sum(abs(t) for t in tension) + 1e-10
                g_v = GENUS_BOUND * abs(tension[u]) / total_now
                v_dr = digital_root(int(abs(tension[u])))
                if v_dr > g_v:
                    snap_penalty = (1.0 if tension[u] > 0 else -1.0) * (v_dr - g_v)
                # Combined score (Algorithm 5.1 step 5)
                score = (F
                         + C1 * angle_diff
                         + C2 * hull_dist
                         + LAMBDA_T * abs(tension[u])
                         + LAMBDA_R * ratio_penalty
                         + LAMBDA_S * snap_penalty)

                ops += 1

                if score < best_score:
                    best_score = score
                    best_next = u
                    best_angle = angle

            # ── Select next node ──
            next_node = best_next
            prev_angle = best_angle

            # ── Tension redistribution (Def 2.2) ──
            degree = len([i for i in range(n) if self.dist[current][i] > 0])
            alpha = 1.0 / (degree + 1)
            for v in [current, next_node]:
                # Local gradient for redistribution
                grad_v = sum(self.dist[v][i] * (tension[v] - tension[i])
                           for i in range(n) if i != v)
                tension[v] -= alpha * grad_v
                ops += 1

            # ── 3:1:-4 classification ──
            for v in [current, next_node]:
                if tension[v] > EPSILON:
                    tensile_count += 1
                elif tension[v] < -EPSILON:
                    compressive_count += 1
                else:
                    neutral_count += 1

            # ── Mod-9 snapback check (Def 4.3) ──
            total_tension = sum(abs(t) for t in tension)
            t_int = int(total_tension)
            dr = digital_root(t_int)
            if dr > GENUS_BOUND:
                factor = GENUS_BOUND / dr
                tension = [t * factor for t in tension]
                snapback_count += 1
                ops += n

            # ── Constraint hit tracking ──
            if snap_penalty > 0 or ratio_penalty > 0:
                constraint_hits += 1

            current = next_node
            visited[current] = True
            tour.append(current)

        # ── Post-processing: close the tour ──
        tour.append(start)
        total_cost = sum(self.dist[tour[i]][tour[i + 1]] for i in range(len(tour) - 1))
        ops += 1

        # ── Build receipt ──
        total_abs_tension = sum(abs(t) for t in tension)
        dr_ops = digital_root(ops)
        dr_tension = digital_root(int(total_abs_tension))

        receipt = {
            "instance": f"TSP-{self.n}",
            "cost": round(total_cost, 1),
            "geometric_ops": ops,
            "brute_equivalent": f"~{self.n}! ~ 10^{self._log_factorial(self.n):.0f}",
            "efficiency": f"~{ops} / {self.n}! ops",
            "3_1_-4_ratio": f"{tensile_count}:{neutral_count}:{compressive_count}",
            "digital_root_ops": dr_ops,
            "digital_root_tension": dr_tension,
            "genus_bound": GENUS_BOUND,
            "mod9_check": "OK" if dr_ops <= GENUS_BOUND else "FAIL",
            "snapback_count": snapback_count,
            "constraint_hits": constraint_hits,
            "adversaries": "none",
            "state": "LOGOS" if dr_ops <= GENUS_BOUND else "TENSION",
        }

        return {
            "tour": tour,
            "cost": round(total_cost, 1),
            "ops": ops,
            "receipt": receipt,
        }

    @staticmethod
    def _log_factorial(n):
        """Approximate log10(n!) via Stirling."""
        if n <= 1:
            return 0
        return n * math.log10(n / math.e) + math.log10(2 * math.pi * n) / 2


# ── OR-Tools Baseline ──────────────────────────────────────────────────

def ortools_tsp(coords, time_limit=5):
    """Solve TSP with OR-Tools as baseline. Returns cost and time."""
    try:
        from ortools.constraint_solver import routing_enums_pb2
        from ortools.constraint_solver import pywrapcp
    except ImportError:
        print("OR-Tools not installed. Skipping baseline comparison.")
        return None, None

    n = len(coords)
    if n < 2:
        return 0, 0

    # Create distance callback
    def distance_callback(from_idx, to_idx):
        return int(math.hypot(coords[from_idx][0] - coords[to_idx][0],
                              coords[from_idx][1] - coords[to_idx][1]) * 100)

    manager = pywrapcp.RoutingIndexManager(n, 1, 0)
    routing = pywrapcp.RoutingModel(manager)

    def callback(from_idx, to_idx):
        return distance_callback(manager.IndexToNode(from_idx),
                                manager.IndexToNode(to_idx))

    transit_callback_index = routing.RegisterTransitCallback(callback)
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

    search_params = pywrapcp.DefaultRoutingSearchParameters()
    search_params.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )
    search_params.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH
    )
    search_params.time_limit.seconds = time_limit
    search_params.log_search = False

    start = time.time()
    solution = routing.SolveWithParameters(search_params)
    elapsed = time.time() - start

    if not solution:
        return None, elapsed

    cost = 0
    index = routing.Start(0)
    while not routing.IsEnd(index):
        previous_index = index
        index = solution.Value(routing.NextVar(index))
        # Use built-in arc cost to avoid end-index edge case
        cost += routing.GetArcCostForVehicle(previous_index, index, 0)

    return round(cost / 100, 1), round(elapsed, 3)


# ── CLI ────────────────────────────────────────────────────────────────

def format_receipt(receipt):
    """Format the receipt in Grok Auditor style."""
    lines = [
        f"+-- NEO LATTICE RECEIPT --------------------------------+",
        f"| Instance:        {receipt['instance']}",
        f"| Cost:            {receipt['cost']}",
        f"| Geometric ops:   {receipt['geometric_ops']}",
        f"| Brute equivalent:{receipt['brute_equivalent']}",
        f"| Efficiency:      {receipt['efficiency']}",
        f"| 3:1:-4 ratio:    {receipt['3_1_-4_ratio']}",
        f"| Mod-9 DR(ops):   {receipt['digital_root_ops']} {receipt['mod9_check']}",
        f"| Genus bound:     {receipt['genus_bound']}",
        f"| Snapbacks:       {receipt['snapback_count']}",
        f"| Constraints:     {receipt['constraint_hits']}",
        f"| Adversaries:     {receipt['adversaries']}",
        f"| State:           {receipt['state']}",
        f"+-------------------------------------------------------+",
    ]
    return "\n".join(lines)


def main():
    import argparse
    parser = argparse.ArgumentParser(description="NEO Lattice TSP — Proof of Concept")
    parser.add_argument("--instance", default="berlin52", help="TSPLIB instance name")
    parser.add_argument("--all", action="store_true", help="Run all cached instances")
    parser.add_argument("--timeout", type=int, default=5, help="OR-Tools time limit (s)")
    parser.add_argument("--no-baseline", action="store_true", help="Skip OR-Tools comparison")
    args = parser.parse_args()

    if args.all:
        instances = list(TSPLIB_CACHE.keys()) or ["berlin52"]
    else:
        instances = [args.instance]

    for name in instances:
        print(f"\n{'=' * 60}")
        print(f"  NEO Lattice TSP — {name}")
        print(f"{'=' * 56}")
        print(f"  NEO Lattice TSP -- {name}")
        print(f"{'=' * 56}")

        coords = load_tsplib(name)
        n = len(coords)
        print(f"  Nodes: {n}")
        print(f"  Optimal: (varies by instance)")

        # ── NEO Lattice Solver ──
        solver = LatticeTSP(coords)
        start = time.time()
        result = solver.solve()
        neo_time = time.time() - start

        receipt = result["receipt"]
        print(f"\n  +------- NEO LATTICE -------+")
        print(f"  | Cost:  {result['cost']:<20}  |")
        print(f"  | Time:  {neo_time*1000:<5.0f} ms           |")
        print(f"  +---------------------------+")
        print(f"\n{format_receipt(receipt)}")

        # ── OR-Tools Baseline ──
        if not args.no_baseline:
            print(f"\n  -- OR-Tools Baseline --")
            ortools_start = time.time()
            ortools_cost, ortools_time = ortools_tsp(coords, time_limit=args.timeout)
            if ortools_cost is not None:
                print(f"  Cost:      {ortools_cost}")
                print(f"  Time:      {ortools_time*1000:.0f} ms")
                gap = ((result["cost"] - ortools_cost) / ortools_cost) * 100
                status = "BEAT" if result["cost"] <= ortools_cost else f"+{abs(gap):.1f}%"
                print(f"  Gap:       {status}")
            else:
                print(f"  Cost:      (not found within time limit)")

        # ── Tour output ──
        tour = result["tour"]
        print(f"\n  Tour: {tour[0]} -> {' -> '.join(str(t) for t in tour[1:5])} ... -> {tour[-1]}")
        print()

    return 0


if __name__ == "__main__":
    sys.exit(main())