import os

# Hope we can get rid of absolute path in training scripts.
root_path = os.path.dirname(os.path.abspath(__file__))

# FENN
from parameter_set.TPP.fenn_parameter_set import training_hyperparameter as feps_t
from parameter_set.TPP.fenn_parameter_set import evaluator_hyperparameter as feps_p

# TFENN
from parameter_set.TPP.tfenn_parameter_set import training_hyperparameter as tfeps_t
from parameter_set.TPP.tfenn_parameter_set import evaluator_hyperparameter as tfeps_p

# FullyNN 
from parameter_set.TPP.fullynn_parameter_set import training_hyperparameter as fps_t
from parameter_set.TPP.fullynn_parameter_set import evaluator_hyperparameter as fps_p

# TFullyNN
from parameter_set.TPP.tfullynn_parameter_set import training_hyperparameter as tfps_t
from parameter_set.TPP.tfullynn_parameter_set import evaluator_hyperparameter as tfps_p

# Naive Marked Temporal Point Processes(Naive)
from parameter_set.TPP.naive_parameter_set import training_hyperparameter as naivehp_t
from parameter_set.TPP.naive_parameter_set import evaluator_hyperparameter as naivehp_p

# Transformer Hawkes Process(THP)
from parameter_set.TPP.thp_parameter_set import training_hyperparameter as thp_t
from parameter_set.TPP.thp_parameter_set import evaluator_hyperparameter as thp_p

# Recurrent Marked Hawkes Process(RMTPP)
from parameter_set.TPP.rmtpp_parameter_set import training_hyperparameter as rmtpp_t
from parameter_set.TPP.rmtpp_parameter_set import evaluator_hyperparameter as rmtpp_p

# Marked Recurrent Marked Hawkes Process(MRMTPP)
from parameter_set.TPP.marked_rmtpp_parameter_set import training_hyperparameter as mrmtpp_t
from parameter_set.TPP.marked_rmtpp_parameter_set import evaluator_hyperparameter as mrmtpp_p

# LogNormMix
from parameter_set.TPP.lognormmix_parameter_set import training_hyperparameter as ifl_t
from parameter_set.TPP.lognormmix_parameter_set import evaluator_hyperparameter as ifl_p

# Marked_LogNormMix
from parameter_set.TPP.marked_lognormmix_parameter_set import training_hyperparameter as mifl_t
from parameter_set.TPP.marked_lognormmix_parameter_set import evaluator_hyperparameter as mifl_p

# Self-attentive Hawkes Process(SAHP)
from parameter_set.TPP.sahp_parameter_set import training_hyperparameter as sahp_t
from parameter_set.TPP.sahp_parameter_set import evaluator_hyperparameter as sahp_p

# Recurrent Hawkes Process(RHP), a.k.a. CTLSTM
from parameter_set.TPP.ctlstm_parameter_set import training_hyperparameter as ctlstm_t
from parameter_set.TPP.ctlstm_parameter_set import evaluator_hyperparameter as ctlstm_p

# Intensity-free Integral-based model-numerical(IFIB-C)
from parameter_set.TPP.ifib_c_parameter_set import training_hyperparameter as ifcps_t
from parameter_set.TPP.ifib_c_parameter_set import evaluator_hyperparameter as ifcps_p

# Intensity-free Integral-based model-numerical(IFIB-N)
from parameter_set.TPP.ifib_n_parameter_set import training_hyperparameter as ifnps_t
from parameter_set.TPP.ifib_n_parameter_set import evaluator_hyperparameter as ifnps_p

# Transformer-powered Intensity-free Integral-based model-numerical(TIFIB-C)
from parameter_set.TPP.tifib_c_parameter_set import training_hyperparameter as tifps_t
from parameter_set.TPP.tifib_c_parameter_set import evaluator_hyperparameter as tifps_p


evaluate_parameter_set = {
    'fenn': {'train': feps_t, 'evaluate': feps_p},
    'tfenn': {'train': tfeps_t, 'evaluate': tfeps_p},
    'fullynn': {'train': fps_t, 'evaluate': fps_p},
    'tfullynn': {'train': tfps_t, 'evaluate': tfps_p},
    'naive': {'train': naivehp_t, 'evaluate': naivehp_p},
    'thp': {'train': thp_t, 'evaluate': thp_p},
    'rmtpp': {'train': rmtpp_t, 'evaluate': rmtpp_p},
    'marked_rmtpp': {'train': rmtpp_t, 'evaluate': rmtpp_p},
    'lognormmix': {'train': ifl_t, 'evaluate': ifl_p},
    'marked_lognormmix': {'train': mifl_t, 'evaluate': mifl_p},
    'sahp': {'train': sahp_t, 'evaluate': sahp_p},
    'ctlstm': {'train': ctlstm_t, 'evaluate': ctlstm_p},
    'ifib_c': {'train': ifcps_t, 'evaluate': ifcps_p},
    'tifib_c': {'train': tifps_t, 'evaluate': tifps_p},
    'ifib_n': {'train': ifnps_t, 'evaluate': ifnps_p},
}

def parameter_retriver(opt):
    model_parameter_set = evaluate_parameter_set[opt.model]
    required_parameter_set = model_parameter_set[opt.script_type][opt.dataset]
    
    return required_parameter_set