# Complete PYOMO implementation

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def optimize_speaker_allocation():
    """Optimize speaker allocation to debates to maximize audience reach."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Data from the database schema
    debates = [
        {'Debate_ID': 1, 'Num_of_Audience': 150},
        {'Debate_ID': 2, 'Num_of_Audience': 250},
        {'Debate_ID': 3, 'Num_of_Audience': 100}
    ]
    
    decision_variables = [
        {'Debate_ID': 1, 'People_ID': 101, 'assignment': True},
        {'Debate_ID': 1, 'People_ID': 102, 'assignment': False},
        {'Debate_ID': 2, 'People_ID': 101, 'assignment': True},
        {'Debate_ID': 2, 'People_ID': 103, 'assignment': True},
        {'Debate_ID': 3, 'People_ID': 102, 'assignment': True}
    ]
    
    constraint_bounds = {
        'Max_Debates_Per_Speaker': 3,
        'Max_Speakers_Per_Debate': 5
    }
    
    # Extract unique speaker and debate IDs
    speaker_ids = list(set(dv['People_ID'] for dv in decision_variables))
    debate_ids = list(set(d['Debate_ID'] for d in debates))
    
    # 3. SETS
    model.S = pyo.Set(initialize=speaker_ids)
    model.D = pyo.Set(initialize=debate_ids)
    
    # 4. PARAMETERS
    audience_dict = {d['Debate_ID']: d['Num_of_Audience'] for d in debates}
    model.Num_of_Audience = pyo.Param(model.D, initialize=audience_dict)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.S, model.D, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.Num_of_Audience[j] * model.x[i, j] for i in model.S for j in model.D)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    # Speaker Participation Limit
    def speaker_limit_rule(model, i):
        return sum(model.x[i, j] for j in model.D) <= constraint_bounds['Max_Debates_Per_Speaker']
    model.speaker_limit = pyo.Constraint(model.S, rule=speaker_limit_rule)
    
    # Debate Capacity Limit
    def debate_capacity_rule(model, j):
        return sum(model.x[i, j] for i in model.S) <= constraint_bounds['Max_Speakers_Per_Debate']
    model.debate_capacity = pyo.Constraint(model.D, rule=debate_capacity_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nVariable values:")
        for i in model.S:
            for j in model.D:
                x_val = pyo.value(model.x[i, j])
                if x_val > 1e-6:  # Only print non-zero values
                    print(f"x[{i},{j}] = {int(x_val)}")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

# Run the optimization
optimize_speaker_allocation()