"""
Консольный скрипт обучения простой ML-модели. Используется как демо для реализации
работы сценария на разных входных данных, сгенерированных в пайплайне Snakemake.

Решается задача анализа зависимости состояния здоровья дерева (health)
от его диаметра ствола (tree_dbh) и типа дерева (spc_common).

Запуск из консоли:

python src/trees-training-workflow.py --file data/compiled-data-01.csv --artifact docs/artifacts/workflow-artifact-01.qmd

"""
import os
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
from tensorflow import keras
from tensorflow.keras import layers

import logging
import argparse

from inc.microfuncs import Microfuncs
from inc.templates import MdTemplates


# Создаем логгер и устанавливаем уровень логгирования
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Создаем обработчик для вывода логов в консоль
handler = logging.StreamHandler()
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)


# Гиперпараметры
EPOCHS = 10
BATCH_SIZE = 32
VALIDATION_SPLIT = 0.2


def main() -> None:
    
    # Получаем аргументы из команды запуска скрипта
    parser = argparse.ArgumentParser(description="")
    parser.add_argument(
        "--file", type=str, required=True, default="", help="Specify the input CSV file"
    )
    parser.add_argument(
        "--artifact", type=str, required=True, default="", help="Specify the output QMD file",
    )
    args = parser.parse_args()

    logger.info(f"Data file name received: {args.file}")
    logger.info(f"Artifact file name received: {args.artifact}")

    # Имена файлов данных и артефактов
    data_file = args.file
    md_file = args.artifact
    plot_file = Microfuncs.replace_file_extension(md_file, "png")

    # Загружаем датасет
    data = pd.read_csv(data_file, low_memory=False)
    logger.info(f"Dataset has been loaded. Shape of data: {data.shape}")

    # Предобработка данных
    data = data[["tree_dbh", "spc_common", "health"]]
    data = data.dropna()

    # Кодирование категориальных признаков
    label_encoder = LabelEncoder()
    data["spc_common"] = label_encoder.fit_transform(data["spc_common"])
    data["health"] = label_encoder.fit_transform(data["health"])

    # Разделение данных на train и test
    X = data[["tree_dbh", "spc_common"]]
    y = data["health"]
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=VALIDATION_SPLIT, random_state=42
    )

    # Нормализация числовых признаков
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Создание модели
    model = keras.Sequential(
        [
            layers.Dense(64, activation="relu", input_shape=(2,)),
            layers.Dense(32, activation="relu"),
            layers.Dense(3, activation="softmax"),
        ]
    )

    # Компиляция модели
    model.compile(
        optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
    )

    # Обучение модели
    history = model.fit(
        X_train_scaled,
        y_train,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        validation_split=VALIDATION_SPLIT,
    )

    # Оценка модели
    test_loss, test_acc = model.evaluate(X_test_scaled, y_test)
    logger.info(f"Training is complete. Accuracy: {test_acc:.3f}")

    # Сохранение графика в файл
    Microfuncs.save_plot_to_file(history, plot_file)
    logger.info(f"The training graph has been saved to a file: {plot_file}")

    # Собираем содержимое md-файла из шаблона с подстановкой значений
    content = Microfuncs.replace_all(
        MdTemplates.workflow_artifact,
        {
            "<ACCURACY>": f"{test_acc:.3f}",
            "<INPUT_FILE>": "data/" + os.path.split(data_file)[1],
            "<PLOT_FILE>": "artifacts/" + os.path.split(plot_file)[1]
        },
    )

    # Сохраняем md-файл, снегерированный из шаблона
    with open(md_file, "w", encoding="utf-8") as f:
        f.writelines(content)

    logger.info(f"Artifact has been saved to a file: {md_file}")


if __name__ == "__main__":
    main()