# ---
# jupyter:
#   jupytext:
#     formats: py:percent,ipynb
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.17.3
#   kernelspec:
#     display_name: .venv
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Исследование и настройка предсказательной модели для цен подержанных автомобилях

# %% [markdown]
# Блокнот использует файл аугментированных данных датасета о подержанных автомобилях, создаваемый блокнотом `eda/cars_eda.py`. См. ниже параметры блокнота для papermill.

# %%
from typing import Optional

# %% tags=["parameters"]
data_aug_pickle_path: Optional[str] = None
# Полный путь к файлу (pickle) для сохранения очищенного датасета. Если не установлен, используется `data/<data_aug_pickle_relpath>`.
data_aug_pickle_relpath: str = 'cars.aug.pickle'
# Путь к файлу (pickle) для сохранения очищенного датасета относительно директории данных `data`. Игнорируется, если установлен data_aug_pickle_path.

model_comment_path: Optional[str] = None
# Полный путь к текстовому файлу с произвольным комментарием для сохранения в MLFlow как артефакт вместе с моделью. Если не установлен, используется `research/<comment_relpath>`.
model_comment_relpath: str = 'comment.txt'
# Путь к текстовому файлу с произвольным комментарием для сохранения в MLFlow как артефакт вместе с моделью относительно директории `research`. Игнорируется, если установлен comment_path.

mlflow_tracking_server_uri: str = 'http://localhost:5000'
# URL tracking-сервера MLFlow.
mlflow_registry_uri: Optional[str] = None
# URL сервера registry MLFlow (если не указан, используется `mlflow_tracking_server_uri`).

mlflow_do_log: bool = False
# Записывать ли прогон (run) в MLFlow; если True, при каждом исполнении блокнота создаётся новый прогон с именем `mlflow_run_name`.
mlflow_experiment_id: Optional[str] = None
# ID эксперимента MLFlow, имеет приоритет над `mlflow_experiment_name`.
mlflow_experiment_name: Optional[str] = 'Current price predicion for used cars'
# Имя эксперимента MLFlow (ниже приоритетом, чем `mlflow_experiment_id`).
mlflow_run_name: str = 'Baseline model'
# Имя нового прогона MLFlow (используется для создания нового прогона, если `mlflow_do_log` установлен в True).

# %%
import os
import pathlib
import pickle
import sys

# %%
import mlflow
import mlflow.models
import mlflow.sklearn
import sklearn.compose
import sklearn.ensemble
import sklearn.metrics
import sklearn.model_selection
import sklearn.pipeline
import sklearn.preprocessing

# %%
BASE_PATH = pathlib.Path('..')

# %%
CODE_PATH = BASE_PATH
sys.path.insert(0, str(CODE_PATH.resolve()))

# %%
from iis_project.sklearn_utils import filter_params
from iis_project.sklearn_utils.compose import COLUMN_TRANSFORMER_PARAMS_COMMON_INCLUDE
from iis_project.sklearn_utils.ensemble import RANDOM_FOREST_REGRESSOR_PARAMS_COMMON_EXCLUDE
from iis_project.sklearn_utils.pipeline import PIPELINE_PARAMS_COMMON_INCLUDE
from iis_project.sklearn_utils.preprocessing import STANDARD_SCALER_PARAMS_COMMON_EXCLUDE

# %%
MODEL_INOUT_EXAMPLE_SIZE = 0x10

# %%
mlflow.set_tracking_uri(mlflow_tracking_server_uri)
if mlflow_registry_uri is not None:
    mlflow.set_registry_uri(mlflow_registry_uri)

# %%
if mlflow_do_log:
    mlflow_experiment = mlflow.set_experiment(experiment_name=mlflow_experiment_name, experiment_id=mlflow_experiment_id)

# %%
DATA_PATH = (
    pathlib.Path(os.path.dirname(data_aug_pickle_path))
    if data_aug_pickle_path is not None
    else (BASE_PATH / 'data')
)

# %% [markdown]
# ## Загрузка и обзор данных

# %%
with open(
    (
        data_aug_pickle_path
        if data_aug_pickle_path is not None
        else (DATA_PATH / data_aug_pickle_relpath)
    ),
    'rb',
) as input_file:
    df_orig = pickle.load(input_file)

# %% [markdown]
# Обзор строк датасета:

# %%
df_orig.head(0x10)

# %% [markdown]
# Размер датасета:

# %%
len(df_orig)

# %% [markdown]
# Количество непустых значений и тип каждого столбца:

# %%
df_orig.info()

# %% [markdown]
# ## Разделение датасета на выборки

# %% [markdown]
# Выделение признаков и целевых переменных:

# %%
feature_columns = (
    'selling_price',
    'driven_kms',
    'fuel_type',
    'selling_type',
    'transmission',
    #'owner',
    'age',
)

target_columns = (
    'present_price',
)

# %%
features_to_scale_to_standard_columns = (
    'selling_price',
    'driven_kms',
    'age',
)
assert all(
    (col in df_orig.select_dtypes(('number',)).columns)
    for col in features_to_scale_to_standard_columns
)

features_to_encode_wrt_target_columns = (
    'fuel_type',
    'selling_type',
    'transmission',
    #'owner',
)
assert all(
    (col in df_orig.select_dtypes(('category', 'object')).columns)
    for col in features_to_encode_wrt_target_columns
)

# %%
df_orig_features = df_orig[list(feature_columns)]
df_target = df_orig[list(target_columns)]

# %% [markdown]
# Разделение на обучающую и тестовую выборки:

# %%
DF_TEST_PORTION = 0.25

# %%
df_orig_features_train, df_orig_features_test, df_target_train, df_target_test = (
    sklearn.model_selection.train_test_split(
        df_orig_features, df_target, test_size=DF_TEST_PORTION, random_state=0x7AE6,
    )
)

# %% [markdown]
# Размеры обучающей и тестовой выборки соответственно:

# %%
tuple(map(len, (df_target_train, df_target_test)))

# %% [markdown]
# ## Создание пайплайнов обработки признаков и обучения модели

# %%
#MODEL_PIP_REQUIREMENTS_PATH = BASE_PATH / 'requirements' / 'requirements-isolated-research-model.txt'

# %% [markdown]
# Сигнатура модели для MLFlow:

# %%
mlflow_model_signature = mlflow.models.infer_signature(model_input=df_orig_features, model_output=df_target)
mlflow_model_signature

# %% [raw] vscode={"languageId": "raw"}
# input_schema = mlflow.types.schema.Schema([
#     mlflow.types.schema.ColSpec("double", "selling_price"),
#     mlflow.types.schema.ColSpec("double", "driven_kms"),
#     mlflow.types.schema.ColSpec("string", "fuel_type"),
#     mlflow.types.schema.ColSpec("string", "selling_type"),
#     mlflow.types.schema.ColSpec("string", "transmission"),
#     mlflow.types.schema.ColSpec("double", "age"),
# ])
#
# output_schema = mlflow.types.schema.Schema([
#     mlflow.types.schema.ColSpec("double", "present_price"),
# ])
#
# mlflow_model_signature = mlflow.models.ModelSignature(inputs=input_schema, outputs=output_schema)

# %% [markdown]
# Пайплайн предобработки признаков:

# %%
preprocess_transformer = sklearn.compose.ColumnTransformer(
    [
        ('scale_to_standard', sklearn.preprocessing.StandardScaler(), features_to_scale_to_standard_columns),
        (
            #'encode_categoricals_one_hot',
            'encode_categoricals_wrt_target',
            #sklearn.preprocessing.OneHotEncoder(),
            sklearn.preprocessing.TargetEncoder(
                target_type='continuous', smooth='auto', shuffle=True, random_state=0x2ED6,
            ),
            features_to_encode_wrt_target_columns,
        ),
    ],
    remainder='drop',
)

# %% [markdown]
# Регрессор &mdash; небольшой случайный лес, цель &mdash; минимизация квадрата ошибки предсказания:

# %%
regressor = sklearn.ensemble.RandomForestRegressor(
    10, criterion='squared_error', max_features='sqrt', random_state=0x016B,
)

# %% [markdown]
# Составной пайплайн:

# %%
pipeline = sklearn.pipeline.Pipeline([
    ('preprocess', preprocess_transformer),
    ('regress', regressor),
])
pipeline

# %%
model_params = filter_params(
    pipeline.get_params(),
    include={
        **{k: True for k in PIPELINE_PARAMS_COMMON_INCLUDE},
        'preprocess': (
            False,
            {
                **{k: True for k in COLUMN_TRANSFORMER_PARAMS_COMMON_INCLUDE},
                'scale_to_standard': True,
                'encode_categorical_wrt_target': True,
            },
        ),
        'regress': (False, True),
    },
    exclude={
        'preprocess': {'scale_to_standard': STANDARD_SCALER_PARAMS_COMMON_EXCLUDE},
        'regress': RANDOM_FOREST_REGRESSOR_PARAMS_COMMON_EXCLUDE,
    },
)
model_params

# %% [markdown]
# ## Baseline модель

# %%
_ = pipeline.fit(df_orig_features_train, df_target_train.iloc[:, 0])

# %%
target_test_predicted = pipeline.predict(df_orig_features_test)

# %% [markdown]
# Метрики качества (MAPE, а также MSE, MAE):

# %%
metrics = {
    'mse': sklearn.metrics.mean_squared_error(df_target_test, target_test_predicted),
    'mae': sklearn.metrics.mean_absolute_error(df_target_test, target_test_predicted),
    'mape': sklearn.metrics.mean_absolute_percentage_error(df_target_test, target_test_predicted),
}

# %%
metrics

# %%
if mlflow_do_log:
    with mlflow.start_run(experiment_id=mlflow_experiment.experiment_id, run_name=mlflow_run_name):
        _ = mlflow.sklearn.log_model(
            pipeline,
            'model',
            signature=mlflow_model_signature,
            input_example=df_orig_features.head(MODEL_INOUT_EXAMPLE_SIZE),
            #pip_requirements=str(MODEL_PIP_REQUIREMENTS_PATH),
        )
        _ = mlflow.log_params(model_params)
        _ = mlflow.log_metrics({k: float(v) for k, v in metrics.items()})
        comment_file_path = (
            model_comment_path
            if model_comment_path is not None
            else (BASE_PATH / 'research' / model_comment_relpath)
        )
        if comment_file_path.exists():
            mlflow.log_artifact(str(comment_file_path))