from ortools.sat.python import cp_model
from utils import read_excel, save_to_excel

def solve_program(preference_mat:list, 
                  want_num_array:list, 
                  is_new_array:list, 
                  is_tech_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_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]
    
    # 这是一个整数规划问题，我们使用 CP-SAT 求解器来解决这个问题
    model = cp_model.CpModel()

    # 定义变量的规模，一共有 N*M 个变量需要求解器求解
    N = len(preference_mat)     # 学生人数
    M = len(preference_mat[0])  # 班次数

    variables = []  
    for i in range(N):
        row_vars = []
        for j in range(M):
            var = model.NewBoolVar(f"choice_{i}_{j}")
            row_vars.append(var)
        variables.append(row_vars)

    # 添加约束：满足每位同学的意愿班次
    for i in range(N):
        model.Add(sum(variables[i]) == want_num_array[i])

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

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

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

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

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

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

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

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

    # 优化目标：最大化同学的满意度
    target_variables = [variables[i][j]*preference_mat[i][j] for i in range(N) for j in range(M)]
    model.Maximize(sum(target_variables))  # Maximize the sum of these variables.

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

    # 输出结果
    variables_return = []
    if status == cp_model.OPTIMAL:
        print("Optimal solution found:")
        for i in range(N):
            row_solution = [solver.Value(variables[i][j]) for j in range(M)]
            variables_return.append(row_solution)
            # print(" ".join(map(str, row_solution)))

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

        return variables_return
    
    elif status == cp_model.FEASIBLE:
        print("A potentially suboptimal solution was found.")
        for i in range(N):
            row_solution = [solver.Value(variables[i][j]) for j in range(M)]
            variables_return.append(row_solution)
            # print(" ".join(map(str, row_solution)))

        # Print the optimized value of the objective function.
        print(f"Suboptimal objective value: {solver.ObjectiveValue()}")

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

        return None