import math
import pandas as pd
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.model_selection import train_test_split

### Деревья решений для регрессии и классификации
### с использованием моделей DecisionTreeRegressor и DecisionTreeClassifier

# Загрузка данных
data = pd.read_csv('CO2.csv')
data = data.dropna()

# Хеширование наименований стран
countries = {}
for country in data['Country']:
    countries[country] = hash(country)
hash_column = []
for country in data['Country']:
    hash_column.append(countries[country])
data.insert(loc=0, column='hashcode', value=hash_column)

# Добавление колонки "процент выбросов от добычи нефти в стране за год" для классификации
procent_oil = []
oils = []
totals = []
for oil in data['Oil']:
    oils.append(oil)
for total in data['Total']:
    totals.append(total)
for i in range(len(oils)):
    procent_oil.append(math.ceil(oils[i]/totals[i]*100))
data.insert(loc=0, column='procent oil', value=procent_oil)

#------ Дерево решений для регрессии ------#

# ЗАДАЧА: предсказание общего объема выбросов СО2 страной за определённый год.

# Необходимые признаки для дерева регрессии
features_for_regr = data[['Year', 'hashcode']]
# Задача дерева регрессии
task_regr = data['Total']

# Разделение данных на обучающий и тестовый наборы для регрессии
X_train_r, X_test_r, \
y_train_r, y_test_r = train_test_split(features_for_regr, task_regr, test_size=0.01, random_state=250)

# Создание и обучение дерева решений для регрессии
model_regr = DecisionTreeRegressor(random_state=250)
model_regr.fit(X_train_r, y_train_r)

# Предсказание на тестовом наборе для регрессии
y_pred_r = model_regr.predict(X_test_r)

# Точечная оценка модели
score_r = model_regr.score(X_test_r, y_test_r)
print("\n\nТочность дерева регрессии:", score_r)

# Важности признаков для дерева регрессии
imp_regr = model_regr.feature_importances_

print("Важность признаков для дерева регрессии: ")
print("Важность 'Year':", imp_regr[0])
print("Важность 'hashcode':", imp_regr[1], "\n\n")

#------ Дерево решений для классификации ------#

# ЗАДАЧА: предсказание процента выбросов СО2 от добычи нефти страной за определённый год
#         с учётом общего объёма выбросов за год (или: какая часть выбросов придётся на добычу нефти).

# Необходимые признаки для дерева классификации
features_for_class = data[['Total', 'Year', 'hashcode']]
# Задача дерева классификации
task_class = data['procent oil']

# Разделение данных на обучающий и тестовый наборы для классификации
X_train_c, X_test_c, \
y_train_c, y_test_c = train_test_split(features_for_class, task_class, test_size=0.01, random_state=250)

# Создание и обучение дерева решений для классификации
model_class = DecisionTreeClassifier(random_state=250)
model_class.fit(X_train_c, y_train_c)

# Предсказание на тестовом наборе для классификации
y_pred_c = model_class.predict(X_test_c)

# Точечная оценка модели
score_c = model_class.score(X_test_c, y_test_c)
print("Точность дерева классификации:", score_c)

# Важности признаков для дерева классификации
imp_class = model_class.feature_importances_

print("Важность признаков для дерева классификации: ")
print("Важность 'Total':", imp_class[0])
print("Важность 'Year':", imp_class[1])
print("Важность 'hashcode':", imp_class[2])