# -*- coding: utf-8 -*-
"""TSP_CW.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1SWb-xRbWCwhjhEH3oP-TA4VK-gjgkXCl
"""

# --- Numba bootstrap: auto-install if missing, safe fallback ---
import sys, subprocess

NUMBA_AVAILABLE = False

def _njit_direct(f):
    # No-op decorator để code vẫn chạy nếu không có numba
    return f

try:
    from numba import njit  # type: ignore
    NUMBA_AVAILABLE = True
except Exception:
    try:
        # Thử cài numba (yêu cầu internet + quyền pip)
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "numba>=0.59"],
            check=False
        )
        from numba import njit  # retry
        NUMBA_AVAILABLE = True
    except Exception:
        # Fallback: chạy không JIT
        njit = _njit_direct  # type: ignore
        # Cảnh báo một lần khi import
        print("[tsp-core] WARNING: 'numba' unavailable; running without JIT.")

# -*- coding: utf-8 -*-
# TSP core (NumPy + Numba, no OOP, no Python lists in JIT code)
# Public Numba API:
#   build_tsp_route(D, algo_id, local_search, lam, lam1) -> np.ndarray[int64]
# Mapping:
#   algo_id = 0  -> Clarke–Wright (CW), lambda ignored
#   algo_id = 1  -> Greedy (cheapest insertion by length) + optional 2-opt, lambda ignored
#   algo_id = 2  -> Greedy with lambda (deform cost using lam, lam1) + optional 2-opt
#
# local_search <= 0: skip 2-opt; local_search > 0: run up to that many passes.
#
# Optional non-JIT wrapper (translate dict to ids):
#   build_tsp_route_from_param(D, param_dict) -> np.ndarray[int64]
#     param["algo"] in {"CW","greedy","greedy_lambda",":greedy with lambda"}
#     param["lambda"] = [lam, lam1] (only for "greedy_lambda")
#     param["local_search"] = int

import numpy as np
from numba import njit

# ==============================
# Helpers (JIT-safe, no lists)
# ==============================
@njit
def tour_length_from_D(D: np.ndarray, tour: np.ndarray) -> float:
    n = tour.size
    s = 0.0
    for i in range(n):
        s += D[tour[i], tour[(i + 1) % n]]
    return s

@njit
def reverse_segment(tour: np.ndarray, i: int, j: int) -> np.ndarray:
    out = tour.copy()
    a, b = i, j
    while a < b:
        t = out[a]; out[a] = out[b]; out[b] = t
        a += 1; b -= 1
    return out

@njit
def two_opt_tour(D: np.ndarray, tour: np.ndarray, max_passes: int) -> np.ndarray:
    n = tour.size
    best = tour.copy()
    best_cost = tour_length_from_D(D, best)
    for _ in range(max_passes):
        improved = False
        for i in range(n - 1):
            for j in range(i + 1, n):
                cand = reverse_segment(best, i, j)
                c = tour_length_from_D(D, cand)
                if c + 1e-12 < best_cost:
                    best = cand
                    best_cost = c
                    improved = True
                    break
            if improved:
                break
        if not improved:
            break
    return best

# ---------- deform cost (window-4) ----------
@njit
def window_diameter4(D: np.ndarray, a: int, b: int, c: int, d: int) -> float:
    m = D[a, b]
    if D[a, c] > m: m = D[a, c]
    if D[a, d] > m: m = D[a, d]
    if D[b, c] > m: m = D[b, c]
    if D[b, d] > m: m = D[b, d]
    if D[c, d] > m: m = D[c, d]
    return m

@njit
def window_pairwise_sum4(D: np.ndarray, a: int, b: int, c: int, d: int) -> float:
    s = 0.0
    s += D[a, b] + D[a, c] + D[a, d]
    s += D[b, c] + D[b, d]
    s += D[c, d]
    return s

@njit
def deform_reg4(D: np.ndarray, tour: np.ndarray, lam: float, lam1: float) -> float:
    n = tour.size
    s = 0.0
    for i in range(n):
        a = tour[(i - 3) % n]
        b = tour[(i - 2) % n]
        c = tour[(i - 1) % n]
        d = tour[i]
        if lam != 0.0:
            s += lam * window_diameter4(D, a, b, c, d)
        if lam1 != 0.0:
            s += lam1 * window_pairwise_sum4(D, a, b, c, d)
    return s

@njit
def deform_total(D: np.ndarray, tour: np.ndarray, lam: float, lam1: float) -> float:
    return tour_length_from_D(D, tour) + deform_reg4(D, tour, lam, lam1)

@njit
def two_opt_deform(D: np.ndarray, tour: np.ndarray, lam: float, lam1: float, max_passes: int) -> np.ndarray:
    n = tour.size
    best = tour.copy()
    best_cost = deform_total(D, best, lam, lam1)
    for _ in range(max_passes):
        improved = False
        for i in range(n - 1):
            for j in range(i + 1, n):
                cand = reverse_segment(best, i, j)
                c = deform_total(D, cand, lam, lam1)
                if c + 1e-12 < best_cost:
                    best = cand
                    best_cost = c
                    improved = True
                    break
            if improved:
                break
        if not improved:
            break
    return best

# ---------- cheapest insertion (length or deform) ----------
@njit
def insert_at(tour: np.ndarray, pos: int, node: int) -> np.ndarray:
    n = tour.size
    out = np.empty(n + 1, dtype=np.int64)
    for i in range(pos + 1):
        out[i] = tour[i]
    out[pos + 1] = node
    for i in range(pos + 1, n):
        out[i + 1] = tour[i]
    return out

@njit
def find_farthest_pair(D: np.ndarray):
    n = D.shape[0]
    best = -1.0
    a = 0
    b = 1
    for i in range(n):
        for j in range(i + 1, n):
            if D[i, j] > best:
                best = D[i, j]
                a, b = i, j
    return a, b

@njit
def initial_triangle(D: np.ndarray, a: int, b: int, lam: float, lam1: float) -> int:
    n = D.shape[0]
    best_c = -1
    best_score = 1e100
    for c in range(n):
        if c == a or c == b:
            continue
        tri = np.array([a, b, c], dtype=np.int64)
        score = deform_total(D, tri, lam, lam1)
        if score < best_score:
            best_score = score
            best_c = c
    return best_c

@njit
def cheapest_insertion_deform(D: np.ndarray, lam: float, lam1: float) -> np.ndarray:
    n = D.shape[0]
    a, b = find_farthest_pair(D)
    c = initial_triangle(D, a, b, lam, lam1)
    tour = np.array([a, b, c], dtype=np.int64)
    used = np.zeros(n, dtype=np.uint8)
    used[a] = 1; used[b] = 1; used[c] = 1
    while np.sum(used) < n:
        base = deform_total(D, tour, lam, lam1)
        best_delta = 1e100
        best_pos = 0
        best_node = -1
        m = tour.size
        for node in range(n):
            if used[node] == 1:
                continue
            for pos in range(m):
                cand = insert_at(tour, pos, node)
                delta = deform_total(D, cand, lam, lam1) - base
                if delta < best_delta:
                    best_delta = delta
                    best_pos = pos
                    best_node = node
        tour = insert_at(tour, best_pos, best_node)
        used[best_node] = 1
    return tour

@njit
def cheapest_insertion_length(D: np.ndarray) -> np.ndarray:
    return cheapest_insertion_deform(D, 0.0, 0.0)

# ---------- Clarke–Wright (no lists) ----------
@njit
def dsu_find(parent: np.ndarray, x: int) -> int:
    xi = x
    while parent[xi] != xi:
        parent[xi] = parent[parent[xi]]
        xi = parent[xi]
    return xi

@njit
def can_attach(deg: np.ndarray, x: int) -> bool:
    return deg[x] < 2

@njit
def link_neighbor(neigh1: np.ndarray, neigh2: np.ndarray, deg: np.ndarray, u: int, v: int):
    if neigh1[u] == -1:
        neigh1[u] = v
    else:
        neigh2[u] = v
    deg[u] += 1

@njit
def other_end(left_end: np.ndarray, right_end: np.ndarray, rep: int, used: int) -> int:
    a = left_end[rep]
    b = right_end[rep]
    if a == used:
        return b
    else:
        return a

@njit
def clarke_wright_tsp(D: np.ndarray, ls_2opt_passes: int) -> np.ndarray:
    n = D.shape[0]
    m = n - 1  # customers 1..n-1

    # savings for 1<=i<j<=n-1
    npairs = m * (m - 1) // 2
    I = np.empty(npairs, dtype=np.int64)
    J = np.empty(npairs, dtype=np.int64)
    S = np.empty(npairs, dtype=np.float64)
    idx = 0
    for ii in range(1, n - 1):
        for jj in range(ii + 1, n):
            I[idx] = ii
            J[idx] = jj
            S[idx] = D[ii, 0] + D[0, jj] - D[ii, jj]
            idx += 1
    order = np.argsort(-S)

    neigh1 = np.full(n, -1, dtype=np.int64)
    neigh2 = np.full(n, -1, dtype=np.int64)
    deg    = np.zeros(n, dtype=np.int64)

    parent   = np.empty(n, dtype=np.int64)
    size     = np.zeros(n, dtype=np.int64)
    left_end = np.empty(n, dtype=np.int64)
    right_end= np.empty(n, dtype=np.int64)

    for v in range(n):
        parent[v] = v
        size[v]   = 0
        left_end[v]  = -1
        right_end[v] = -1
    for v in range(1, n):
        size[v] = 1
        left_end[v]  = v
        right_end[v] = v

    comps = m
    for t in range(npairs):
        if comps <= 1:
            break
        e = order[t]
        i = I[e]
        j = J[e]
        if not can_attach(deg, i) or not can_attach(deg, j):
            continue
        ri = dsu_find(parent, i)
        rj = dsu_find(parent, j)
        if ri == rj:
            continue
        if not (i == left_end[ri] or i == right_end[ri]):
            continue
        if not (j == left_end[rj] or j == right_end[rj]):
            continue

        link_neighbor(neigh1, neigh2, deg, i, j)
        link_neighbor(neigh1, neigh2, deg, j, i)

        newA = other_end(left_end, right_end, ri, i)
        newB = other_end(left_end, right_end, rj, j)

        if size[ri] >= size[rj]:
            parent[rj] = ri
            size[ri] = size[ri] + size[rj]
            left_end[ri]  = newA
            right_end[ri] = newB
        else:
            parent[ri] = rj
            size[rj] = size[ri] + size[rj]
            left_end[rj]  = newA
            right_end[rj] = newB
        comps -= 1

    rep = -1
    for v in range(1, n):
        rep = dsu_find(parent, v)
        break

    start = left_end[rep]
    end   = right_end[rep]

    chain = np.empty(m, dtype=np.int64)
    pos = 0
    cur = start
    prev = -1
    while True:
        chain[pos] = cur
        pos += 1
        if cur == end:
            break
        n1 = neigh1[cur]
        n2 = neigh2[cur]
        nxt = -1
        if n1 != -1 and n1 != prev:
            nxt = n1
        elif n2 != -1 and n2 != prev:
            nxt = n2
        prev = cur
        cur = nxt

    tsp_cycle = np.empty(n, dtype=np.int64)
    tsp_cycle[0] = 0
    for k in range(m):
        tsp_cycle[k + 1] = chain[k]

    if ls_2opt_passes > 0:
        tsp_cycle = two_opt_tour(D, tsp_cycle, max_passes=ls_2opt_passes)
    return tsp_cycle

# ==============================
# Public Numba API
# ==============================
@njit
def rotate_start_at_zero(route: np.ndarray) -> np.ndarray:
    n = route.size
    pos0 = 0
    for i in range(n):
        if route[i] == 0:
            pos0 = i
            break
    if pos0 == 0:
        return route
    rot = np.empty(n, dtype=np.int64)
    k = 0
    for i in range(pos0, n):
        rot[k] = route[i]; k += 1
    for i in range(0, pos0):
        rot[k] = route[i]; k += 1
    return rot

@njit
def build_tsp_route(D: np.ndarray, algo_id: int, local_search: int, lam=2.4, lam1=1.7) -> np.ndarray:
    """
    algo_id: 0 = CW; 1 = greedy(length); 2 = greedy_lambda (deform with lam, lam1)
    local_search: <=0 skip 2-opt; >0 run up to 'local_search' passes
    lam, lam1: used only if algo_id == 2
    Returns: np.int64 cycle of length n, rotated so route[0] == 0
    """
    n = D.shape[0]
    if n < 3:
        r = np.empty(n, dtype=np.int64)
        for i in range(n):
            r[i] = i
        return r

    if algo_id == 0:  # CW
        ls = local_search if local_search > 0 else 0
        route = clarke_wright_tsp(D, ls)
    elif algo_id == 1:  # greedy by length
        route = cheapest_insertion_length(D)
        if local_search > 0:
            route = two_opt_tour(D, route, local_search)
    else:  # algo_id == 2: greedy with lambda
        route = cheapest_insertion_deform(D, lam, lam1)
        if local_search > 0:
            route = two_opt_deform(D, route, lam, lam1, local_search)

    route = rotate_start_at_zero(route)
    return route

# ==============================
# Optional: non-JIT wrapper (param dict → ids)
# ==============================
def build_tsp_route_from_param(D: np.ndarray, param: dict) -> np.ndarray:
    algo = str(param.get("algo", "CW")).strip().lower()
    if algo in (":greedy with lambda", ":greedy_with_lambda", "greedy-with-lambda"):
        algo = "greedy_lambda"
    if algo == "cw":
        algo_id = 0
    elif algo == "greedy":
        algo_id = 1
    elif algo == "greedy_lambda":
        algo_id = 2
    else:
        raise ValueError('param["algo"] must be one of {"CW","greedy","greedy_lambda"}')

    local_search = int(param.get("local_search", 200))
    lam = 0.0
    lam1 = 0.0
    if algo_id == 2:
        lamv = param.get("lambda", [0.0, 0.0])
        lam = float(lamv[0]) if len(lamv) >= 1 else 0.0
        lam1 = float(lamv[1]) if len(lamv) >= 2 else 0.0

    return build_tsp_route(D, algo_id, local_search, lam, lam1)


# ==============================
# Minimal self-test (optional)
# ==============================

# -*- coding: utf-8 -*-
import numpy as np

def tsp_help():
    """
    TSP core (NumPy + Numba, no OOP, no Python lists in JIT code)

    Public Numba API
    ----------------
    build_tsp_route(D, algo_id, local_search, lam=2.4, lam1=1.7) -> np.ndarray[int64]
        - D: np.ndarray[float64] shape (n,n), ma trận khoảng cách đối xứng, D[i,i]=0.
        - algo_id:
            0: Clarke–Wright (CW), lam/lam1 bị bỏ qua.
            1: Greedy (cheapest insertion theo độ dài) + optional 2-opt.
            2: Greedy với lambda (biến dạng chi phí bằng lam, lam1) + optional 2-opt.
        - local_search:
            <=0: bỏ qua 2-opt;
            >0 : chạy tối đa 'local_search' lượt 2-opt.
        - lam, lam1: chỉ dùng khi algo_id=2.
        - Trả về: chu trình Hamilton (np.int64) độ dài n, quay sao cho route[0]==0.

    Wrapper không JIT
    -----------------
    build_tsp_route_from_param(D, param: dict) -> np.ndarray[int64]
        - param["algo"] ∈ {"CW","greedy","greedy_lambda",":greedy with lambda"}
        - param["local_search"] = int
        - param["lambda"] = [lam, lam1] (chỉ cho "greedy_lambda")

    Điều kiện đầu vào
    -----------------
    - D phải là ma trận (n,n) float64, không âm, đối xứng (khuyến nghị), và D[i,i]=0.
    - n >= 3 để ý nghĩa TSP; n < 3 sẽ trả về dãy [0], [0,1].
    - Không dùng Python list trong các hàm JIT; chỉ NumPy ndarray.

    Ví dụ nhanh (I/O)
    -----------------
    >>> rng = np.random.default_rng(0)
    >>> pts = rng.random((8, 2))
    >>> D = np.sqrt(((pts[:,None,:]-pts[None,:,:])**2).sum(axis=2))
    >>> route_cw = build_tsp_route(D, algo_id=0, local_search=100)
    >>> route_gr = build_tsp_route(D, algo_id=1, local_search=100)
    >>> route_gl = build_tsp_route(D, algo_id=2, local_search=100, lam=2.0, lam1=1.0)
    >>> int(route_cw[0]) == 0  # luôn start tại 0
    True

    Gợi ý hiệu năng
    ---------------
    - Biên dịch (JIT) lần đầu sẽ chậm; lần sau nhanh.
    - Dùng float64 nhất quán để tránh boxing.
    - 2-opt có độ phức tạp O(P * n^2) (P = số passes); cân bằng với chất lượng tour.

    Sai sót thường gặp
    ------------------
    - D không đối xứng hoặc chứa NaN/Inf -> hành vi không xác định.
    - D không phải float64 -> ép kiểu trước khi truyền.
    - algo_id=2 nhưng không set lam/lam1 (mặc định vẫn hợp lệ: lam=2.4, lam1=1.7).
    """
    text = tsp_help.__doc__
    if text is None:
        text = "No help available."
    print(text)
if __name__ == "__main__":
    rng = np.random.default_rng(0)
    pts = rng.random((40, 2))
    D = np.zeros((40, 40), dtype=np.float64)
    for i in range(40):
        for j in range(i+1, 40):
            d = np.sqrt(np.sum((pts[i]-pts[j])**2))
            D[i,j] = d; D[j,i] = d

    rCW = build_tsp_route_from_param(D, {"algo":"CW", "local_search":200})
    rGR = build_tsp_route_from_param(D, {"algo":"greedy", "local_search":200})
    rGL = build_tsp_route_from_param(D, {"algo":"greedy_lambda", "lambda":[2.0,1.0], "local_search":200})
    rCWNB = build_tsp_route(D, algo_id=0, local_search= 200)
    print("len CW :", float(tour_length_from_D(D, rCW)))
    print("len GR :", float(tour_length_from_D(D, rGR)))
    print("len GL :", float(tour_length_from_D(D, rGL)))
    print("len CW NB :", float(tour_length_from_D(D, rCWNB)))