import argparse
import sys

from _test_functions.realworld_functions import *
from _test_functions.synthetic_functions import (BaseBraninCurrin, BaseDTLZ1,
                                                 BaseDTLZ2, BaseDTLZ3,
                                                 BaseDTLZ4, BaseDTLZ7,
                                                 BaseOKA1, BaseOKA2, BaseZDT1,
                                                 BaseZDT2, BaseZDT3, BaseVLMOP2,
                                                 BaseVLMOP3)


def get_args():
    parser = argparse.ArgumentParser(description='Get function name')
    parser.add_argument('-f', '--func_name', type=str, default='', help='Objective name')
    parser.add_argument('-d', '--dim', type=int, default=5, help='Number of dimension')
    parser.add_argument('-m', '--n_objs', type=int, default=2, help='Number of objectives')
    parser.add_argument('-n', '--max_evals', type=int, default=200, help='Number of function evaluations')
    parser.add_argument('-b', '--batch', type=int, default=10, help='Number of batch')
    parser.add_argument('--n_init', type=int, default=20, help='Number of initial points')
    parser.add_argument('--n_beta', type=int, default=20, help='Number of discretized points on approximated CHIM')

    parser.add_argument('--seed', type=int, default=1, help='Random seed')
    args, unknown_args = parser.parse_known_args()

    input_dicts = vars(args)
    input_dicts['objective_function'] = get_function(args.func_name, dim=args.dim, num_objectives=args.n_objs)
    unknown_args_dict = {}
    for i in range(0, len(unknown_args), 2):
        unknown_args_dict[unknown_args[i].strip('--')] = unknown_args[i+1]
    input_dicts['unknown_args'] = unknown_args_dict
    return input_dicts


def get_function(objective_name, **kwargs):
    if objective_name == "dtlz1":
        return BaseDTLZ1(**kwargs)
    elif objective_name == "dtlz2":
        return BaseDTLZ2(**kwargs)
    elif objective_name == "dtlz3":
        return BaseDTLZ3(**kwargs)
    elif objective_name == "dtlz4":
        return BaseDTLZ4(**kwargs)
    elif objective_name == "dtlz7":
        return BaseDTLZ7(**kwargs)
    
    elif objective_name == "zdt1":
        return BaseZDT1(**kwargs)
    elif objective_name == "zdt2":
        return BaseZDT2(**kwargs)
    elif objective_name == "zdt3":
        return BaseZDT3(**kwargs)
    
    elif objective_name == "oka1":
        return BaseOKA1(**kwargs)
    elif objective_name == "oka2":
        return BaseOKA2(**kwargs)
    
    elif objective_name == "vlmop2":
        return BaseVLMOP2(**kwargs)
    elif objective_name == "vlmop3":
        return BaseVLMOP3(**kwargs)

    elif objective_name == "branin-currin":
        return BaseBraninCurrin(**kwargs)
    
    elif objective_name == "vehicle-design":
        return BaseVehicleDesign(**kwargs)
    elif objective_name == "four-bar-truss":
        return BaseFourBarTruss(**kwargs)
    elif objective_name == "concrete-beam":
        return BaseReinforcedConcreteBeam(**kwargs)
    elif objective_name == "pressure-vessel":
        return BasePressureVessel(**kwargs)
    elif objective_name == "hatch-cover":
        return BaseHatchCover(**kwargs)
    elif objective_name == "coil-spring":
        return BaseCoilCompressionSpring(**kwargs)
    elif objective_name == "two-bar-truss":
        return BaseTwoBarTruss(**kwargs)
    elif objective_name == "welded-beam":
        return BaseWeldedBeam(**kwargs)
    elif objective_name == "disc-brake":
        return BaseDiscBrake(**kwargs)
    elif objective_name == "speed-reducer":
        return BaseSpeedReducer(**kwargs)
    elif objective_name == "gear-train":
        return BaseGearTrain(**kwargs)
    elif objective_name == "rocket-injector":
        return BaseRocketInjector(**kwargs)
    elif objective_name == "car-impact":
        return BaseCarSideImpact(**kwargs)
    elif objective_name == "marine-design":
        return BaseConceptualMarineDesign(**kwargs)
    elif objective_name == "water-planning":
        return BaseWaterResourcePlanning(**kwargs)
    elif objective_name == "car-cab-design":
        return BaseCarCabDesign(**kwargs)
    
    elif objective_name == "mazda-car":
        return BaseMazdaCarsOptimization(**kwargs)
    else:
        raise ValueError(f"Unknown objective name: {objective_name}")
