#!/usr/bin/env python3
"""
DOCplex implementation for roller coaster optimization problem
"""

from docplex.mp.model import Model

def roller_coaster_optimization():
    """Optimize roller coaster distribution across parks to maximize visitor satisfaction"""
    
    # 1. MODEL & DATA SETUP
    mdl = Model(name="roller_coaster_optimization")
    
    # Data from the problem context
    parks = [1, 2, 3]
    coaster_types = ['Wooden', 'Steel', 'Inverted']
    
    # Visitor satisfaction scores
    visitor_satisfaction = {
        (1, 'Wooden'): 8.5,
        (1, 'Steel'): 9.0,
        (1, 'Inverted'): 7.5,
        (2, 'Wooden'): 8.0,
        (2, 'Steel'): 9.5,
        (2, 'Inverted'): 8.0,
        (3, 'Wooden'): 7.0,
        (3, 'Steel'): 9.0,
        (3, 'Inverted'): 8.5
    }
    
    # Park budgets
    park_budgets = {1: 1000000, 2: 1500000, 3: 2000000}
    
    # Park available space
    park_available_space = {1: 10000, 2: 15000, 3: 20000}
    
    # Park maximum roller coasters
    park_max_coasters = {1: 5, 2: 7, 3: 10}
    
    # Fixed cost and space requirements
    cost_per_coaster = 500000
    space_per_coaster = 2000
    
    # 2. VARIABLES
    x = {(p, c): mdl.integer_var(name=f"x_{p}_{c}", lb=0) for p in parks for c in coaster_types}
    
    # 3. OBJECTIVE FUNCTION
    objective = mdl.sum(visitor_satisfaction[(p, c)] * x[(p, c)] for p in parks for c in coaster_types)
    mdl.maximize(objective)
    
    # 4. CONSTRAINTS
    
    # Budget Constraint
    for p in parks:
        budget_expr = mdl.sum(cost_per_coaster * x[(p, c)] for c in coaster_types)
        mdl.add_constraint(budget_expr <= park_budgets[p], ctname=f"budget_{p}")
    
    # Space Constraint
    for p in parks:
        space_expr = mdl.sum(space_per_coaster * x[(p, c)] for c in coaster_types)
        mdl.add_constraint(space_expr <= park_available_space[p], ctname=f"space_{p}")
    
    # Maximum Roller Coasters Constraint
    for p in parks:
        max_coasters_expr = mdl.sum(x[(p, c)] for c in coaster_types)
        mdl.add_constraint(max_coasters_expr <= park_max_coasters[p], ctname=f"max_coasters_{p}")
    
    # 5. SOLVING & RESULTS
    solution = mdl.solve()
    
    if solution:
        print(f"Optimal value: {solution.objective_value}")
        for p in parks:
            for c in coaster_types:
                value = solution.get_value(x[(p, c)])
                if value > 1e-6:
                    print(f"Park {p}, Coaster Type {c}: {value:.0f}")
    else:
        print("No solution found")
        print(f"Status: {mdl.solve_details.status}")
    
    return mdl

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