# Complete GUROBIPY implementation

import gurobipy as gp
from gurobipy import GRB

def optimize_budget_allocation():
    """Optimize marketing budget allocation across cities to maximize rating improvement."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("restaurant_budget_allocation")
    
    # Data: Coefficients for rating improvement per unit of budget
    city_ids = [1, 2, 3]
    coefficients = [0.15, 0.25, 0.10]
    total_budget = 100000
    
    # CRITICAL: Validate array lengths before loops
    assert len(city_ids) == len(coefficients), "Array length mismatch"
    
    # 2. VARIABLES
    # Variable dictionaries for budget allocations
    allocation = {city_id: model.addVar(vtype=GRB.CONTINUOUS, name=f"allocation_{city_id}", lb=0) 
                  for city_id in city_ids}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total expected improvement in restaurant ratings
    model.setObjective(gp.quicksum(coefficients[i] * allocation[city_ids[i]] for i in range(len(city_ids))), GRB.MAXIMIZE)
    
    # 4. CONSTRAINTS - CORRECT SYNTAX PATTERNS
    
    # Total budget constraint
    model.addConstr(gp.quicksum(allocation[city_id] for city_id in city_ids) <= total_budget, name="budget_constraint")
    
    # Non-negativity constraints are implicitly handled by the lower bound in variable definition
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for city_id in city_ids:
            if allocation[city_id].x > 1e-6:
                print(f"Allocation for city {city_id}: {allocation[city_id].x:.2f}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Run the optimization
optimize_budget_allocation()