#!/usr/bin/env python3
"""
DOCplex implementation for swimming event assignment optimization
"""

from docplex.mp.model import Model

def swimming_event_assignment():
    """Optimize swimmer assignments to events to maximize total performance score"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="swimming_event_assignment")
    
    # Data from the problem description
    swimmer_performance = [
        {'swimmer_id': 1, 'event_id': 1, 'score': 85.5},
        {'swimmer_id': 2, 'event_id': 1, 'score': 78.3},
        {'swimmer_id': 3, 'event_id': 1, 'score': 90.0}
    ]
    
    event_stadium = [
        {'event_id': 1, 'stadium_id': 1},
        {'event_id': 2, 'stadium_id': 2},
        {'event_id': 3, 'stadium_id': 3}
    ]
    
    stadium_capacity = [
        {'stadium_id': 1, 'capacity': 100, 'event_id': 1},
        {'stadium_id': 2, 'capacity': 150, 'event_id': 2},
        {'stadium_id': 3, 'capacity': 200, 'event_id': 3}
    ]
    
    # Create dictionaries for easy access
    swimmer_event_score = {(sp['swimmer_id'], sp['event_id']): sp['score'] for sp in swimmer_performance}
    event_stadium_map = {es['event_id']: es['stadium_id'] for es in event_stadium}
    stadium_event_capacity = {se['event_id']: se['capacity'] for se in stadium_capacity}
    
    # Get unique swimmers and events
    swimmers = list(set(sp['swimmer_id'] for sp in swimmer_performance))
    events = list(set(es['event_id'] for es in event_stadium))
    
    # 2. VARIABLES
    x = {(s, e): mdl.binary_var(name=f"x_{s}_{e}") for s in swimmers for e in events}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(swimmer_event_score[(s, e)] * x[(s, e)] for s in swimmers for e in events if (s, e) in swimmer_event_score)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Swimmer Assignment Limit: Each swimmer can be assigned to at most one event
    for s in swimmers:
        mdl.add_constraint(mdl.sum(x[(s, e)] for e in events if (s, e) in swimmer_event_score) <= 1, ctname=f"swimmer_limit_{s}")
    
    # Event Capacity Constraint: Number of swimmers assigned to each event must not exceed the capacity
    for e in events:
        mdl.add_constraint(mdl.sum(x[(s, e)] for s in swimmers if (s, e) in swimmer_event_score) <= stadium_event_capacity[e], ctname=f"event_capacity_{e}")
    
    # Performance Score Threshold: Only swimmers with a score of at least 80 can be assigned
    for s in swimmers:
        for e in events:
            if (s, e) in swimmer_event_score and swimmer_event_score[(s, e)] < 80:
                mdl.add_constraint(x[(s, e)] == 0, ctname=f"score_threshold_{s}_{e}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for s in swimmers:
            for e in events:
                if (s, e) in x and solution.get_value(x[(s, e)]) > 0.5:
                    print(f"Swimmer {s} assigned to Event {e} with score {swimmer_event_score[(s, e)]}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

if __name__ == "__main__":
    swimming_event_assignment()