import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
import math
import openjij as oj
from itertools import permutations

# ============================
# 預設距離矩陣：10-city TSP
# ============================
D = np.array([
    [0, 2, 8, 5, 7, 6, 3, 9, 4, 2],
    [2, 0, 5, 7, 3, 8, 4, 6, 9, 1],
    [8, 5, 0, 2, 6, 7, 3, 4, 1, 5],
    [5, 7, 2, 0, 4, 3, 8, 6, 2, 7],
    [7, 3, 6, 4, 0, 5, 9, 2, 7, 3],
    [6, 8, 7, 3, 5, 0, 4, 7, 9, 6],
    [3, 4, 3, 8, 9, 4, 0, 5, 6, 2],
    [9, 6, 4, 6, 2, 7, 5, 0, 8, 3],
    [4, 9, 1, 2, 7, 9, 6, 8, 0, 5],
    [2, 1, 5, 7, 3, 6, 2, 3, 5, 0]
])

N = D.shape[0]
penalty = 20

def idx(i, t, N):
    return i * N + t

# ============================
# KDE 函數
# ============================
def kde_1d(samples, num_points=200):
    if len(samples) == 0:
        return np.array([]), np.array([])
    
    xs = np.linspace(min(samples), max(samples), num_points)
    n = len(samples)
    if n < 2:
        return xs, np.zeros_like(xs)
    std = np.std(samples)
    if std == 0:
        std = 1.0
    h = 1.06 * std * (n ** (-1/5))
    if h == 0:
        h = 1.0
    ys = []
    inv_sqrt_2pi = 1.0 / math.sqrt(2.0 * math.pi)
    for x in xs:
        s = 0.0
        for xi in samples:
            z = (x - xi) / h
            s += math.exp(-0.5 * z * z) * inv_sqrt_2pi
        ys.append(s / (n * h))
    return xs, np.array(ys)

# ============================
# Route diversity matrix 函數
# ============================
def route_diversity_matrix(routes, N):
    freq = np.zeros((N, N), dtype=float)
    valid_routes = 0
    for route in routes:
        if len(route) == N and all(0 <= city < N for city in route):
            for t, city in enumerate(route):
                if 0 <= city < N:  # 確保city有效
                    freq[city, t] += 1
            valid_routes += 1
    if valid_routes > 0:
        freq /= valid_routes
    return freq

# ============================
# 建立 QUBO：位置編碼 + 固定起點 0
# ============================
def build_qubo(D, penalty, fix_start=True):
    N = D.shape[0]
    Q = {}

    # 距離項
    for t in range(N):
        t2 = (t + 1) % N
        for i in range(N):
            for j in range(N):
                if i != j:
                    u = idx(i, t, N)
                    v = idx(j, t2, N)
                    Q[(u, v)] = Q.get((u, v), 0) + D[i, j]

    # 約束A: 每個 time slot 有且只有一個城市
    for t in range(N):
        for i in range(N):
            u = idx(i, t, N)
            Q[(u, u)] = Q.get((u, u), 0) - penalty
            for j in range(i + 1, N):
                v = idx(j, t, N)
                Q[(u, v)] = Q.get((u, v), 0) + 2 * penalty

    # 約束B: 每個城市被訪問一次
    for i in range(N):
        for t in range(N):
            u = idx(i, t, N)
            Q[(u, u)] = Q.get((u, u), 0) - penalty
            for t2 in range(t + 1, N):
                v = idx(i, t2, N)
                Q[(u, v)] = Q.get((u, v), 0) + 2 * penalty

    # 固定起點：t=0 必須是 city 0
    if fix_start:
        big = 9999.0
        # 禁止 i!=0 在 t=0
        for i in range(1, N):
            u = idx(i, 0, N)
            Q[(u, u)] = Q.get((u, u), 0) + big
        # 獎勵 city 0 在 t=0
        u0 = idx(0, 0, N)
        Q[(u0, u0)] = Q.get((u0, u0), 0) - big

    return Q, N

# ============================
# Route decode & cost
# ============================
def decode_route(sample, N):
    route = []
    for t in range(N):
        city_for_t = None
        for i in range(N):
            if sample.get(idx(i, t, N), 0) == 1:
                city_for_t = i
                break
        if city_for_t is None:
            city_for_t = -1
        route.append(city_for_t)
    return route

def compute_cost(route, D):
    if -1 in route:
        return 99999  # 無效路徑懲罰
    N = len(route)
    total = 0
    for k in range(N):
        a = route[k]
        b = route[(k + 1) % N]
        total += D[a, b]
    return total

# ============================
# 兩台UAV TSP分配策略
# ============================
def split_route_for_two_uavs(route):
    """
    將單一TSP路徑分配給兩台UAV
    策略：交替分配城市，確保負載平衡
    """
    if len(route) == 0:
        return [], []
    
    # 移除無效城市
    valid_route = [city for city in route if city >= 0 and city < N]
    
    if len(valid_route) == 0:
        return [], []
    
    # 確保起點是0
    if valid_route[0] != 0:
        if 0 in valid_route:
            valid_route.remove(0)
            valid_route = [0] + valid_route
        else:
            valid_route = [0] + valid_route
    
    # 策略1: 奇偶分配
    uav1_path = []
    uav2_path = []
    
    for i, city in enumerate(valid_route):
        if i % 2 == 0:
            uav1_path.append(city)
        else:
            uav2_path.append(city)
    
    return uav1_path, uav2_path

def compute_two_uav_cost(uav1_path, uav2_path, D):
    """計算兩台UAV的總成本"""
    def path_cost(path):
        if len(path) <= 1:
            return 0
        cost = 0
        for i in range(len(path)):
            current = path[i]
            next_city = path[(i + 1) % len(path)]
            cost += D[current, next_city]
        return cost
    
    cost1 = path_cost(uav1_path)
    cost2 = path_cost(uav2_path)
    return cost1 + cost2, cost1, cost2

# ============================
# 經典 TSP 解法：暴力枚舉 (限制小規模)
# ============================
def classic_tsp_two_uav(D, max_permutations=1000):
    """
    經典TSP求解並分配給兩台UAV
    由於10城市的排列數太大(9!=362880)，限制搜索數量
    """
    cities = list(range(1, len(D)))  # 排除起點0
    best_total_cost = float('inf')
    best_route = None
    best_uav1 = None
    best_uav2 = None
    
    count = 0
    for perm in permutations(cities):
        if count >= max_permutations:
            break
        count += 1
        
        route = [0] + list(perm)  # 固定0為起點
        
        # 分配給兩台UAV
        uav1_path, uav2_path = split_route_for_two_uavs(route)
        total_cost, cost1, cost2 = compute_two_uav_cost(uav1_path, uav2_path, D)
        
        if total_cost < best_total_cost:
            best_total_cost = total_cost
            best_route = route
            best_uav1 = uav1_path
            best_uav2 = uav2_path

    return best_uav1, best_uav2, best_total_cost, best_route

# ============================
# 跑 SA or SQA 多次
# ============================
def run_algorithm(name, sampler, Q, D, N, num_runs=20, num_reads=20):
    all_costs = []
    all_routes = []
    all_uav1_paths = []
    all_uav2_paths = []
    all_uav_costs = []

    print(f"\n=== Running {name} ===")
    for r in range(num_runs):
        result = sampler.sample_qubo(Q, num_reads=num_reads)
        best = result.first.sample
        route = decode_route(best, N)
        cost = compute_cost(route, D)
        
        # 分配給兩台UAV
        uav1_path, uav2_path = split_route_for_two_uavs(route)
        total_uav_cost, cost1, cost2 = compute_two_uav_cost(uav1_path, uav2_path, D)
        
        all_costs.append(cost)
        all_routes.append(tuple(route))
        all_uav1_paths.append(tuple(uav1_path))
        all_uav2_paths.append(tuple(uav2_path))
        all_uav_costs.append(total_uav_cost)
        
        print(f"{name} Run {r+1:02d}: Route={route}")
        print(f"  TSP Cost={cost}, UAV1={uav1_path}(cost={cost1}), UAV2={uav2_path}(cost={cost2}), Total={total_uav_cost}")

    return all_costs, all_routes, all_uav1_paths, all_uav2_paths, all_uav_costs

# ============================
# 創建綜合結果圖表
# ============================
def create_comprehensive_charts(sa_results, sqa_results, classic_results):
    """創建綜合分析圖表"""
    fig = plt.figure(figsize=(18, 12))
    
    # 解包結果
    sa_costs, sa_routes, sa_uav1, sa_uav2, sa_uav_costs = sa_results
    sqa_costs, sqa_routes, sqa_uav1, sqa_uav2, sqa_uav_costs = sqa_results
    classic_uav1, classic_uav2, classic_cost, classic_route = classic_results
    
    # 1. TSP成本分布比較 (左上)
    ax1 = plt.subplot(2, 3, 1)
    bins = range(min(sa_costs + sqa_costs), max(sa_costs + sqa_costs) + 2)
    ax1.hist(sa_costs, bins=bins, alpha=0.6, label="SA", density=True, color='lightblue')
    ax1.hist(sqa_costs, bins=bins, alpha=0.6, label="SQA", density=True, color='lightcoral')
    
    # 添加KDE
    if len(sa_costs) > 1:
        xs_sa, ys_sa = kde_1d(sa_costs)
        ax1.plot(xs_sa, ys_sa, label="SA KDE", color='blue', linewidth=2)
    if len(sqa_costs) > 1:
        xs_sqa, ys_sqa = kde_1d(sqa_costs)
        ax1.plot(xs_sqa, ys_sqa, label="SQA KDE", color='red', linewidth=2)
    
    ax1.axvline(x=classic_cost, color='green', linestyle='--', 
                label=f'Classic: {classic_cost}', linewidth=2)
    ax1.set_title("TSP Cost Distribution")
    ax1.set_xlabel("Tour Cost")
    ax1.set_ylabel("Density")
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 2. UAV總成本分布比較 (右上)
    ax2 = plt.subplot(2, 3, 2)
    uav_bins = range(min(sa_uav_costs + sqa_uav_costs), max(sa_uav_costs + sqa_uav_costs) + 2)
    ax2.hist(sa_uav_costs, bins=uav_bins, alpha=0.6, label="SA UAV", density=True, color='lightblue')
    ax2.hist(sqa_uav_costs, bins=uav_bins, alpha=0.6, label="SQA UAV", density=True, color='lightcoral')
    
    classic_uav_total, _, _ = compute_two_uav_cost(classic_uav1, classic_uav2, D)
    ax2.axvline(x=classic_uav_total, color='green', linestyle='--', 
                label=f'Classic UAV: {classic_uav_total}', linewidth=2)
    ax2.set_title("Two-UAV Total Cost Distribution")
    ax2.set_xlabel("Total UAV Cost")
    ax2.set_ylabel("Density")
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    # 3. 箱型圖比較 (中上)
    ax3 = plt.subplot(2, 3, 3)
    box_data = [sa_costs, sqa_costs, sa_uav_costs, sqa_uav_costs]
    bp = ax3.boxplot(box_data, labels=["SA\nTSP", "SQA\nTSP", "SA\nUAV", "SQA\nUAV"], 
                     patch_artist=True, showmeans=True)
    
    colors = ['lightblue', 'lightcoral', 'skyblue', 'salmon']
    for patch, color in zip(bp['boxes'], colors):
        patch.set_facecolor(color)
    
    ax3.set_title("Cost Comparison Box Plot")
    ax3.set_ylabel("Cost")
    ax3.grid(True, alpha=0.3)
    
    # 4. Route diversity heatmap - SA (左下)
    ax4 = plt.subplot(2, 3, 4)
    sa_mat = route_diversity_matrix(sa_routes, N)
    im1 = ax4.imshow(sa_mat, aspect='auto', origin='lower', cmap='viridis')
    ax4.set_title("SA Route Diversity")
    ax4.set_xlabel("Position")
    ax4.set_ylabel("City")
    plt.colorbar(im1, ax=ax4, shrink=0.8)
    
    # 5. Route diversity heatmap - SQA (右下)
    ax5 = plt.subplot(2, 3, 5)
    sqa_mat = route_diversity_matrix(sqa_routes, N)
    im2 = ax5.imshow(sqa_mat, aspect='auto', origin='lower', cmap='viridis')
    ax5.set_title("SQA Route Diversity")
    ax5.set_xlabel("Position")
    ax5.set_ylabel("City")
    plt.colorbar(im2, ax=ax5, shrink=0.8)
    
    # 6. 最佳路徑可視化 (中下)
    ax6 = plt.subplot(2, 3, 6)
    
    # 找出最佳量子解
    best_sa_idx = sa_costs.index(min(sa_costs))
    best_sqa_idx = sqa_costs.index(min(sqa_costs))
    best_sa_route = list(sa_routes[best_sa_idx])
    best_sqa_route = list(sqa_routes[best_sqa_idx])
    
    # 簡單的圓形佈局
    angles = np.linspace(0, 2*np.pi, N, endpoint=False)
    x_pos = np.cos(angles)
    y_pos = np.sin(angles)
    
    # 繪製城市
    ax6.scatter(x_pos, y_pos, s=200, c='red', zorder=5)
    for i, (x, y) in enumerate(zip(x_pos, y_pos)):
        ax6.annotate(str(i), (x, y), ha='center', va='center', 
                    fontsize=10, fontweight='bold', color='white', zorder=6)
    
    # 選擇更好的路徑來顯示
    if min(sa_costs) <= min(sqa_costs):
        display_route = best_sa_route
        route_name = "SA"
        route_cost = min(sa_costs)
    else:
        display_route = best_sqa_route
        route_name = "SQA"
        route_cost = min(sqa_costs)
    
    # 繪製路徑
    if all(0 <= city < N for city in display_route):
        route_x = [x_pos[city] for city in display_route] + [x_pos[display_route[0]]]
        route_y = [y_pos[city] for city in display_route] + [y_pos[display_route[0]]]
        ax6.plot(route_x, route_y, 'b-', linewidth=2, alpha=0.8, zorder=3)
    
    ax6.set_title(f"Best {route_name} Route\nCost: {route_cost}")
    ax6.set_xlim(-1.3, 1.3)
    ax6.set_ylim(-1.3, 1.3)
    ax6.set_aspect('equal')
    ax6.axis('off')
    
    plt.tight_layout()
    plt.savefig("tsp_two_uav_comprehensive_analysis.png", dpi=300, bbox_inches='tight')
    print(f"\n📊 綜合分析圖表已保存: tsp_two_uav_comprehensive_analysis.png")
    plt.show()

# ============================
# Main
# ============================
if __name__ == "__main__":
    print("🚀 開始10城市兩台UAV TSP分析...")
    print(f"Distance Matrix ({N}x{N}):")
    print(D)
    
    Q, N = build_qubo(D, penalty, fix_start=True)

    # Samplers
    sqa_sampler = oj.SQASampler()
    sa_sampler = oj.SASampler()

    # 參數設定
    NUM_RUNS = 15
    NUM_READS = 20

    # 運行量子退火算法
    print("🔥 Running Quantum Annealing Algorithms...")
    sa_results = run_algorithm("SA", sa_sampler, Q, D, N, NUM_RUNS, NUM_READS)
    sqa_results = run_algorithm("SQA", sqa_sampler, Q, D, N, NUM_RUNS, NUM_READS)

    # 運行經典TSP解法
    print("\n🎯 Running Classic TSP (limited search)...")
    classic_results = classic_tsp_two_uav(D, max_permutations=5000)
    classic_uav1, classic_uav2, classic_total_cost, classic_route = classic_results
    
    print(f"Classic TSP Results:")
    print(f"  Best Route: {classic_route}")
    print(f"  UAV1 Path: {classic_uav1}")
    print(f"  UAV2 Path: {classic_uav2}")
    print(f"  Total Cost: {classic_total_cost}")

    # 統計分析
    sa_costs, sa_routes, sa_uav1, sa_uav2, sa_uav_costs = sa_results
    sqa_costs, sqa_routes, sqa_uav1, sqa_uav2, sqa_uav_costs = sqa_results

    print(f"\n📊 Algorithm Performance Summary:")
    print(f"{'Algorithm':<15} {'TSP Cost':<20} {'UAV Total Cost':<20}")
    print(f"{'-'*55}")
    print(f"{'SA':<15} min={min(sa_costs):<3} mean={np.mean(sa_costs):<6.1f} "
          f"min={min(sa_uav_costs):<3} mean={np.mean(sa_uav_costs):<6.1f}")
    print(f"{'SQA':<15} min={min(sqa_costs):<3} mean={np.mean(sqa_costs):<6.1f} "
          f"min={min(sqa_uav_costs):<3} mean={np.mean(sqa_uav_costs):<6.1f}")
    print(f"{'Classic':<15} {classic_total_cost:<23} {classic_total_cost:<20}")

    # 創建綜合圖表
    print("\n🎨 Creating comprehensive analysis charts...")
    create_comprehensive_charts(sa_results, sqa_results, classic_results)

    # 最終比較
    best_sa_uav = min(sa_uav_costs)
    best_sqa_uav = min(sqa_uav_costs)
    
    print(f"\n🏆 Final Two-UAV TSP Results:")
    print(f"SA best UAV cost: {best_sa_uav}")
    print(f"SQA best UAV cost: {best_sqa_uav}")
    print(f"Classic UAV cost: {classic_total_cost}")
    
    if best_sa_uav <= best_sqa_uav and best_sa_uav <= classic_total_cost:
        print("🥇 Winner: SA")
    elif best_sqa_uav <= best_sa_uav and best_sqa_uav <= classic_total_cost:
        print("🥇 Winner: SQA")
    else:
        print("🥇 Winner: Classic Algorithm")