from ortools.linear_solver import pywraplp

def solve_program(preference_mat:list, 
                  want_num_array:list, 
                  is_new_array:list, 
                  is_tech_array:list,
                  is_hr_array:list,
                  num_min=None,
                  num_max=None,
                  num_tech_min=None,
                  num_tech_max=None,
                  num_old_min=None,
                  num_old_max=None,
                  num_hr_min=None,
                  num_hr_max=None,
                  num_new_min=None,
                  num_new_max=None
                  ):

    is_old_array = [not is_new for is_new in is_new_array]
    
    # 使用 SCIP 求解器求解组合优化问题
    solver = pywraplp.Solver.CreateSolver("SCIP")
    if not solver:
        return

    N = len(preference_mat)     # 学生人数
    M = len(preference_mat[0])  # 班次数
    avg_num = sum(want_num_array) / M
    print(f"平均人数：{avg_num}")

    # 创建决策变量：variables[i][j] 表示第i个学生是否选择第j个班次
    variables = [[solver.IntVar(0.0, 1.0, f"choice_{i}_{j}") 
                  for j in range(M)] for i in range(N)]
    
    # 创建辅助变量用于优化目标（均衡班次人数）
    aux_vars = [solver.NumVar(0.0, solver.infinity(), f"aux_{j}") 
                for j in range(M)]
    print("Number of variables =", solver.NumVariables())

    # 约束：每个同学意愿一定要满足
    for i in range(N):
        for j in range(M):
            solver.Add(variables[i][j] <= preference_mat[i][j])

    # 约束：辅助变量用于计算与平均人数的偏差（绝对值）
    for j in range(M):
        actual_num = sum(variables[i][j] for i in range(N))
        solver.Add(aux_vars[j] >= actual_num - avg_num)
        solver.Add(aux_vars[j] >= avg_num - actual_num)

    # 约束：满足每位同学的意愿班次数
    for i in range(N):
        total_shifts = sum(variables[i])
        if want_num_array[i] == 3:
            solver.Add(total_shifts >= 2)
            solver.Add(total_shifts <= 3)
        else:
            solver.Add(total_shifts == want_num_array[i])

    # 添加班次人数约束的辅助函数
    def add_shift_constraint(array, min_val, max_val):
        for j in range(M):
            shift_count = sum(variables[i][j] * array[i] for i in range(N))
            if min_val is not None:
                solver.Add(shift_count >= min_val)
            if max_val is not None:
                solver.Add(shift_count <= max_val)

    # 约束：每个班次的总人数
    add_shift_constraint([1] * N, num_min, num_max)
    
    # 约束：每班次电脑或电器的老人数量
    tech_old_array = [is_old_array[i] * is_tech_array[i] for i in range(N)]
    add_shift_constraint(tech_old_array, num_tech_min, num_tech_max)
    
    # 约束：每班次老人数量
    add_shift_constraint(is_old_array, num_old_min, num_old_max)
    
    # 约束：每班次人资部小朋友数量
    add_shift_constraint(is_hr_array, num_hr_min, num_hr_max)
    
    # 约束：每班次小朋友数量
    add_shift_constraint(is_new_array, num_new_min, num_new_max)

    # 优化目标：最小化各班次人数与平均人数的偏差总和
    solver.Minimize(sum(aux_vars))

    # 求解
    status = solver.Solve()

    if status == pywraplp.Solver.OPTIMAL:
        print("Optimal solution found:")
        variables_return = [[variables[i][j].solution_value() for j in range(M)] 
                           for i in range(N)]
        aux_vars_return = [aux_vars[j].solution_value() for j in range(M)]
        
        print(f"Optimized objective value: {solver.Objective().Value()}")
        print(aux_vars_return)
        return variables_return
    else:
        print("No solution found.")
        return None