# Complete PYOMO implementation

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

def optimize_basketball_selection():
    """Optimize basketball team selection for the tournament."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Example data (replace with actual data from database)
    team_ids = [1, 2, 3, 4, 5]
    win_percents = {1: 85.0, 2: 78.5, 3: 72.0, 4: 80.5, 5: 75.0}
    conference_indicators = {1: 'East', 2: 'West', 3: 'South', 4: 'East', 5: 'West'}
    location_indicators = {1: 'New York', 2: 'Los Angeles', 3: 'Chicago', 4: 'New York', 5: 'Los Angeles'}
    
    # Business configuration
    total_teams = 3
    min_east = 1
    min_west = 1
    min_south = 1
    max_ny = 1
    max_la = 1
    max_chicago = 1
    
    # CRITICAL: Validate array lengths before indexing
    assert len(team_ids) == len(win_percents) == len(conference_indicators) == len(location_indicators), "Array length mismatch"
    
    # 3. SETS (Pyomo way to define indices)
    model.I = pyo.Set(initialize=team_ids)
    
    # 4. PARAMETERS (data containers)
    model.win_percent = pyo.Param(model.I, initialize=win_percents)
    model.conference_indicator = pyo.Param(model.I, initialize=conference_indicators)
    model.location_indicator = pyo.Param(model.I, initialize=location_indicators)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.I, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.win_percent[i] * model.x[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    # Total Teams Selected
    def total_teams_rule(model):
        return sum(model.x[i] for i in model.I) == total_teams
    model.total_teams_constraint = pyo.Constraint(rule=total_teams_rule)
    
    # East Conference Diversity
    def east_conference_rule(model):
        return sum(model.x[i] for i in model.I if model.conference_indicator[i] == 'East') >= min_east
    model.east_conference_constraint = pyo.Constraint(rule=east_conference_rule)
    
    # West Conference Diversity
    def west_conference_rule(model):
        return sum(model.x[i] for i in model.I if model.conference_indicator[i] == 'West') >= min_west
    model.west_conference_constraint = pyo.Constraint(rule=west_conference_rule)
    
    # South Conference Diversity
    def south_conference_rule(model):
        return sum(model.x[i] for i in model.I if model.conference_indicator[i] == 'South') >= min_south
    model.south_conference_constraint = pyo.Constraint(rule=south_conference_rule)
    
    # New York Geographical Limit
    def ny_limit_rule(model):
        return sum(model.x[i] for i in model.I if model.location_indicator[i] == 'New York') <= max_ny
    model.ny_limit_constraint = pyo.Constraint(rule=ny_limit_rule)
    
    # Los Angeles Geographical Limit
    def la_limit_rule(model):
        return sum(model.x[i] for i in model.I if model.location_indicator[i] == 'Los Angeles') <= max_la
    model.la_limit_constraint = pyo.Constraint(rule=la_limit_rule)
    
    # Chicago Geographical Limit
    def chicago_limit_rule(model):
        return sum(model.x[i] for i in model.I if model.location_indicator[i] == 'Chicago') <= max_chicago
    model.chicago_limit_constraint = pyo.Constraint(rule=chicago_limit_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)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nSelected Teams:")
        for i in model.I:
            if pyo.value(model.x[i]) > 0.5:  # Only print selected teams
                print(f"Team {i} selected with win percentage {model.win_percent[i]}")
        
    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
if __name__ == "__main__":
    optimize_basketball_selection()