# Complete PYOMO implementation - Retry Attempt 1

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

def music_festival_optimization():
    """Optimize artist selection for maximum exposure at a music festival."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Data from the problem statement
    artist_ids = [1, 2, 3]
    weeks_on_top = {1: 12, 2: 8, 3: 15}
    availability_status = {1: True, 2: False, 3: True}
    capacity = 2
    
    # 3. SETS
    model.I = pyo.Set(initialize=artist_ids)
    
    # 4. PARAMETERS
    model.weeks_on_top = pyo.Param(model.I, initialize=weeks_on_top)
    model.availability_status = pyo.Param(model.I, initialize=availability_status)
    model.capacity = pyo.Param(initialize=capacity)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.I, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.weeks_on_top[i] * model.x[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    # Capacity constraint
    def capacity_rule(model):
        return sum(model.x[i] for i in model.I) <= model.capacity
    model.capacity_constraint = pyo.Constraint(rule=capacity_rule)
    
    # Availability constraints
    def availability_rule(model, i):
        return model.x[i] <= int(model.availability_status[i])
    model.availability_constraint = pyo.Constraint(model.I, rule=availability_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nSelected artists:")
        for i in model.I:
            if pyo.value(model.x[i]) > 0.5:  # Binary variable, check if selected
                print(f"Artist {i} selected")
        
    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}")

# Execute the optimization function
music_festival_optimization()