#!/usr/bin/env python3
"""
Pyomo 6.9.2 Implementation for Musical Actor-Role Assignment Problem
"""

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def optimize_actor_role_assignment():
    """Optimize actor-role assignments to maximize audience engagement."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Actors data: actor_id, age, performance_duration
    actors_data = {
        1: {'age': 28, 'performance_duration': 15},
        2: {'age': 35, 'performance_duration': 20},
        3: {'age': 22, 'performance_duration': 10}
    }
    
    # Roles data: role_id
    roles_data = {1, 2, 3}
    
    # Weights for engagement metric
    age_weight = 0.6
    duration_weight = 0.4
    
    # Maximum total duration limit
    max_total_duration = 100
    
    # 3. SETS
    model.A = pyo.Set(initialize=actors_data.keys())  # Set of actors
    model.R = pyo.Set(initialize=roles_data)         # Set of roles
    
    # 4. PARAMETERS
    model.age = pyo.Param(model.A, initialize={a: actors_data[a]['age'] for a in model.A})
    model.duration = pyo.Param(model.A, initialize={a: actors_data[a]['performance_duration'] for a in model.A})
    
    # 5. VARIABLES
    model.x = pyo.Var(model.A, model.R, within=pyo.Binary)  # Binary decision variable
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum((age_weight * model.age[a] + duration_weight * model.duration[a]) * model.x[a, r] for a in model.A for r in model.R)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    # Actor Assignment Limit: Each actor can be assigned to at most one role
    def actor_assignment_rule(model, a):
        return sum(model.x[a, r] for r in model.R) <= 1
    model.actor_assignment_constraint = pyo.Constraint(model.A, rule=actor_assignment_rule)
    
    # Role Fulfillment: Each role must be filled by exactly one actor
    def role_fulfillment_rule(model, r):
        return sum(model.x[a, r] for a in model.A) == 1
    model.role_fulfillment_constraint = pyo.Constraint(model.R, rule=role_fulfillment_rule)
    
    # Total Duration Limit: Combined duration must not exceed max_total_duration
    def total_duration_rule(model):
        return sum(model.duration[a] * model.x[a, r] for a in model.A for r in model.R) <= max_total_duration
    model.total_duration_constraint = pyo.Constraint(rule=total_duration_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nOptimal Assignments:")
        for a in model.A:
            for r in model.R:
                if pyo.value(model.x[a, r]) > 0.5:  # Only print assignments where x[a, r] = 1
                    print(f"Actor {a} assigned to Role {r}")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

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