#!/usr/bin/env python3
"""
DOCPLEX implementation for wrestler selection optimization problem
"""

from docplex.mp.model import Model

def wrestler_selection_optimization():
    """Optimize wrestler selection to maximize total days held while adhering to constraints"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="wrestler_selection")
    
    # Data from the provided SQL inserts
    days_held = [150, 250, 100]
    eliminations = [8, 12, 5]
    teams = ['Team A', 'Team B', 'Team A']
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(days_held) == len(eliminations) == len(teams), "Array length mismatch"
    safe_range = range(min(len(days_held), len(eliminations), len(teams)))  # Safe indexing
    
    # 2. VARIABLES
    # Binary decision variables for each wrestler
    x = {i: mdl.binary_var(name=f"x_{i}") for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total days held by selected wrestlers
    objective = mdl.sum(days_held[i] * x[i] for i in safe_range)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Total Wrestlers Constraint: Select at most 6 wrestlers
    total_wrestlers = mdl.sum(x[i] for i in safe_range)
    mdl.add_constraint(total_wrestlers <= 6, ctname="total_wrestlers_limit")
    
    # Total Eliminations Constraint: Total eliminations <= 30
    total_eliminations = mdl.sum(eliminations[i] * x[i] for i in safe_range)
    mdl.add_constraint(total_eliminations <= 30, ctname="total_eliminations_limit")
    
    # Team Balance Constraint: Equal number of wrestlers from Team A and Team B
    team_a = [i for i in safe_range if teams[i] == 'Team A']
    team_b = [i for i in safe_range if teams[i] == 'Team B']
    team_a_selected = mdl.sum(x[i] for i in team_a)
    team_b_selected = mdl.sum(x[i] for i in team_b)
    mdl.add_constraint(team_a_selected == team_b_selected, ctname="team_balance")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            if solution.get_value(x[i]) > 0.5:
                print(f"Wrestler {i} selected: Days held = {days_held[i]}, Eliminations = {eliminations[i]}, Team = {teams[i]}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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