# Complete PYOMO implementation

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

def bodybuilder_optimization():
    """Optimization for maximizing total weight lifted by bodybuilders"""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Data from the database
    snatch_impact = [1.2, 1.3, 1.1]
    clean_jerk_impact = [1.5, 1.6, 1.4]
    performance_targets = [300.0, 320.0, 340.0]
    total_limits = [190.0, 210.0, 220.0]
    
    # Validate array lengths
    n_bodybuilders = len(performance_targets)
    assert len(snatch_impact) == len(clean_jerk_impact) == len(total_limits) == n_bodybuilders, "Array length mismatch"
    
    # 3. SETS
    model.I = pyo.RangeSet(1, n_bodybuilders)  # 1-based indexing for bodybuilders
    
    # 4. PARAMETERS
    model.snatch_impact = pyo.Param(model.I, initialize={i+1: snatch_impact[i] for i in range(n_bodybuilders)})
    model.clean_jerk_impact = pyo.Param(model.I, initialize={i+1: clean_jerk_impact[i] for i in range(n_bodybuilders)})
    model.performance_target = pyo.Param(model.I, initialize={i+1: performance_targets[i] for i in range(n_bodybuilders)})
    model.total_limit = pyo.Param(model.I, initialize={i+1: total_limits[i] for i in range(n_bodybuilders)})
    
    # 5. VARIABLES
    model.x1 = pyo.Var(model.I, within=pyo.NonNegativeReals)  # Snatch weights
    model.x2 = pyo.Var(model.I, within=pyo.NonNegativeReals)  # Clean & Jerk weights
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.x1[i] + model.x2[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)
    
    # 7. CONSTRAINTS
    
    # Performance Target Constraints
    def performance_target_rule(model, i):
        return model.x1[i] + model.x2[i] <= model.performance_target[i]
    model.performance_target_constraint = pyo.Constraint(model.I, rule=performance_target_rule)
    
    # Training Impact Constraints
    def training_impact_rule(model, i):
        return model.snatch_impact[i] * model.x1[i] + model.clean_jerk_impact[i] * model.x2[i] <= model.total_limit[i]
    model.training_impact_constraint = pyo.Constraint(model.I, rule=training_impact_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("\nVariable values:")
        for i in model.I:
            x1_val = pyo.value(model.x1[i])
            x2_val = pyo.value(model.x2[i])
            print(f"Bodybuilder {i}: Snatch = {x1_val:.3f}, Clean & Jerk = {x2_val:.3f}")
        
    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
bodybuilder_optimization()