import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
import numpy as np

# 1 часть лабораторной работы
# Вариант 20. Pclass, Parch, Fare
def part_one():
    data = pd.read_csv('titanic.csv', index_col='Passengerid')
    # выгрузка непустых данных
    data = data.loc[(np.isnan(data['Pclass']) == False) & (np.isnan(data['Fare']) == False) & (np.isnan(data['Parch']) == False) & (np.isnan(data['Survived']) == False)]
    # отбор нужных столбцов
    corr = data[['Pclass', 'Parch', 'Fare']]
    # респечатка первых 5 строк данных
    print(corr.head())
    # определение целевой переменной
    y = data['Survived']
    # создание и обучение дерева решений
    clf = DecisionTreeClassifier(random_state=241)
    clf.fit(corr, y)
    # получение и распечатка важностей признаков
    importances = clf.feature_importances_
    print(importances)
    top_importances = importances.argsort()[-2:][::-1]
    print("Наиболее важные признаки:", corr.columns[top_importances][0], "и", corr.columns[top_importances][1])

# функция для приведения типа мобильного устройства к числу
def device_to_bool(device):
    if device == "Computer":
        return 0
    elif device == "Mobile":
        return 1
    elif device == "Tab":
        return 2

# функция для приведения уровня гибкости к числу
def flexibility_level_to_bool(flexibility_level):
    if flexibility_level == "Low":
        return 0
    elif flexibility_level == "Moderate":
        return 1
    elif flexibility_level == "High":
        return 2

# 2 часть лабораторной работы
# Вариант 20. Зависимость уровня гибкости от возраста и устройства, с которого человек работает
def part_two():
    data = pd.read_csv('dataset.csv')
    # приведение типа мобильного устройства к числу
    data['Device'] = data['Device'].apply(device_to_bool)
    # приведение уровня гибкости к числу
    data['Flexibility Level'] = data['Flexibility Level'].apply(flexibility_level_to_bool)
    # отбор нужных столбцов
    X = data[['Age', 'Device']]
    # респечатка первых 5 строк данных
    print(X.head())
    # определение целевой переменной
    y = data['Flexibility Level']
    # разделение данных на обучающую и тестовую выборки
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.01, random_state=42)
    # создание и обучение дерева регрессии
    tree_reg = DecisionTreeRegressor()
    tree_reg.fit(X_train, y_train)
    # получение и распечатка важностей признаков
    importances = tree_reg.feature_importances_
    print(importances)
    top_importances = importances.argsort()[-2:][::-1]
    print("Наиболее важные признаки:", X.columns[top_importances][0], "и", X.columns[top_importances][1])
    # предсказание на тестовых данных
    y_pred = tree_reg.predict(X_test)
    # оценка точности модели
    mse = mean_squared_error(y_test, y_pred)
    print("Средняя квадратичная ошибка:", mse)

print("---ПЕРВАЯ ЧАСТЬ ЛАБОРАТОРНОЙ РАБОТЫ---")
part_one()
print("\n---ВТОРАЯ ЧАСТЬ ЛАБОРАТОРНОЙ РАБОТЫ---")
part_two()