#!/usr/bin/env python3
"""
DOCPLEX implementation for the musical optimization problem
"""

from docplex.mp.model import Model

def musical_optimization():
    """Optimize actor assignments to maximize audience engagement"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="musical_optimization")
    
    # Data from the problem
    actors = [
        {"actor_id": 1, "age": 28, "performance_duration": 15},
        {"actor_id": 2, "age": 35, "performance_duration": 20},
        {"actor_id": 3, "age": 22, "performance_duration": 10}
    ]
    
    roles = [
        {"role_id": 1, "musical_id": 1, "role_name": "Lead"},
        {"role_id": 2, "musical_id": 1, "role_name": "Supporting"},
        {"role_id": 3, "musical_id": 2, "role_name": "Chorus"}
    ]
    
    # Weights and maximum duration
    age_weight = 0.6
    duration_weight = 0.4
    max_duration = 100
    
    # CRITICAL: Validate array lengths
    assert len(actors) > 0 and len(roles) > 0, "Empty actors or roles list"
    
    # 2. VARIABLES
    # Binary decision variables: x[a][r] = 1 if actor a is assigned to role r
    x = {(a["actor_id"], r["role_id"]): mdl.binary_var(name=f"x_{a['actor_id']}_{r['role_id']}")
         for a in actors for r in roles}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total engagement: sum over all actors and roles of (0.6 * age + 0.4 * duration) * x[a][r]
    engagement = mdl.sum((age_weight * a["age"] + duration_weight * a["performance_duration"]) * x[(a["actor_id"], r["role_id"])]
                         for a in actors for r in roles)
    mdl.maximize(engagement)
    
    # 4. CONSTRAINTS
    
    # Actor Assignment Limit: Each actor can be assigned to at most one role
    for a in actors:
        mdl.add_constraint(mdl.sum(x[(a["actor_id"], r["role_id"])] for r in roles) <= 1,
                          ctname=f"actor_limit_{a['actor_id']}")
    
    # Role Fulfillment: Each role must be filled by exactly one actor
    for r in roles:
        mdl.add_constraint(mdl.sum(x[(a["actor_id"], r["role_id"])] for a in actors) == 1,
                          ctname=f"role_fulfillment_{r['role_id']}")
    
    # Total Duration Limit: Combined duration of all performances must not exceed max_duration
    total_duration = mdl.sum(a["performance_duration"] * x[(a["actor_id"], r["role_id"])]
                             for a in actors for r in roles)
    mdl.add_constraint(total_duration <= max_duration, ctname="total_duration_limit")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for a in actors:
            for r in roles:
                if solution.get_value(x[(a["actor_id"], r["role_id"])]) > 0.5:
                    print(f"Actor {a['actor_id']} assigned to Role {r['role_id']}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
if __name__ == "__main__":
    musical_optimization()