import random
import math
import yaml
import numpy as np
from utils import if_valid_partition, GA_solver
from itertools import product, combinations
import json
from tqdm import tqdm

np.random.seed(42)
random.seed(42)
best_T = float('inf')
best_solution = None
best_row_boundaries = None
best_col_boundaries = None


# ---------------------------
# 需要修改的超参数
# ---------------------------
params_file = 'params_50_50_3'


with open(params_file + '.yml', 'r', encoding='utf-8') as file:
    params = yaml.safe_load(file)

H = params['H']
W = params['W']
k = params['num_cars']

flight_time_factor = params['flight_time_factor']
comp_time_factor = params['comp_time_factor']
trans_time_factor = params['trans_time_factor']
car_time_factor = params['car_time_factor']
bs_time_factor = params['bs_time_factor']

flight_energy_factor = params['flight_energy_factor']
comp_energy_factor = params['comp_energy_factor']
trans_energy_factor = params['trans_energy_factor']
battery_energy_capacity = params['battery_energy_capacity']

# # 定义数字列表
# numbers = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]

row_cuts_set = [[0.3, 0.48, 0.77]]
col_cuts_set = [[0.5]]

for row_cuts in row_cuts_set:
    for col_cuts in col_cuts_set:
        row_boundaries = [0.0] + list(row_cuts) + [1.0]
        col_boundaries = [0.0] + list(col_cuts) + [1.0]

        # 这里面的距离不再是比例，而是真实距离！
        rectrangles = if_valid_partition(
            row_boundaries, col_boundaries, params)
        if not rectrangles:
            continue
        else:
            # 使用遗传算法求出每一种网格划分的可行解，然后选择其中的最优解
            current_solution, current_time, to_process_idx = GA_solver(
                rectrangles, params)

            if current_time < best_T:
                best_T = current_time
                best_solution = current_solution
                best_row_boundaries = row_boundaries
                best_col_boundaries = col_boundaries

                # 将best_solution分解成每个车队的路径
                found_start_points_indices = []
                for i in range(len(best_solution)):
                    if best_solution[i] in to_process_idx:
                        found_start_points_indices.append(i)
                car_paths = []
                for j in range(len(found_start_points_indices) - 1):
                    from_index = found_start_points_indices[j]
                    end_index = found_start_points_indices[j + 1]
                    car_path = []
                    for k in range(from_index, end_index + 1):
                        rectrangle_idx = best_solution[k]
                        if rectrangle_idx not in to_process_idx:
                            car_path.append(rectrangle_idx - 1)
                    if car_path:
                        car_paths.append(car_path)

# 输出最佳方案
print("Best solution:", best_solution)
print("Time:", best_T)
print("Row boundaries:", best_row_boundaries)
print("Col boundaries:", best_col_boundaries)
print("Car Paths:", car_paths)