# Complete PYOMO implementation - Retry Attempt 1

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

def mountain_photos_optimization():
    """Optimize the selection of mountain photos for maximum diversity."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Assuming the data is provided as lists for simplicity
    photo_ids = [1, 2, 3, 4, 5]
    diversity_scores = [0.85, 0.75, 0.9, 0.65, 0.8]
    
    # Constraints parameters
    P = 3  # Total number of photos allowed
    M_max = {1: 2, 2: 1, 3: 1, 4: 1, 5: 1}  # Max photos per mountain (assuming each photo is from a different mountain)
    
    # CRITICAL: Validate array lengths before indexing
    assert len(photo_ids) == len(diversity_scores), "Array length mismatch"
    
    # 3. SETS
    model.I = pyo.Set(initialize=photo_ids)
    
    # 4. PARAMETERS
    model.diversity_score = pyo.Param(model.I, initialize={photo_ids[i]: diversity_scores[i] for i in range(len(photo_ids))})
    
    # 5. VARIABLES
    model.x = pyo.Var(model.I, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.diversity_score[i] * model.x[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    
    # Total number of photos constraint
    def total_photos_rule(model):
        return sum(model.x[i] for i in model.I) <= P
    model.total_photos_constraint = pyo.Constraint(rule=total_photos_rule)
    
    # Photos per mountain constraint
    def photos_per_mountain_rule(model, i):
        return model.x[i] <= M_max[i]
    model.photos_per_mountain_constraint = pyo.Constraint(model.I, rule=photos_per_mountain_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Solve the model
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective):.2f}")
        
        # Extract variable values
        print("\nSelected photos:")
        for i in model.I:
            if pyo.value(model.x[i]) > 0.5:  # Binary variable, check if selected
                print(f"Photo ID {i} selected with diversity score {model.diversity_score[i]:.2f}")
        
    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

# Run the optimization
mountain_photos_optimization()