#!/usr/bin/env python3
"""
Gurobipy 12.0.2 Implementation for Gymnast Team Selection Problem
"""

import gurobipy as gp
from gurobipy import GRB

def optimize_gymnast_team():
    """Optimize gymnast team selection to maximize total points across events."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("gymnast_team_selection")
    
    # Data from the problem
    gymnasts = [1, 2, 3]
    events = ['Floor_Exercise', 'Pommel_Horse', 'Rings', 'Vault', 'Parallel_Bars', 'Horizontal_Bar']
    
    # Performance scores for each gymnast in each event
    performance_scores = {
        1: {'Floor_Exercise': 9.5, 'Pommel_Horse': 8.7, 'Rings': 9.0, 'Vault': 9.2, 'Parallel_Bars': 8.9, 'Horizontal_Bar': 9.1},
        2: {'Floor_Exercise': 8.8, 'Pommel_Horse': 9.3, 'Rings': 8.7, 'Vault': 8.9, 'Parallel_Bars': 9.0, 'Horizontal_Bars': 8.6},
        3: {'Floor_Exercise': 9.0, 'Pommel_Horse': 8.5, 'Rings': 9.2, 'Vault': 9.1, 'Parallel_Bars': 8.8, 'Horizontal_Bar': 9.3}
    }
    
    # Constraints
    max_events_per_gymnast = 3
    max_gymnasts_on_team = 10
    
    # CRITICAL: Validate array lengths before loops
    assert len(gymnasts) == len(performance_scores), "Gymnast and performance scores length mismatch"
    
    # 2. VARIABLES
    # Decision variables: x[i,j] = 1 if gymnast i participates in event j, 0 otherwise
    x = model.addVars(gymnasts, events, vtype=GRB.BINARY, name="x")
    
    # Auxiliary variables: y[i] = 1 if gymnast i is on the team, 0 otherwise
    y = model.addVars(gymnasts, vtype=GRB.BINARY, name="y")
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total points scored by the team across all events
    model.setObjective(gp.quicksum(performance_scores[i][j] * x[i,j] for i in gymnasts for j in events), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS
    
    # Event Participation Limit: Each gymnast can participate in a maximum of 3 events
    for i in gymnasts:
        model.addConstr(gp.quicksum(x[i,j] for j in events) <= max_events_per_gymnast, name=f"event_limit_{i}")
    
    # Team Size Limit: The total number of gymnasts on the team cannot exceed 10
    model.addConstr(gp.quicksum(y[i] for i in gymnasts) <= max_gymnasts_on_team, name="team_size_limit")
    
    # Link y[i] to x[i,j]: y[i] = 1 if gymnast i participates in any event
    for i in gymnasts:
        model.addConstr(y[i] <= gp.quicksum(x[i,j] for j in events), name=f"link_y_{i}")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in gymnasts:
            for j in events:
                if x[i,j].x > 0.5:
                    print(f"Gymnast {i} participates in {j}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

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