#!/usr/bin/env python3
"""
Gurobipy 12.0.2 Implementation for Budget Allocation Problem
"""

import gurobipy as gp
from gurobipy import GRB

def budget_allocation_optimization():
    """Optimize budget allocation to maximize population served."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("budget_allocation")
    
    # Data from the problem
    cities = [1, 2, 3]
    populations = {1: 150000, 2: 250000, 3: 100000}
    districts = {1: 1, 2: 2, 3: 3}
    total_budget = 1000000
    min_city_allocation = 50000
    min_district_allocation = 200000
    
    # CRITICAL: Validate array lengths before loops
    assert len(cities) == len(populations) == len(districts), "Array length mismatch"
    
    # 2. VARIABLES
    # Decision variables: budget allocation to each city
    x = {i: model.addVar(vtype=GRB.CONTINUOUS, name=f"x_{i}", lb=min_city_allocation) 
         for i in cities}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize total population served
    model.setObjective(gp.quicksum(populations[i] * x[i] for i in cities), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS
    
    # Total Budget Limit
    model.addConstr(gp.quicksum(x[i] for i in cities) <= total_budget, name="total_budget_limit")
    
    # Minimum Allocation per District
    for i in cities:
        model.addConstr(x[i] >= min_district_allocation, name=f"district_min_{i}")
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in cities:
            print(f"x[{i}] = {x[i].x:.3f}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Execute the optimization
if __name__ == "__main__":
    budget_allocation_optimization()