#!/usr/bin/env python3
"""
Gurobipy Implementation for Roller Coaster Optimization Problem
"""

import gurobipy as gp
from gurobipy import GRB

def roller_coaster_optimization():
    """Optimize roller coaster distribution across theme parks"""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("roller_coaster_optimization")
    
    # Data from the problem
    parks = [1, 2, 3]
    coaster_types = ['Wooden', 'Steel', 'Inverted']
    
    # Visitor satisfaction scores
    satisfaction_scores = {
        (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_space = {1: 10000, 2: 15000, 3: 20000}
    
    # Park maximum roller coasters
    park_max_coasters = {1: 5, 2: 7, 3: 10}
    
    # Fixed cost and space per roller coaster
    cost_per_coaster = 500000
    space_per_coaster = 2000
    
    # CRITICAL: Validate array lengths before loops
    assert len(parks) == len(park_budgets) == len(park_space) == len(park_max_coasters), "Array length mismatch"
    
    # 2. VARIABLES
    x = model.addVars(parks, coaster_types, vtype=GRB.INTEGER, name="x", lb=0)
    
    # 3. OBJECTIVE FUNCTION
    model.setObjective(
        gp.quicksum(satisfaction_scores[p, c] * x[p, c] for p in parks for c in coaster_types),
        GRB.MAXIMIZE
    )
    
    # 4. CONSTRAINTS
    
    # Budget Constraint
    for p in parks:
        model.addConstr(
            gp.quicksum(cost_per_coaster * x[p, c] for c in coaster_types) <= park_budgets[p],
            name=f"budget_{p}"
        )
    
    # Space Constraint
    for p in parks:
        model.addConstr(
            gp.quicksum(space_per_coaster * x[p, c] for c in coaster_types) <= park_space[p],
            name=f"space_{p}"
        )
    
    # Maximum Roller Coasters Constraint
    for p in parks:
        model.addConstr(
            gp.quicksum(x[p, c] for c in coaster_types) <= park_max_coasters[p],
            name=f"max_coasters_{p}"
        )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for p in parks:
            for c in coaster_types:
                if x[p, c].x > 1e-6:
                    print(f"x[{p}, {c}] = {x[p, c].x:.0f}")
    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__":
    roller_coaster_optimization()