#!/usr/bin/env python3
"""
DOCplex implementation for Gymnast Team Selection Optimization
"""

from docplex.mp.model import Model

def gymnast_optimization():
    """Optimize gymnast team selection to maximize total points"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="gymnast_team_selection")
    
    # Performance scores for each gymnast in each event
    performance_scores = [
        [9.5, 8.7, 9.0, 9.2, 8.9, 9.1],  # Gymnast 1
        [8.8, 9.3, 8.7, 8.9, 9.0, 8.6],  # Gymnast 2
        [9.0, 8.5, 9.2, 9.1, 8.8, 9.3]   # Gymnast 3
    ]
    
    # Constraints
    max_events_per_gymnast = 3
    max_gymnasts_on_team = 10
    
    # Number of gymnasts and events
    n_gymnasts = len(performance_scores)
    n_events = len(performance_scores[0]) if n_gymnasts > 0 else 0
    
    # CRITICAL: Validate array lengths
    assert all(len(scores) == n_events for scores in performance_scores), "Performance scores array length mismatch"
    
    # Safe range for indexing
    safe_gymnasts = range(n_gymnasts)
    safe_events = range(n_events)
    
    # 2. VARIABLES
    # Decision variables: x[i][j] = 1 if gymnast i participates in event j
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in safe_gymnasts for j in safe_events}
    
    # Auxiliary variables: y[i] = 1 if gymnast i is on the team
    y = {i: mdl.binary_var(name=f"y_{i}") for i in safe_gymnasts}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total points scored by the team
    objective = mdl.sum(performance_scores[i][j] * x[i, j] for i in safe_gymnasts for j in safe_events)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Event Participation Limit: Each gymnast can participate in at most 3 events
    for i in safe_gymnasts:
        mdl.add_constraint(mdl.sum(x[i, j] for j in safe_events) <= max_events_per_gymnast, ctname=f"event_limit_{i}")
    
    # Team Size Limit: Total number of gymnasts on the team cannot exceed 10
    mdl.add_constraint(mdl.sum(y[i] for i in safe_gymnasts) <= max_gymnasts_on_team, ctname="team_size_limit")
    
    # Link y[i] to x[i,j]: y[i] = 1 if gymnast i participates in any event
    for i in safe_gymnasts:
        mdl.add_constraint(mdl.sum(x[i, j] for j in safe_events) <= n_events * y[i], ctname=f"link_{i}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_gymnasts:
            for j in safe_events:
                if solution.get_value(x[i, j]) > 0.5:
                    print(f"Gymnast {i+1} participates in Event {j+1}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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