# Complete DOCPLEX implementation - Retry Attempt 2

from docplex.mp.model import Model

def storm_resource_allocation():
    # Model setup
    mdl = Model(name="storm_resource_allocation")
    
    # Data setup
    storm_ids = [1, 2, 3]
    region_ids = [101, 102, 103]
    
    # Storm details
    storm_details = {
        1: {'storm_speed': 110, 'Damage_millions_USD': 15.5, 'Number_Deaths': 55},
        2: {'storm_speed': 90, 'Damage_millions_USD': 10.0, 'Number_Deaths': 40},
        3: {'storm_speed': 130, 'Damage_millions_USD': 20.0, 'Number_Deaths': 70}
    }
    
    # Allocation costs
    allocation_costs = {
        (1, 101): 5000,
        (1, 102): 4000,
        (2, 103): 6000
    }
    
    # Weights
    w1 = 1.0  # Weight for damage
    w2 = 1.0  # Weight for loss of life
    
    # Budget
    total_budget = 1500000
    
    # Validate array lengths
    assert len(storm_ids) == len(storm_details), "Storm IDs and details mismatch"
    assert len(allocation_costs) == len([(s, r) for s in storm_ids for r in region_ids if (s, r) in allocation_costs]), "Allocation costs mismatch"
    
    # Decision variables
    is_allocated = mdl.binary_var_dict([(s, r) for s in storm_ids for r in region_ids], name="is_allocated")
    
    # Objective function
    total_damage = mdl.sum(storm_details[s]['Damage_millions_USD'] for s in storm_ids)
    total_deaths = mdl.sum(storm_details[s]['Number_Deaths'] for s in storm_ids)
    objective = w1 * total_damage + w2 * total_deaths
    mdl.minimize(objective)
    
    # Constraints
    # Budget constraint
    budget_constraint = mdl.sum(allocation_costs.get((s, r), 0) * is_allocated[s, r] for s in storm_ids for r in region_ids) <= total_budget
    mdl.add_constraint(budget_constraint, ctname="budget_constraint")
    
    # Speed constraint (only for storms with speed <= 120 km/h)
    for s in storm_ids:
        if storm_details[s]['storm_speed'] <= 120:
            mdl.add_constraint(mdl.sum(is_allocated[s, r] for r in region_ids) >= 0, ctname=f"speed_constraint_{s}")
    
    # Cities affected constraint
    cities_affected = mdl.sum(is_allocated[s, r] for s in storm_ids for r in region_ids)
    mdl.add_constraint(cities_affected <= 5, ctname="cities_affected_constraint")
    
    # Solve the model
    solution = mdl.solve()
    
    # Output results
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for s in storm_ids:
            for r in region_ids:
                if (s, r) in is_allocated and solution.get_value(is_allocated[s, r]) > 0:
                    print(f"Resources allocated to storm {s} in region {r}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Execute the function
storm_resource_allocation()