# Complete DOCPLEX implementation

from docplex.mp.model import Model

def city_record_optimization():
    """Optimization for match allocation to maximize GDP impact"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="city_record")
    
    # Data: Cities, GDP, and Temperature suitability
    cities = [1, 2, 3]
    gdp = {1: 1500.0, 2: 2500.0, 3: 1800.0}
    temperature_suitability = {
        (1, 'Jan'): (16.0, 24.0),
        (2, 'Jan'): (18.0, 26.0),
        (3, 'Jan'): (15.0, 23.0),
        (1, 'Feb'): (17.0, 25.0),
        (2, 'Feb'): (19.0, 27.0),
        (3, 'Feb'): (16.0, 24.0),
        (1, 'Mar'): (18.0, 26.0),
        (2, 'Mar'): (20.0, 28.0),
        (3, 'Mar'): (17.0, 25.0)
    }
    
    matches = [1, 2, 3]
    months = ['Jan', 'Feb', 'Mar']
    
    # CRITICAL: Validate array lengths to prevent IndexError
    assert len(cities) == len(gdp), "Array length mismatch for cities and GDP"
    
    # 2. VARIABLES
    x = {(m, c): mdl.binary_var(name=f"x_{m}_{c}") for m in matches for c in cities}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(gdp[c] * x[m, c] for m in matches for c in cities)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Each match is allocated to at most one city
    for m in matches:
        mdl.add_constraint(mdl.sum(x[m, c] for c in cities) <= 1, ctname=f"match_{m}_allocation")
    
    # Temperature suitability constraint
    for m in matches:
        for c in cities:
            for month in months:
                min_temp, max_temp = temperature_suitability[(c, month)]
                if min_temp < 15 or max_temp > 27:
                    mdl.add_constraint(x[m, c] == 0, ctname=f"temp_suitability_{m}_{c}_{month}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for m in matches:
            for c in cities:
                value = solution.get_value(x[m, c])
                if value > 0.5:  # Since x is binary, check if it's 1
                    print(f"Match {m} is allocated to City {c}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

# Run the optimization
city_record_optimization()