import numpy as np
from ga import GA


def if_valid_partition(row_boundaries, col_boundaries, params):
    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']

    # 根据分割边界生成所有矩形任务
    rectangles = []
    for i in range(len(row_boundaries) - 1):
        for j in range(len(col_boundaries) - 1):
            r1 = row_boundaries[i]
            r2 = row_boundaries[i + 1]
            c1 = col_boundaries[j]
            c2 = col_boundaries[j + 1]
            d = (r2 - r1) * H * (c2 - c1) * W  # 任务的照片数量（矩形面积）

            # 求解rho
            rho_time_limit = (flight_time_factor - trans_time_factor) / \
                (comp_time_factor - trans_time_factor)
            rho_energy_limit = (battery_energy_capacity - flight_energy_factor * d - trans_energy_factor * d) / \
                (comp_energy_factor * d - trans_energy_factor * d)
            if rho_energy_limit < 0:
                return []

            rho = min(rho_time_limit, rho_energy_limit)
            flight_time = flight_time_factor * d
            comp_time = comp_time_factor * rho * d
            trans_time = trans_time_factor * (1 - rho) * d
            bs_time = bs_time_factor * (1 - rho) * d

            # 计算任务矩形中心，用于后续车辆移动时间计算
            center_r = (r1 + r2) / 2.0 * H
            center_c = (c1 + c2) / 2.0 * W

            rectangles.append({
                # 'r1': r1, 'r2': r2, 'c1': c1, 'c2': c2,
                'd': d,
                'rho': rho,
                'flight_time': flight_time,
                'comp_time': comp_time,
                'trans_time': trans_time,
                'bs_time': bs_time,
                'center': (center_r, center_c)
            })
    return rectangles


def GA_solver(rectangles, params):
    num_city = len(rectangles) + 1    # 划分好的区域中心点+整个区域的中心
    k = params['num_cars']

    # 初始化坐标 (第一个点是整个区域的中心)
    center_data = [[params['H'] / 2.0, params['W'] / 2.0]]
    for rec in rectangles:
        center_data.append(rec['center'])
    center_data = np.array(center_data)

    # 关键：有k架无人机，则再增加N-1个`点` (坐标是起始点)，这些点之间的距离是inf
    for d in range(k - 1):
        center_data = np.vstack([center_data, center_data[0]])
        num_city += 1  # 增加欺骗城市

    to_process_idx = [0]
    # print("start point:", location[0])
    for d in range(1, k):  # 1, ... drone-1
        # print("added base point:", location[num_city - d])
        to_process_idx.append(num_city - d)

    model = GA(num_drones=k, num_city=num_city, num_total=20, data=center_data.copy(
    ), to_process_idx=to_process_idx, rectangles=rectangles)
    Best_path, Best = model.run()

    # 根据最佳路径计算各系统任务分配
    if Best_path[0] not in to_process_idx:
        Best_path.insert(0, 0)

    if Best_path[-1] not in to_process_idx:
        Best_path.append(0)
    
    return Best_path, Best, to_process_idx