# Complete PYOMO implementation - Retry Attempt 2

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

def storm_resource_allocation():
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Define sets for storms and regions
    storms = [1, 2, 3]
    regions = [101, 102, 103]
    
    # Define parameters
    damage_millions_usd = {1: 15.5, 2: 10.0, 3: 20.0}
    number_deaths = {1: 55, 2: 40, 3: 70}
    cost_per_allocation = {(1, 101): 5000, (1, 102): 4000, (2, 103): 6000}
    storm_speed = {1: 110, 2: 90, 3: 130}
    
    # Define weights
    w1 = 1.0  # Weight for damage
    w2 = 1.0  # Weight for loss of life
    
    # Define total budget
    total_budget = 1500000
    
    # 3. SETS
    model.storms = pyo.Set(initialize=storms)
    model.regions = pyo.Set(initialize=regions)
    
    # 4. PARAMETERS
    model.damage_millions_usd = pyo.Param(model.storms, initialize=damage_millions_usd)
    model.number_deaths = pyo.Param(model.storms, initialize=number_deaths)
    model.cost_per_allocation = pyo.Param(model.storms, model.regions, initialize=cost_per_allocation, default=0)
    model.storm_speed = pyo.Param(model.storms, initialize=storm_speed)
    
    # 5. VARIABLES
    model.is_allocated = pyo.Var(model.storms, model.regions, within=pyo.Binary)
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return w1 * sum(model.damage_millions_usd[storm_id] for storm_id in model.storms) + \
               w2 * sum(model.number_deaths[storm_id] for storm_id in model.storms)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    # Budget Constraint
    def budget_rule(model):
        return sum(model.cost_per_allocation[storm_id, region_id] * model.is_allocated[storm_id, region_id] 
                   for storm_id in model.storms for region_id in model.regions) <= total_budget
    model.budget_constraint = pyo.Constraint(rule=budget_rule)
    
    # Speed Constraint
    def speed_rule(model, storm_id):
        return model.storm_speed[storm_id] <= 120
    model.speed_constraint = pyo.Constraint(model.storms, rule=speed_rule)
    
    # Cities Affected Constraint
    def cities_affected_rule(model):
        return sum(model.is_allocated[storm_id, region_id] 
                   for storm_id in model.storms for region_id in model.regions) <= 5
    model.cities_affected_constraint = pyo.Constraint(rule=cities_affected_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    results = solver.solve(model, tee=True)
    
    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print(f"Optimal value: {pyo.value(model.objective)}")
    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 function
storm_resource_allocation()