import os
import re
from ortools.sat.python import cp_model

def solve_with_or_tools(items, capacity, timeout_seconds=60.0):
    """
    使用 Google OR-Tools 的 CP-SAT 求解器来寻找装箱问题的最优解.

    参数:
        items (list[int]): 物品大小的列表.
        capacity (int): 每个箱子的容量.
        timeout_seconds (float): 求解器的最大运行时间(秒).

    返回:
        一个包含最终箱子分配方案的列表, 或者在未找到解时返回 None.
    """
    num_items = len(items)
    # 理论上箱子数量的上限是物品的数量
    num_bins = num_items

    # 创建 CP-SAT 模型
    model = cp_model.CpModel()

    # --- 变量定义 ---
    # x[i][j] 是一个布尔变量, 如果物品 i 被放入箱子 j, 则为 1, 否则为 0
    x = {}
    for i in range(num_items):
        for j in range(num_bins):
            x[i, j] = model.NewBoolVar(f"x_{i}_{j}")

    # y[j] 是一个布尔变量, 如果箱子 j 被使用, 则为 1, 否则为 0
    y = [model.NewBoolVar(f"y_{j}") for j in range(num_bins)]

    # --- 约束条件 ---
    # 1. 每个物品必须且只能被放入一个箱子.
    for i in range(num_items):
        model.AddExactlyOne([x[i, j] for j in range(num_bins)])

    # 2. 对于每个箱子, 其内部所有物品的总大小不能超过其容量.
    #    如果箱子 j 未被使用 (y[j] == 0), 那么它的容量约束为 0.
    for j in range(num_bins):
        model.Add(sum(items[i] * x[i, j] for i in range(num_items)) <= capacity * y[j])

    # --- 优化目标 ---
    # 目标是最小化被使用的箱子总数.
    model.Minimize(sum(y))

    # --- 求解 ---
    solver = cp_model.CpSolver()
    solver.parameters.max_time_in_seconds = timeout_seconds
    status = solver.Solve(model)

    # --- 返回结果 ---
    if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
        solution = []
        for j in range(num_bins):
            # 如果箱子被使用
            if solver.Value(y[j]) == 1:
                bin_items = [i for i in range(num_items) if solver.Value(x[i, j]) == 1]
                solution.append(bin_items)
        return solution
    else:
        return None


def process_bpp_folder(folder_path):
    """
    遍历指定文件夹中的所有 .txt 文件, 使用 OR-Tools 求解并将最优解写回文件.
    此函数会覆盖文件中已有的旧解.
    """
    print(f"--- 开始处理文件夹: {folder_path} ---")
    for filename in sorted(os.listdir(folder_path)):
        if not filename.endswith(".txt"):
            continue
        
        filepath = os.path.join(folder_path, filename)
        print(f"处理文件: {filename}...")

        try:
            # --- 1. 读取和解析文件 ---
            with open(filepath, "r", encoding='utf-8') as f:
                full_content = f.read()
            
            # 从完整内容中分离出数据部分, 忽略可能存在的旧解部分
            data_content = re.split(r"^\s*===", full_content, flags=re.MULTILINE)[0]
            
            # 【关键修正】从文件内容解析数据
            numeric_lines = [line.strip() for line in data_content.splitlines() if line.strip().isdigit()]

            if len(numeric_lines) < 3:
                print(f"  -> 警告: 在 '{filename}' 中未找到足够的数字行 (至少3行), 已跳过.")
                continue
            
            # 第一行是物品数量 (求解器不需要, 但格式要求)
            # item_count = int(numeric_lines[0])
            # 第二行是容量
            capacity = int(numeric_lines[1])
            # 从第三行开始是物品列表
            items = [int(n) for n in numeric_lines[2:]]

            if not items:
                print(f"  -> 警告: 在 '{filename}' 中未找到物品数据, 已跳过.")
                continue

            # --- 2. 使用 OR-Tools 求解 ---
            print(f"  -> 正在求解... (容量: {capacity}, 物品数: {len(items)})")
            solution = solve_with_or_tools(items, capacity)

            # --- 3. 构建解的文本并写回文件 ---
            solution_str = f"\n\n=== Bin Packing Solution (OR-Tools) ===\n"
            solution_str += f"Bin capacity: {capacity}\n"
            if solution is None:
                solution_str += "No solution found.\n"
                print(f"  -> 未找到解.")
            else:
                solution_str += f"Total bins used: {len(solution)}\n"
                for i, bin_indices in enumerate(solution, 1):
                    bin_sum = sum(items[item_idx] for item_idx in bin_indices)
                    item_values = ", ".join(str(items[item_idx]) for item_idx in bin_indices)
                    solution_str += f"Bin {i} (sum={bin_sum}): {item_values}\n"
                print(f"  -> 求解完成, 使用了 {len(solution)} 个箱子.")

            # 使用 'w' 模式写入, 这样会用纯数据+新解覆盖整个文件
            with open(filepath, "w", encoding='utf-8') as f:
                f.write(data_content.strip() + solution_str)

        except Exception as e:
            print(f"  -> 处理 '{filename}' 时发生错误: {e}")

    print("--- 所有文件处理完毕 ---")


if __name__ == "__main__":
    # !! 重要: 请将此路径修改为您存放 BPP 数据文件的实际文件夹路径 !!
    INPUT_FOLDER = '/home/liuyihong/openevolve/examples/BPP/input'
    
    if os.path.isdir(INPUT_FOLDER):
        process_bpp_folder(INPUT_FOLDER)
    else:
        print(f"错误: 文件夹 '{INPUT_FOLDER}' 不存在, 请检查路径.")

