#!/usr/bin/env python3
"""
DOCPLEX Implementation for Budget Allocation Optimization Problem
"""

from docplex.mp.model import Model

def budget_allocation_optimization():
    """Optimize budget allocation to maximize population served."""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="budget_allocation")
    
    # Data from the problem
    cities = [1, 2, 3]
    populations = {1: 150000, 2: 250000, 3: 100000}
    districts = {1: 'North', 2: 'South', 3: 'East'}
    city_district = {1: 1, 2: 2, 3: 3}
    
    # Parameters
    total_budget = 1000000
    min_city_allocation = 50000
    min_district_allocation = 200000
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(cities) == len(populations) == len(city_district), "Array length mismatch"
    safe_range = range(min(len(cities), len(populations), len(city_district)))  # Safe indexing
    
    # 2. VARIABLES
    # Decision variables: budget allocation to each city
    x = {i: mdl.continuous_var(name=f"x_{i}", lb=0) for i in safe_range}
    
    # 3. OBJECTIVE FUNCTION
    # Maximize the total population served
    objective = mdl.sum(populations[i] * x[i] for i in safe_range)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Total Budget Limit
    total_allocation = mdl.sum(x[i] for i in safe_range)
    mdl.add_constraint(total_allocation <= total_budget, ctname="total_budget_limit")
    
    # Minimum Allocation per City
    for i in safe_range:
        mdl.add_constraint(x[i] >= min_city_allocation, ctname=f"min_city_allocation_{i}")
    
    # Minimum Allocation per District
    for i in safe_range:
        district_id = city_district[i]
        mdl.add_constraint(x[i] >= min_district_allocation, ctname=f"min_district_allocation_{district_id}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for i in safe_range:
            value = solution.get_value(x[i])
            print(f"x[{i}] = {value:.3f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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