'''
Скрипт train.py должен:

- получить данные из двух таблиц postgres tit3_norm и tit3_alive;
- корректно разделить их на тренировочную и тестовую части;
- обучить модель используя библиотеку catboost;
- оценить качество результата;
- сохранить обученную модель в файл.

Источник: https://www.phind.com

Перед исполнением скрипта установить нужные модули так:

pip install psycopg2-binary catboost pandas numpy pickle
'''

import psycopg2
from catboost import CatBoostClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import pickle

def load_data():
    """Загрузка данных из PostgreSQL"""
    conn = psycopg2.connect(
        dbname='titanic',
        user="pupkin",
        password="1q2w3e"
    )

    cur = conn.cursor()

    # Загружаем данные из обеих таблиц
    cur.execute("""
        SELECT
            n.age_norm,
            n.pclass_norm,
            CASE WHEN n.sex = 'male' THEN 0 ELSE 1 END as sex,
            n.sibsp_norm,
            n.parch_norm,
            n.fare_norm,
            CASE WHEN n.embarked = 'S' THEN 0
                 WHEN n.embarked = 'C' THEN 1
                 ELSE 2 END as embarked,
            a.survived
        FROM tit3_norm n
        JOIN tit3_alive a ON n.id = a.id
    """)

    data = cur.fetchall()
    conn.close()

    return data

def train_and_save():
    """Обучение модели и её сохранение"""
    # Загружаем данные
    data = load_data()

    # Преобразуем в формат для обучения
    X = [row[:-1] for row in data]
    y = [row[-1] for row in data]

    # Разделяем на обучающую и тестовую выборки
    X_train, X_test, y_train, y_test = train_test_split(
      X, y, test_size=0.2, random_state=42)

    # Создаём и обучаем модель
    model = CatBoostClassifier(
        iterations=1000,
        depth=6,
        learning_rate=0.1,
        eval_metric='Accuracy'
    )

    model.fit(X_train, y_train)

    # Оцениваем модель
    y_pred = model.predict(X_test)
    print("\nМетрики качества:")
    print("Точность:", accuracy_score(y_test, y_pred))
    print("\nДетальный отчёт классификации:")
    print(classification_report(y_test, y_pred))

    # Сохраняем модель
    model.save_model('titanic_survival_model.cbm')

if __name__ == "__main__":
    train_and_save()


# End train.py