import itertools

import numpy as np
import matplotlib.pyplot as plt
from pyqubo import Array, Constraint, Placeholder

import vrplib

instance_name = "CMT1"
path_instance = "cvrp.vrp"
path_solution = "cvrp.sol"

vrplib.download_instance(instance_name, path_instance)
vrplib.download_solution(instance_name, path_solution)

vrp_instance = vrplib.read_instance(path_instance)
vrp_solution = vrplib.read_solution(path_solution)

customers = vrp_instance["node_coord"]
customers_nodepot = customers[1:]  # 0: depot
cus_num = len(customers_nodepot)

# 距離行列の生成
D_uv = np.zeros((cus_num, cus_num))
for u in range(cus_num):
    for v in range(u + 1, cus_num):
        distance_uv = np.linalg.norm(customers_nodepot[u] - customers_nodepot[v])
        D_uv[u][v] = D_uv[v][u] = distance_uv

D_uv_normed = D_uv / np.max(D_uv)
print(D_uv_normed)

demands = vrp_instance["demand"][1:]  # 0: depot
print(demands)

# 車両の設定
car_num = 5  # 車の台数
capa = vrp_instance["capacity"]  # トラックの容量

def make_pyqubo_model1(car_num, cus_num, capa, D_uv, demands):
    x = Array.create("x", shape=(cus_num, car_num), vartype="BINARY")
    y = Array.create("y", shape=(capa, car_num), vartype="BINARY")

    H_A_1 = np.sum(
        [
            (1 - np.sum([y[i - 1, k] for i in range(1, capa + 1)])) ** 2
            for k in range(car_num)
        ]
    )
    H_A_2 = np.sum(
        [
            (
                np.sum([i * y[i - 1, k] for i in range(1, capa + 1)])
                - np.sum([demands[v] * x[v, k] for v in range(cus_num)])
            )
            ** 2
            for k in range(car_num)
        ]
    )
    H_B = np.sum(
        [(1 - np.sum([x[v, k] for k in range(car_num)])) ** 2 for v in range(cus_num)]
    )
    H_C = np.sum(
        [
            np.sum(
                [
                    D_uv[u, v] * x[u, k] * x[v, k]
                    for u, v in itertools.combinations(range(cus_num), 2)
                ]
            )
            for k in range(car_num)
        ]
    )

    H = Placeholder("X") * Constraint(H_A_1, "H_A_1") + Placeholder("A") * Constraint(
        H_A_2, "H_A_2"
    )
    H = H + Placeholder("B") * Constraint(H_B, "H_B") + Placeholder("C") * H_C

    return H.compile()

model1 = make_pyqubo_model1(car_num, cus_num, capa, D_uv, demands)

feed_dict1 = dict(X=2500.0, A=50.0, B=2500.0, C=1.0)  # いろいろ係数を調整してみてください
bqm1 = model1.to_bqm(feed_dict=feed_dict1)
print("num_variables:", bqm1.num_variables)

# SAの場合
from neal import SimulatedAnnealingSampler

sampler = SimulatedAnnealingSampler()

# D-Waveマシンの場合
# from dwave.system import DWaveSampler, EmbeddingComposite
# sampler_config = dict(solver="Advantage_system6.2", token="YOUR_TOKEN")
# sampler = EmbeddingComposite(DWaveSampler(**sampler_config))

sampler_params1 = dict(num_reads=10, num_sweeps=10000)
sampleset1 = sampler.sample(bqm1, **sampler_params1).aggregate()
sampleset1.to_pandas_dataframe().head()

# 制約を満たしているかどうかを確認する
decoded_samples1 = model1.decode_sampleset(sampleset1, feed_dict1)

for sample in decoded_samples1:
    print(sample.energy)
    print(sample.constraints(only_broken=True))

# 顧客の要求量に関して正規化する
demands_normed = demands / capa
print(demands_normed)

def make_pyqubo_model2(car_num, cus_num, capa, D_uv_normed, demands):
    x = Array.create("x", shape=(cus_num, car_num), vartype="BINARY")

    H_B = np.sum(
        [(1 - np.sum([x[v, k] for k in range(car_num)])) ** 2 for v in range(cus_num)]
    )  # one-hot
    H_C = np.sum(
        [
            np.sum(
                [
                    D_uv_normed[u, v] * x[u, k] * x[v, k]
                    for u, v in itertools.combinations(range(cus_num), 2)
                ]
            )
            for k in range(car_num)
        ]
    )  # 距離
    H_D = np.sum(
        [
            (
                np.sum(demands) / car_num
                - np.sum([demands[v] * x[v, k] for v in range(cus_num)])
            )
            ** 2
            for k in range(car_num)
        ]
    )  # balance

    H = (
        Placeholder("B") * Constraint(H_B, "H_B")
        + Placeholder("C") * H_C
        + Placeholder("D") * H_D
    )
    return H.compile()

model2 = make_pyqubo_model2(car_num, cus_num, capa, D_uv_normed, demands_normed)

feed_dict2 = dict(B=3, C=1, D=100)
bqm2 = model2.to_bqm(feed_dict=feed_dict2)
print("num_variables:", bqm2.num_variables)

# SAの場合
sampler = SimulatedAnnealingSampler()

# D-Waveマシンの場合 サイズ50だと大きいため解けない
# from dwave.system import DWaveSampler, EmbeddingComposite
# sampler_config = dict(solver="Advantage_system6.2", token="")
# sampler = EmbeddingComposite(DWaveSampler(**sampler_config))

sampler_params2 = dict(num_reads=10, num_sweeps=10000)
sampleset2 = sampler.sample(bqm2, **sampler_params2).aggregate()
sampleset1.to_pandas_dataframe().head()

decoded_samples2 = model2.decode_sampleset(sampleset2, feed_dict2)
for sample in decoded_samples2:
    print(sample.energy)
    print(sample.constraints(only_broken=True))

def sample_to_assignment(sample, car_num, cus_num):
    assignment = []
    for m in range(cus_num):
        for n in range(car_num):
            if sample[f"x[{m}][{n}]"] == 1:
                assignment.append((m, n))
    return assignment

sample2 = decoded_samples2[0]
assignment2 = sample_to_assignment(sample2.sample, car_num, cus_num)
print(assignment2)

# 顧客がどのクラスターに属するかを示すリストを作成
belonged_cluster = [j for _, j in assignment2]
print(belonged_cluster)

def weight_check(belonged_cluster, demands, capa):
    weights = [0] * len(set(belonged_cluster))
    for i, cluster_i in enumerate(belonged_cluster):
        weights[cluster_i] += int(demands[i])

    for i, weight_i in enumerate(weights):
        if weight_i > capa:
            print(f"キャパオーバーです。(超過:{weight_i - capa})")
    print("容量制約を満たしています。")


weight_check(belonged_cluster, demands, capa)

# クラスタリングの結果を出力
plt.scatter(customers_nodepot[:, 0], customers_nodepot[:, 1], c=belonged_cluster)
plt.show()

# depotが入っている状態での顧客番号のリストを作成
clusters = [[0] for i in range(car_num)]  # depotの番号は初めから追加しておく
for i in range(cus_num):
    clusters[belonged_cluster[i]].append(i + 1)  # depotが含まれている状態に戻すために数字を+1する
print(clusters)

# 同一クラスターにおける顧客どうしの距離を計算する
distance_matrices = []
for i in range(car_num):
    cluster_size = len(clusters[i])
    distance_matrix = np.zeros((cluster_size, cluster_size))
    for j in range(cluster_size):
        for k in range(j + 1, cluster_size):
            distance_jk = np.linalg.norm(
                customers[clusters[i][j]] - customers[clusters[i][k]]
            )
            distance_matrix[j][k] = distance_matrix[k][j] = distance_jk

    distance_matrices.append(distance_matrix.tolist())

# ortoolsを利用して巡回セールスマン問題を解く
from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp


def print_solution(manager, routing, solution, j):
    index = routing.Start(0)
    plan_output = f"Route for vehicle {j}:\n"
    route_distance = 0
    while not routing.IsEnd(index):
        plan_output += " {} ->".format(manager.IndexToNode(index))
        previous_index = index
        index = solution.Value(routing.NextVar(index))
        route_distance += routing.GetArcCostForVehicle(previous_index, index, 0)
    plan_output += " {}\n".format(manager.IndexToNode(index))
    plan_output += "Route distance: {}miles\n".format(route_distance)
    print(plan_output)
    return route_distance


def get_routes(solution, routing, manager):
    routes = []
    for route_nbr in range(routing.vehicles()):
        index = routing.Start(route_nbr)
        route = [manager.IndexToNode(index)]
        while not routing.IsEnd(index):
            index = solution.Value(routing.NextVar(index))
            route.append(manager.IndexToNode(index))
        routes.append(route)
    return routes

solutions = []
route_dis = []

for j in range(len(clusters)):
    distance_matrix = distance_matrices[j]

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

    def distance_callback(from_index, to_index):
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return int(distance_matrix[from_node][to_node])  # ortoolsの仕様でintを返す

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

    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )

    solution = routing.SolveWithParameters(search_parameters)

    if solution is None:
        print("No solution found !")
        continue

    # 解を表示する
    route_dis.append(print_solution(manager, routing, solution, j + 1))

    routes = get_routes(solution, routing, manager)
    for i, route in enumerate(routes):
        solutions.append(route)

import networkx as nx


# 解を視覚的に描画する関数
def print_routes(car_num, solutions, belonged_cluster, customers):
    # 元のデータ番号での呼び出し
    solutions2 = [[] for i in range(car_num)]
    for i in range(len(solutions)):
        for j in range(len(solutions[i])):
            solutions2[i].append(clusters[i][solutions[i][j]])

    # 各顧客について色を指定する
    cm = plt.get_cmap("tab20")
    node_list = [(0, {"color": cm(0)})]
    for i in range(len(belonged_cluster)):
        node_list.append((i + 1, {"color": cm(belonged_cluster[i] + 1)}))

    # depotおよび顧客の番号と位置に関する辞書を作成
    pos = {}
    for i in range(len(customers)):
        pos[i] = (customers[i][0], customers[i][1])

    # ルーティングで得られた解をもとに実際に通る辺のリストを作成
    edge_list = []
    for i in range(len(solutions2)):
        for j in range(len(solutions2[i]) - 1):
            edge_list.append((solutions2[i][j], solutions2[i][j + 1]))

    g = nx.DiGraph()
    g.add_nodes_from(node_list)
    g.add_edges_from(edge_list)

    node_color = [node["color"] for node in g.nodes.values()]
    plt.figure(figsize=(8, 8))  # 見やすいようにサイズを変更してください
    nx.draw_networkx(g, pos, with_labels=True, node_color=node_color)
    plt.show()


print_routes(car_num, solutions, belonged_cluster, customers)

# 各ルートの距離を計算する
distance_sum = []
for i in range(len(solutions)):
    sum = 0
    for j in range(len(solutions[i]) - 1):
        sum += distance_matrices[i][solutions[i][j]][solutions[i][j + 1]]
    distance_sum.append(sum)

for i in range(len(distance_sum)):
    print(f"{i+1}番目のルートの距離の総和: {distance_sum[i]}")
print(np.sum(distance_sum))

bks = vrp_solution["cost"]  # 厳密解
ans = np.sum(distance_sum)
print((ans - bks) / bks * 100, "%")