import random
import pyepo
import numpy as np
import pyomo.environ as pe
from pyomo.environ import ConcreteModel, Var, Objective, Constraint, SolverFactory
from models import ialloyProduction
from pyepo.model.grb import knapsack, shortestpath, tsp, portfolio
from pyepo.model.omo import optOmoModel
from pyepo import EPO

# Define variables

# 入口加模型选择
def opt_model(model_type, para1, para2, para3):
    if model_type == 'knapsack':
        optmodel = knapsack.knapsackModel(para1, para2)  # weights, capacity
    elif model_type == 'shortestpath':
        optmodel = shortestpath.shortestPathModel(para1)  # build model
    elif model_type == 'tsp':
        optmodel = tsp.tspDFJModel(para1)  # build model
    elif model_type == 'portfolio':
        optmodel = portfolio.portfolioModel(para1, para2)  # build model
    elif model_type == 'alloyproduction':
        optmodel = ialloyProduction.alloyProductionModel(para1, para2)  # cons, req
    return optmodel
