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]
    is_not_tech_array = [not is_tech for is_tech in is_tech_array]
    
    # 这是一个组合优化问题，我们使用 SCIP 求解器来解决这个问题
    solver = pywraplp.Solver.CreateSolver("SCIP")
    if not solver:
        return

    # 定义变量的规模，一共有 N*M 个变量需要求解器求解
    N = len(preference_mat)     # 学生人数
    M = len(preference_mat[0])  # 班次数
    avg_num = sum(want_num_array) / M  # 每班次的平均人数
    print(f"平均人数：{avg_num}")

    variables = []
    aux_vars = []  # 辅助变量
    infinity = solver.infinity()
    for i in range(N):
        row_vars = []
        for j in range(M):
            var = solver.IntVar(0.0, 1.0, f"choice_{i}_{j}")
            row_vars.append(var)
        variables.append(row_vars)
    for j in range(M):
        aux_var = solver.NumVar(0.0, infinity, f"aux_{j}")
        aux_vars.append(aux_var)
    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):
        if want_num_array[i] == 3:
            solver.Add(sum(variables[i]) <= 3)
            solver.Add(sum(variables[i]) >= 2)
        else:
            solver.Add(sum(variables[i]) == want_num_array[i])

    # 添加约束：每个班次至少有？位同学
    if num_min is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j] for i in range(N)) >= num_min)

    # 添加约束：每个班次至多有？位同学
    if num_max is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j] for i in range(N)) <= num_max)

    # 添加约束：每班次最少包含？个电脑或电器的老人
    if num_tech_min is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_old_array[i]*is_tech_array[i] for i in range(N)) >= num_tech_min)

    # 添加约束：每班次最多包含？个电脑或电器的老人
    if num_tech_max is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_old_array[i]*is_tech_array[i] for i in range(N)) <= num_tech_max)

    # 添加约束：每班次至少包含？个老人
    if num_old_min is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_old_array[i] for i in range(N)) >= num_old_min)

    # 添加约束：每班次至多包含？个老人
    if num_old_max is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_old_array[i] for i in range(N)) <= num_old_max)

    # 添加约束：每班次至少包含？个人资部小朋友
    if num_hr_min is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_hr_array[i] for i in range(N)) >= num_hr_min)

    # 添加约束：每班次至多包含？个人资部小朋友
    if num_hr_max is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_hr_array[i] for i in range(N)) <= num_hr_max)

    # 添加约束：每班次至少包含？个小朋友
    if num_new_min is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_new_array[i] for i in range(N)) >= num_new_min)

    # 添加约束：每班次至多包含？个小朋友
    if num_new_max is not None:
        for j in range(M):
            solver.Add(sum(variables[i][j]*is_new_array[i] for i in range(N)) <= num_new_max)

    # 优化目标：每班次人数尽可能平均
    solver.Minimize(sum(aux_vars))  # Maximize the sum of these variables.

    # 求解优化问题
    status = solver.Solve()

    # 输出结果
    variables_return = []
    aux_vars_return = []
    if status == pywraplp.Solver.OPTIMAL:
        print("Optimal solution found:")
        for i in range(N):
            row_solution = [variables[i][j].solution_value() for j in range(M)]
            variables_return.append(row_solution)

        for j in range(M):
            aux_vars_return.append(aux_vars[j].solution_value())

        # Print the optimized value of the objective function.
        print(f"Optimized objective value: {solver.Objective().Value()}")
        print(aux_vars_return)

        return variables_return
    
    else:
        print("No solution found.")

        return None