# !/usr/bin/env python
# -*- coding: utf-8 -*-

#  Reference: https://zhuanlan.zhihu.com/p/142691916
#  Usage: Example for solving transport problem with pyomo.

from pyomo.environ import *
from pyomo.opt import SolverFactory

model = ConcreteModel()

# --------- 定义参数 --------- #
# 定义生产地和销售地
model.i = Set(initialize=['seattle', 'san-diego'], doc='Canning plans')
model.j = Set(initialize=['new-york', 'chicago', 'topeka'], doc='Markets')
# 定义生产地的生产量和销售地的需求量
model.a = Param(model.i, initialize={'seattle': 350, 'san-diego': 600}, doc='Capacity of plant i in cases')
model.b = Param(model.j, initialize={'new-york': 325, 'chicago': 300, 'topeka': 275}, doc='Demand at market j in cases')
# 定义生产地到销售地的距离
dtab = {
    ('seattle',  'new-york'): 2.5,
    ('seattle',  'chicago'): 1.7,
    ('seattle',  'topeka'): 1.8,
    ('san-diego','new-york'): 2.5,
    ('san-diego','chicago'): 1.8,
    ('san-diego','topeka'): 1.4,
    }
model.d = Param(model.i, model.j, initialize=dtab, doc='Distance in thousands of miles')
# 定义运输单价
model.f = Param(initialize=90, doc='Freight in dollars per case per thousand miles')
# 计算单位运价
def c_init(model, i, j):
    return model.f * model.d[i, j] / 1000
model.c = Param(model.i, model.j, initialize=c_init, doc='Transport cost in thousands of dollar per case')

# --------- 定义决策变量 --------- #
model.x = Var(model.i, model.j, bounds=(0.0, None), doc='Shipment quantities in case')

# --------- 定义约束条件 --------- #
# 生产地供应量不大于生产量
def supply_rule(model, i):
    return sum(model.x[i, j] for j in model.j) <= model.a[i]
model.supply = Constraint(model.i, rule=supply_rule, doc='Observe supply limit at plant i')

# 销售地得到的供应量不小于需求量
def demand_rule(model, j):
    return sum(model.x[i, j] for i in model.i) >= model.b[j]
model.demand = Constraint(model.j, rule=demand_rule, doc='Satisfy demand at market j')

# --------- 定义目标函数 --------- #
def objective_rule(model):
    return sum(model.c[i, j] * model.x[i, j] for i in model.i for j in model.j)
model.objective = Objective(rule=objective_rule, sense=minimize, doc='Define objective function')

# --------- 模型求解 --------- #
opt = SolverFactory("glpk")
results = opt.solve(model)
model.x.display()