# Complete GUROBIPY implementation

import gurobipy as gp
from gurobipy import GRB

def bike_optimization():
    """Optimize bike allocation across stations to minimize penalty costs."""
    
    # 1. MODEL & DATA SETUP
    model = gp.Model("bike_optimization")
    
    # Data: Penalty costs, dock counts, expected demands, and total bikes
    shortage_penalty = [5, 4, 6]  # Example penalty costs for shortages
    excess_penalty = [3, 2, 4]    # Example penalty costs for excesses
    dock_counts = [15, 20, 10]
    expected_demands = [12, 18, 8]
    total_bikes = 3  # Total number of bikes in shortage and excess
    
    # Validate array lengths
    n_stations = len(dock_counts)
    assert len(shortage_penalty) == len(excess_penalty) == len(expected_demands) == n_stations, "Array length mismatch"
    
    # 2. VARIABLES
    # Decision variables for shortages and excesses at each station
    shortages = model.addVars(n_stations, vtype=GRB.CONTINUOUS, name="shortage", lb=0)
    excesses = model.addVars(n_stations, vtype=GRB.CONTINUOUS, name="excess", lb=0)
    
    # 3. OBJECTIVE FUNCTION
    # Minimize total penalty costs
    model.setObjective(
        gp.quicksum(shortage_penalty[i] * shortages[i] + excess_penalty[i] * excesses[i] for i in range(n_stations)),
        GRB.MINIMIZE
    )
    
    # 4. CONSTRAINTS
    # Total bikes constraint
    model.addConstr(
        gp.quicksum(shortages[i] + excesses[i] for i in range(n_stations)) == total_bikes,
        name="total_bikes"
    )
    
    # Dock capacity constraints
    model.addConstrs(
        (shortages[i] + expected_demands[i] <= dock_counts[i] for i in range(n_stations)),
        name="dock_capacity"
    )
    
    # 5. SOLVING & RESULTS
    model.optimize()
    
    if model.status == GRB.OPTIMAL:
        print(f"Optimal value: {model.objVal}")
        for i in range(n_stations):
            print(f"Station {i+1}: Shortage = {shortages[i].x:.3f}, Excess = {excesses[i].x:.3f}")
    elif model.status == GRB.INFEASIBLE:
        print("Problem is infeasible")
    elif model.status == GRB.UNBOUNDED:
        print("Problem is unbounded")
    
    return model

# Run the optimization
bike_optimization()