import numpy as np
from {{import_library}} import {{ model_name }}

{% if model_arg == "HPO_noRandomSeed" %}
model = {{ model_name }}(**best_params)
{% elif model_arg == "HPO_RandomSeed" %}
random_state_model = {{ pipeline.config.seed_for_model}}
model = {{ model_name }}(random_state=random_state_model, **best_params)
{% elif model_arg == "noHPO_noRandomSeed" %}
model = {{ model_name }}({{ params }})
{% elif model_arg == "noHPO_RandomSeed" %}
random_state_model = {{ pipeline.config.seed_for_model}}
model = {{ model_name }}(random_state=random_state_model, {{ params }})
{% endif %}

{% if is_multioutput_regression%}
from sklearn.multioutput import MultiOutputRegressor

model = MultiOutputRegressor(model)
{% elif is_multioutput_classification %}
from sklearn.multioutput import MultiOutputClassifier

model = MultiOutputClassifier(model)
{% endif %}
{% set xgbclassifier = "XGBClassifier" %}
{% if is_multioutput_classification %}
from sklearn.preprocessing import LabelEncoder
label_encoders = {}
for i, column in enumerate(target_train.columns):
    le = LabelEncoder()
    target_train[column] = le.fit_transform(target_train[column])
    label_encoders[column] = le
with open('target_LabelEncoder.pkl', 'wb') as f:
    pickle.dump(label_encoders, f)
{% elif model_name == xgbclassifier %}
from sklearn.preprocessing import LabelEncoder

label_encoder = LabelEncoder()
target_train = pd.DataFrame(label_encoder.fit_transform(target_train), columns=TARGET_COLUMNS)
with open('target_LabelEncoder.pkl', 'wb') as f:
    pickle.dump(label_encoder, f)

{% endif %}
{% if pipeline.task.target_columns|length == 1 %}
model.fit(feature_train, target_train.values.ravel())
{% else %}
model.fit(feature_train, target_train)
{% endif %}
with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)