# Complete PYOMO implementation

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def bike_optimization():
    """Optimize bike redistribution to minimize unmet demands and movement costs."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    stations = [1, 2, 3]
    cost_per_bike_movement = 3.0
    unmet_demands = {1: 2, 2: 0, 3: 1}
    dock_capacities = {1: 15, 2: 20, 3: 10}
    initial_bikes = {1: 10, 2: 15, 3: 8}  # Assumed initial bike counts
    
    # Validate array lengths
    assert len(stations) == len(unmet_demands) == len(dock_capacities) == len(initial_bikes), "Array length mismatch"
    
    # 3. SETS
    model.I = pyo.Set(initialize=stations)  # Stations
    
    # 4. PARAMETERS
    model.cost_per_bike = pyo.Param(initialize=cost_per_bike_movement)
    model.unmet_demand = pyo.Param(model.I, initialize=unmet_demands)
    model.dock_capacity = pyo.Param(model.I, initialize=dock_capacities)
    model.initial_bikes = pyo.Param(model.I, initialize=initial_bikes)
    
    # 5. VARIABLES
    model.x = pyo.Var(model.I, model.I, within=pyo.NonNegativeIntegers)  # Bikes moved from i to j
    model.u = pyo.Var(model.I, within=pyo.NonNegativeIntegers)  # Unmet demands at station i
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return model.cost_per_bike * sum(model.x[i, j] for i in model.I for j in model.I if i != j) + sum(model.u[i] for i in model.I)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    
    # Station Capacity Constraint
    def station_capacity_rule(model, j):
        return sum(model.x[i, j] for i in model.I if i != j) <= model.dock_capacity[j]
    model.station_capacity_constraint = pyo.Constraint(model.I, rule=station_capacity_rule)
    
    # Initial Bike Availability Constraint
    def initial_bikes_rule(model, i):
        return sum(model.x[i, j] for j in model.I if i != j) <= model.initial_bikes[i]
    model.initial_bikes_constraint = pyo.Constraint(model.I, rule=initial_bikes_rule)
    
    # Unmet Demand Constraint
    def unmet_demand_rule(model, i):
        return model.u[i] >= 0
    model.unmet_demand_constraint = pyo.Constraint(model.I, rule=unmet_demand_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print("Optimal solution found!")
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nBike movements:")
        for i in model.I:
            for j in model.I:
                if i != j:
                    x_val = pyo.value(model.x[i, j])
                    if x_val > 1e-6:  # Only print non-zero values
                        print(f"Bikes moved from station {i} to station {j}: {int(x_val)}")
        
        print("\nUnmet demands:")
        for i in model.I:
            u_val = pyo.value(model.u[i])
            if u_val > 1e-6:  # Only print non-zero values
                print(f"Unmet demands at station {i}: {int(u_val)}")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

# Run the optimization
if __name__ == "__main__":
    bike_optimization()