from catboost import CatBoostClassifier
import pandas as pd
import numpy as np
from typing import Dict, Any

# Загрузка модели
modelBagOfHeroes = CatBoostClassifier()
modelBagOfHeroes.load_model("artifacts/model_bag_of_heroes.cbm")

# Загрузка порядка фич
def load_feature_order(path: str) -> list[str]:
    fo = pd.read_csv(path)
    first_col = fo.columns[0]
    return fo[first_col].tolist()

FEATURE_ORDER_BAG: list[str] = load_feature_order("artifacts/feature_order_bag_of_heroes.csv")

def build_bag_of_heroes_features(payload: Dict[str, Any]) -> pd.DataFrame:
    """
    Конвертирует payload в bag-of-heroes формат.

    payload содержит:
    - is_first_pick_radiant
    - r_h1, r_h2, r_h3, r_h4, r_h5
    - d_h1, d_h2, d_h3, d_h4, d_h5

    Возвращает DataFrame с колонками:
    - is_first_pick_radiant
    - radiant_hero_{1-145}
    - dire_hero_{1-145}
    """
    # Получаем героев из payload
    radiant_heroes = []
    dire_heroes = []

    for i in range(1, 6):
        r_hero = payload.get(f"r_h{i}", -1)
        d_hero = payload.get(f"d_h{i}", -1)

        if r_hero and r_hero != -1:
            radiant_heroes.append(int(r_hero))
        if d_hero and d_hero != -1:
            dire_heroes.append(int(d_hero))

    # Создаем словарь признаков
    features = {feat: 0 for feat in FEATURE_ORDER_BAG}

    # Устанавливаем is_first_pick_radiant
    features["is_first_pick_radiant"] = int(payload.get("is_first_pick_radiant", 0))

    # Устанавливаем бинарные признаки для героев Radiant
    for hero_id in radiant_heroes:
        feat_name = f"radiant_hero_{hero_id}"
        if feat_name in features:
            features[feat_name] = 1

    # Устанавливаем бинарные признаки для героев Dire
    for hero_id in dire_heroes:
        feat_name = f"dire_hero_{hero_id}"
        if feat_name in features:
            features[feat_name] = 1

    # Создаем DataFrame с правильным порядком колонок
    return pd.DataFrame([features], columns=FEATURE_ORDER_BAG)

def predict_bag_of_heroes(payload: Dict[str, Any]) -> Dict[str, float]:
    """
    Делает предсказание с использованием bag-of-heroes модели.

    Возвращает:
    {
        "radiant_win": вероятность победы Radiant (0-100),
        "dire_win": вероятность победы Dire (0-100)
    }
    """
    X = build_bag_of_heroes_features(payload)
    proba = modelBagOfHeroes.predict_proba(X)[0, 1]

    radiant_win = round(float(np.clip(proba * 100.0, 0.0, 100.0)))
    dire_win = 100.0 - radiant_win

    return {
        "radiant_win": radiant_win,
        "dire_win": dire_win
    }