#!/usr/bin/env python3
"""
DOCplex implementation for the riding club optimization problem
"""

from docplex.mp.model import Model

def riding_club_optimization():
    """Optimize player assignments to coaches to maximize total points"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="riding_club_optimization")
    
    # Player data
    players = [1, 2, 3]
    player_votes = {1: 12, 2: 18, 3: 9}
    
    # Coach data
    coaches = [1, 2, 3]
    coach_ranks = {1: 3, 2: 2, 3: 1}
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(players) == len(player_votes), "Player data length mismatch"
    assert len(coaches) == len(coach_ranks), "Coach data length mismatch"
    
    # 2. VARIABLES
    # Binary decision variables: x[p, c] = 1 if player p is assigned to coach c
    x = {(p, c): mdl.binary_var(name=f"x_{p}_{c}") for p in players for c in coaches}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total points: sum(player_votes[p] * coach_ranks[c] * x[p, c])
    objective = mdl.sum(player_votes[p] * coach_ranks[c] * x[p, c] for p in players for c in coaches)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Single Assignment Constraint: Each player must be assigned to exactly one coach
    for p in players:
        mdl.add_constraint(mdl.sum(x[p, c] for c in coaches) == 1, ctname=f"single_assignment_{p}")
    
    # Coach Capacity Constraint: Each coach can handle a maximum of four players
    for c in coaches:
        mdl.add_constraint(mdl.sum(x[p, c] for p in players) <= 4, ctname=f"coach_capacity_{c}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for p in players:
            for c in coaches:
                if solution.get_value(x[p, c]) > 0.5:
                    print(f"Player {p} assigned to Coach {c}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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