#!/usr/bin/env python3
"""
DOCPLEX Implementation for Bodybuilder Team Selection Problem
"""

from docplex.mp.model import Model

def bodybuilder_team_selection():
    """Optimization model to select a team of bodybuilders maximizing performance score."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="bodybuilder_team_selection")
    
    # Data from the database
    snatch_scores = [150.5, 160.0, 170.5]
    clean_jerk_scores = [200.0, 210.5, 220.0]
    heights = [175.0, 180.0, 185.0]
    weights = [90.0, 95.0, 100.0]
    
    # Business configuration
    max_team_size = 5
    min_avg_height = 170.0
    max_avg_weight = 100.0
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(snatch_scores) == len(clean_jerk_scores) == len(heights) == len(weights), "Array length mismatch"
    safe_range = range(min(len(snatch_scores), len(clean_jerk_scores), len(heights), len(weights)))  # Safe indexing
    
    # 2. VARIABLES
    # Binary decision variables for each bodybuilder
    x = {i: mdl.binary_var(name=f"x_{i}") for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total performance score
    total_score = mdl.sum((snatch_scores[i] + clean_jerk_scores[i]) * x[i] for i in safe_range)
    mdl.maximize(total_score)
    
    # 4. CONSTRAINTS
    
    # Team Size Limit
    team_size = mdl.sum(x[i] for i in safe_range)
    mdl.add_constraint(team_size <= max_team_size, ctname="team_size_limit")
    
    # Minimum Average Height
    total_height = mdl.sum(heights[i] * x[i] for i in safe_range)
    mdl.add_constraint(total_height >= min_avg_height * team_size, ctname="min_avg_height")
    
    # Maximum Average Weight
    total_weight = mdl.sum(weights[i] * x[i] for i in safe_range)
    mdl.add_constraint(total_weight <= max_avg_weight * team_size, ctname="max_avg_weight")
    
    # 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"Bodybuilder {i} selected: {value:.0f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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