import pandas as pd
import numpy as np
import rectools as rt
from rectools.models import PopularModel, UserKNNModel
from rectools.dataset import Dataset
from rectools.metrics import precision_at_k, recall_at_k, map_at_k
from sklearn.model_selection import train_test_split

# Пути к файлам
DATA_PATH = "dataset/"
TRAIN_FILE = DATA_PATH + "train.csv"
TEST_FILE = DATA_PATH + "test.csv"
SUBMISSION_FILE = "submission.csv"

# Загрузка данных
train_df = pd.read_csv(TRAIN_FILE)
test_df = pd.read_csv(TEST_FILE)

# Разделение данных на train/val
train_data, val_data = train_test_split(train_df, test_size=0.2, random_state=42)

# Создание датасета
dataset = Dataset.construct(train_data, user_col="user_id", item_col="item_id", feedback_col="rating")
val_dataset = Dataset.construct(val_data, user_col="user_id", item_col="item_id", feedback_col="rating")

# Инициализация моделей
pop_model = PopularModel()
pop_model.fit(dataset)

knn_model = UserKNNModel(K=10, similarity="cosine")
knn_model.fit(dataset)

# Функция предсказания рекомендаций
def predict(model):
    user_ids = test_df["user_id"].unique()
    recommendations = model.recommend(user_ids, dataset, k=10)  # Топ-10 рекомендаций
    return recommendations

# Оценка моделей на валидации
def evaluate_model(model):
    user_ids = val_data["user_id"].unique()
    recs = model.recommend(user_ids, dataset, k=10)
    precision = precision_at_k(val_dataset, recs, k=10)
    recall = recall_at_k(val_dataset, recs, k=10)
    map_score = map_at_k(val_dataset, recs, k=10)
    print(f"Precision@10: {precision:.4f}, Recall@10: {recall:.4f}, MAP@10: {map_score:.4f}")

# Сохранение предсказаний в CSV
def save_predictions(predictions, filename=SUBMISSION_FILE):
    predictions.to_csv(filename, index=False)
    print(f"Predictions saved to {filename}")

# Запуск
if __name__ == "__main__":
    print("Evaluating Popular Model...")
    evaluate_model(pop_model)
    
    print("Evaluating UserKNN Model...")
    evaluate_model(knn_model)
    
    print("Generating final predictions...")
    preds = predict(knn_model)  # Используем UserKNN для финальных рекомендаций
    save_predictions(preds)