import pyomo.environ as pyo


def slickoil():
    model = pyo.ConcreteModel(name='Slick_Oil')

    model.wells = ['Well1', 'Well2', 'Well3', 'Well4', 'Well5', 'Well6']
    model.refineries = ['RefineryA', 'RefineryB', 'RefineryC', 'RefineryD', 'RefineryE']

    model.pipelines = [('Well1', 'RefineryA'), ('Well1', 'RefineryB'),
                       ('Well2', 'RefineryB'), ('Well2', 'RefineryC'), ('Well2', 'RefineryD'),
                       ('Well3', 'RefineryB'), ('Well3', 'RefineryC'), ('Well3', 'RefineryE'),
                       ('Well4', 'RefineryA'), ('Well4', 'RefineryE'),
                       ('Well5', 'RefineryD'), ('Well5', 'RefineryE'),
                       ('Well6', 'RefineryC'), ('Well6', 'RefineryE')]

    model.flows = pyo.Var(model.pipelines, within=pyo.NonNegativeReals)

    model.constraint_values = {'Wells': {'Well1': 40, 'Well2': 50, 'Well3': 40,
                                         'Well4': 100, 'Well5': 100, 'Well6': 100},
                               'Refineries': {'RefineryA': 100, 'RefineryB': 40, 'RefineryC': 80,
                                              'RefineryD': 100, 'RefineryE': 80},
                               'Demand': 100}

    well_costs = {'Well1': 3, 'Well2': 12, 'Well3': 4, 'Well4': 6, 'Well5': 10, 'Well6': 2.5}
    refinery_costs = {'RefineryA': 4, 'RefineryB': 2, 'RefineryC': 6, 'RefineryD': 1, 'RefineryE': 8}

    # Define the objective as minimize the costs of getting the oil to the customer
    model.obj = pyo.Objective(expr=sum(model.flows[(well, refinery)] * (well_costs[well] + refinery_costs[refinery])
                                       for (well, refinery) in model.pipelines),
                              sense=pyo.minimize)

    def demand_constraint(mdl):
        return (sum(mdl.flows[(well, refinery)]
                    for (well, refinery) in mdl.pipelines) >= 100)
    model.demand_constraint = pyo.Constraint(rule=demand_constraint)

    def well_constraints(mdl, input_well):
        return (sum(mdl.flows[(well, refinery)]
                    for (well, refinery) in mdl.pipelines if well == input_well)
                <= mdl.constraint_values['Wells'][input_well])
    model.well_constraints = pyo.Constraint(model.wells, rule=well_constraints)

    def refinery_constraints(mdl, input_refinery):
        return (sum(mdl.flows[(well, refinery)]
                    for (well, refinery) in mdl.pipelines if refinery == input_refinery)
                <= mdl.constraint_values['Refineries'][input_refinery])
    model.refinery_constraints = pyo.Constraint(model.refineries, rule=refinery_constraints)

    model.write("slick_oil.mps", io_options={"symbolic_solver_labels": True})


if __name__ == '__main__':
    slickoil()