#!/usr/bin/env python3
"""
DOCplex implementation for university_basketball optimization problem
"""

from docplex.mp.model import Model

def optimize_basketball_selection():
    """Optimize basketball team selection for tournament"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="university_basketball")
    
    # Business configuration parameters
    Total_Teams = 3
    Min_East = 1
    Min_West = 1
    Min_South = 1
    Max_NY = 1
    Max_LA = 1
    Max_Chicago = 1
    
    # Data from database
    Team_ID = [1, 2, 3, 4, 5]
    Win_Percent = [85.0, 78.5, 72.0, 80.5, 75.0]
    Conference_Indicator = ['East', 'West', 'South', 'East', 'West']
    Location_Indicator = ['New York', 'Los Angeles', 'Chicago', 'New York', 'Los Angeles']
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(Team_ID) == len(Win_Percent) == len(Conference_Indicator) == len(Location_Indicator), "Array length mismatch"
    safe_range = range(min(len(Team_ID), len(Win_Percent), len(Conference_Indicator), len(Location_Indicator)))  # Safe indexing
    
    # 2. VARIABLES
    x = {i: mdl.binary_var(name=f"x_{Team_ID[i]}") for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(Win_Percent[i] * x[i] for i in safe_range)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Total Teams Selected
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range) == Total_Teams, ctname="total_teams")
    
    # Conference Diversity
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Conference_Indicator[i] == 'East') >= Min_East, ctname="min_east")
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Conference_Indicator[i] == 'West') >= Min_West, ctname="min_west")
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Conference_Indicator[i] == 'South') >= Min_South, ctname="min_south")
    
    # Geographical Limits
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Location_Indicator[i] == 'New York') <= Max_NY, ctname="max_ny")
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Location_Indicator[i] == 'Los Angeles') <= Max_LA, ctname="max_la")
    mdl.add_constraint(mdl.sum(x[i] for i in safe_range if Location_Indicator[i] == 'Chicago') <= Max_Chicago, ctname="max_chicago")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            value = solution.get_value(x[i])
            if value > 1e-6:
                print(f"Team {Team_ID[i]} selected: {value:.0f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Execute the optimization
if __name__ == "__main__":
    optimize_basketball_selection()