# Complete DOCPLEX implementation

from docplex.mp.model import Model

def riding_club_optimization():
    """Riding Club Optimization using DOCPLEX"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="riding_club")
    
    # Data from the problem statement
    compatibility_scores = [
        (1, 1, 0.85),
        (1, 2, 0.65),
        (2, 1, 0.75),
        (2, 3, 0.6),
        (3, 2, 0.9),
        (3, 3, 0.7)
    ]
    
    coach_constraints = {
        1: 5,
        2: 4,
        3: 6
    }
    
    min_compatibility_score = 0.7
    
    # Extract unique player and coach IDs
    player_ids = set(cs[0] for cs in compatibility_scores)
    coach_ids = set(cs[1] for cs in compatibility_scores)
    
    # 2. VARIABLES
    # Binary decision variables for player-coach assignments
    x = {(i, j): mdl.binary_var(name=f"x_{i}_{j}") for i in player_ids for j in coach_ids}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total compatibility score from player-coach assignments
    objective = mdl.sum(cs[2] * x[cs[0], cs[1]] for cs in compatibility_scores)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Coach Capacity Constraints
    for j in coach_ids:
        mdl.add_constraint(mdl.sum(x[i, j] for i in player_ids if (i, j) in x) <= coach_constraints[j], ctname=f"coach_capacity_{j}")
    
    # Player Assignment Constraints
    for i in player_ids:
        mdl.add_constraint(mdl.sum(x[i, j] for j in coach_ids if (i, j) in x) <= 1, ctname=f"player_assignment_{i}")
    
    # Compatibility Threshold Constraints
    for cs in compatibility_scores:
        i, j, score = cs
        if score < min_compatibility_score:
            mdl.add_constraint(x[i, j] == 0, ctname=f"compatibility_threshold_{i}_{j}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for (i, j) in x:
            value = solution.get_value(x[i, j])
            if value > 1e-6:
                print(f"x[{i},{j}] = {value:.3f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
riding_club_optimization()