# Complete PYOMO implementation - Retry Attempt 1

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

def club_optimization():
    """Pyomo model for optimizing student-club assignments"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Data from the database schema
    preference_scores_data = {
        (1, 101): 8,
        (1, 102): 3,
        (2, 101): 5,
        (2, 103): 7,
        (3, 102): 6,
        (3, 103): 4
    }
    
    club_capacities_data = {
        101: 60,
        102: 40,
        103: 50
    }
    
    # Extract unique students and clubs
    students = set(stu for stu, _ in preference_scores_data.keys())
    clubs = set(club for _, club in preference_scores_data.keys())
    
    # CRITICAL: Validate data completeness
    assert all((stu, club) in preference_scores_data for stu in students for club in clubs), "Missing preference scores"
    
    # 3. SETS
    model.Students = pyo.Set(initialize=students)
    model.Clubs = pyo.Set(initialize=clubs)
    
    # 4. PARAMETERS
    model.PreferenceScore = pyo.Param(model.Students, model.Clubs, initialize=preference_scores_data)
    model.Capacity = pyo.Param(model.Clubs, initialize=club_capacities_data)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.Students, model.Clubs, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def objective_rule(model):
        return sum(model.PreferenceScore[stu, club] * model.x[stu, club] for stu in model.Students for club in model.Clubs)
    model.objective = pyo.Objective(rule=objective_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    def capacity_constraint_rule(model, club):
        return sum(model.x[stu, club] for stu in model.Students) <= model.Capacity[club]
    model.capacity_constraint = pyo.Constraint(model.Clubs, rule=capacity_constraint_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        optimal_value = pyo.value(model.objective)
        print(f"Optimal value: {optimal_value}")
        
        # Extract and print variable values
        print("\nAssignment decisions:")
        for stu in model.Students:
            for club in model.Clubs:
                if pyo.value(model.x[stu, club]) > 0.5:  # Binary variable, check if assigned
                    print(f"Student {stu} assigned to Club {club}")
        
    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}")

# Execute the optimization
club_optimization()