import pathlib
import warnings
import sys
sys.path.append(".")
sys.path.append("../..")
from src.core.meta.bet_name_meta import BetName  # noqa
from src.data_manager.preprocess_tools import DataPreProcessor  # noqa
from src.data_manager.data_loader import DataLoader  # noqa

warnings.filterwarnings("ignore")

root_dir = pathlib.Path(".").absolute().parent.parent
dbpath = root_dir / "data" / "keibadata.db"
start_year = 2000  # DBが持つ最古の年を指定
split_year = 2014  # 学習対象期間の開始年を指定
target_year = 2019  # テスト対象期間の開始年を指定
end_year = 2023  # テスト対象期間の終了年を指定 (当然DBに対象年のデータがあること)

# 各種インスタンスの作成
data_loader = DataLoader(
    start_year,
    end_year,
    dbpath=dbpath  # dbpathは各種環境に合わせてパスを指定してください。絶対パス推奨
)

dataPreP = DataPreProcessor(
    # キャッシュ機能を使用する場合にTrueを指定。デフォルト：True
    use_cache=True,
    cache_dir=pathlib.Path("./data")
)

df = data_loader.load_racedata()
dfblood = data_loader.load_horseblood()

df = dataPreP.exec_pipeline(
    df,
    dfblood,
    blood_set=["s", "b", "bs", "bbs", "ss", "sss", "ssss", "bbbs"],
    lagN=5
)

from src.data_manager.dataset_tools import DatasetGenerator

dataset_generator = DatasetGenerator(split_year, target_year, end_year)


target_category = [["stallionId"], ["breedId"],
                   ["bStallionId"], ["b2StallionId"]]
target_sub_category = ["field", "dist_cat"]

cat_list = ["field", "place", "dist_cat",
            "condition", "inoutside", "direction", "horseId", 'weather',]
for cat in cat_list:
    df[cat] = df[cat].astype("category")

dataset_mapping = dataset_generator.make_dataset_mapping(df)
dataset_generator.calcurate_category_winrate(
    df, dataset_mapping, target_category, target_sub_category, [1])


# セカンドモデルの特徴量
feature_columns = [
    'distance',
    'number',
    'boxNum',
    'odds',
    'favorite',
    'age',
    'jweight',
    'weight',
    'gl',
    'race_span',
    "raceGrade",
] + dataPreP.encoding_columns + [
    "stallionId", "breedId", "bStallionId", "b2StallionId"
] + ['winR_stallion', 'winR_breed', 'winR_bStallion', 'winR_b2Stallion']

import pandas as pd

alllist = []
for key, dataset in dataset_mapping.items():
    dflist = []
    for mode in ["train", "valid", "test"]:
        idf = dataset.__dict__[
            mode][feature_columns].isna().sum(axis=0).rename(mode)
        dflist += [idf[idf > 0]]
    alllist += [pd.concat(dflist,
                          axis=1).rename(columns=lambda x: f"{key}_{x}").T]
display(pd.concat(alllist))

for dataset in dataset_mapping.values():
    for mode in ["train", "valid", "test"]:
        idf = dataset.__dict__[mode]
        idf["race_span_fill"] = idf[["age", "race_span"]].apply(lambda row: int(
            row["age"])*365 if pd.isna(row["race_span"]) else row["race_span"], axis=1)
        dataset.__dict__[mode] = idf.copy()

feature_columns += ["race_span_fill"]
feature_columns.remove("race_span")

alllist = []
for key, dataset in dataset_mapping.items():
    dflist = []
    for mode in ["train", "valid", "test"]:
        idf = dataset.__dict__[
            mode][feature_columns].isna().sum(axis=0).rename(mode)
        dflist += [idf[idf > 0]]
    alllist += [pd.concat(dflist,
                          axis=1).rename(columns=lambda x: f"{key}_{x}").T]
display(pd.concat(alllist))

dataset.train["race_span_fill"].describe()

# 念のため欠損値補完前の分布を確認
display(df["race_span"].describe().to_frame(name="補完前").T)

dftype = dataset.train[feature_columns].dtypes.to_frame("TypeInfo")
dftype["TypeName"] = dftype["TypeInfo"].apply(lambda x: x.name)
category_list = dftype[dftype["TypeName"].isin(["category"])].index.tolist()
numerous_list = dftype[~dftype["TypeName"].isin(["category"])].index.tolist()

numerous_list

# よって、標準化する特徴量をまとめる
# レース全体で標準化する特徴量
sd_by_all_race = ["distance", "age"]

# レースごとに標準化する特徴量
sd_by_a_race = ["number", "boxNum", "odds", "favorite", "jweight", "weight"]

for key, dataset in dataset_mapping.items():
    for sd_col in sd_by_all_race + ["race_span_fill"]:
        sdcol = sd_col.split("_fill")[0]
        idf = pd.concat([dataset.train[sdcol], dataset.valid[sdcol]])

        mean_train = dataset.train[sdcol].mean()
        mean_tr_va = idf.mean()
        std_train = dataset.train[sdcol].std()
        std_tr_va = idf.std()

        dataset.train[sd_col +
                      "_dev"] = (dataset.train[sd_col]-mean_train)/std_train
        dataset.valid[sd_col +
                      "_dev"] = (dataset.valid[sd_col]-mean_train)/std_train
        dataset.test[sd_col +
                     "_dev"] = (dataset.test[sd_col]-mean_tr_va)/std_tr_va

    for mode in ["train", "valid", "test"]:
        dataset.__dict__[mode]["gl"+"_dev"] = (dataset.__dict__[
            mode]["weight"]+dataset.__dict__[mode]["gl"])/dataset.__dict__[mode]["weight"]

    for sdcol in sd_by_a_race:
        for mode in ["train", "valid", "test"]:
            mean_dict = dataset.__dict__[mode][["raceId", sdcol]].groupby("raceId")[
                sdcol].mean()
            std_dict = dataset.__dict__[mode][["raceId", sdcol]].groupby("raceId")[
                sdcol].std()
            dataset.__dict__[mode][sdcol+"_dev"] = (dataset.__dict__[mode][sdcol] - dataset.__dict__[
                mode]["raceId"].map(mean_dict))/dataset.__dict__[mode]["raceId"].map(std_dict)
            # 標準偏差が0になるパターンもあり得るので、その場合は全て0とする
            dataset.__dict__[mode][sdcol+"_dev"].fillna(0, inplace=True)

dataset.train[[n+"_dev" for n in numerous_list[:-5] +
               numerous_list[-1:]] + numerous_list[-5:-1]]

# まずは、カテゴリ特徴量の種類を確認していきましょう。
category_list

# すべての期間でカテゴリ数が変わらない定常なもの
stationary_category = ["raceGrade", "place_en",
                       "field_en", "sex_en", "condition_en", "dist_cat_en"]

# カテゴリ数に変動があるもの
non_stationary_category = ["jockeyId", "teacherId", "horseId",
                           "stallionId", "breedId", "bStallionId", "b2StallionId"]

for key, dataset in dataset_mapping.items():
    for cat in non_stationary_category:
        cat_unique = dataset.train[cat].unique()

        # カテゴリの振り番のマッピングを作成
        cat_map = {c: n for n, c in enumerate(cat_unique)}
        # 学習・検証・テストデータにマッピングを実施
        dataset.train[cat+"_en"] = dataset.train[cat].map(cat_map)
        dataset.valid[cat+"_en"] = dataset.valid[cat].map(cat_map)
        dataset.test[cat+"_en"] = dataset.test[cat].map(cat_map)
        # 割り振れなかったものを、Unkownのカテゴリとしてセット
        dataset.valid[cat+"_en"].fillna(len(cat_map), inplace=True)
        dataset.test[cat+"_en"].fillna(len(cat_map), inplace=True)
        # データ型をintにしておく
        dataset.train[cat+"_en"] = dataset.train[cat+"_en"].astype("int64")
        dataset.valid[cat+"_en"] = dataset.valid[cat+"_en"].astype("int64")
        dataset.test[cat+"_en"] = dataset.test[cat+"_en"].astype("int64")

    for cat in stationary_category:
        # データ型をintにしておく
        dataset.train[cat] = dataset.train[cat].astype("int64")
        dataset.valid[cat] = dataset.valid[cat].astype("int64")
        dataset.test[cat] = dataset.test[cat].astype("int64")

dataPreP.encoding_columns

# 特徴量のカラム名が変わったので再定義しておく
# 量的変数の特徴量
num_feas = [
    'distance_dev',
    'number_dev',
    'boxNum_dev',
    'odds_dev',
    'favorite_dev',
    'age_dev',
    'jweight_dev',
    'weight_dev',
    'gl_dev',
    'race_span_fill_dev',
] + ['winR_stallion', 'winR_breed', 'winR_bStallion', 'winR_b2Stallion']
# 質的変数の特徴量
cat_feas = [
    'place_en',
    'field_en',
    'sex_en',
    'condition_en',
    'jockeyId_en',
    'teacherId_en',
    'dist_cat_en',
    'horseId_en',
    "raceGrade", "stallionId_en", "breedId_en", "bStallionId_en", "b2StallionId_en"
]

from torch.utils.data import Dataset
import torch


# Datasetクラスを継承してカスタムDatasetクラスを作成
class CustomKaibaAIDataset(Dataset):
    def __init__(self, dfnum: pd.DataFrame, dfcat: pd.DataFrame, dflabel: pd.Series) -> None:
        self.numerous = dfnum
        self.cat = dfcat
        self.label = dflabel

    def __len__(self):
        return len(self.label)

    def __getitem__(self, index):
        num_fea = torch.tensor(self.numerous.loc[index], dtype=torch.float32)
        # カテゴリ特徴量は一つ一つベクトル埋め込み層に突っ込むので、特徴量ごとに分けてtensor化しておく
        cat_feas = torch.tensor(
            [self.cat[c].loc[index] for c in self.cat.columns], dtype=torch.float32)

        label = torch.tensor(self.label[index], dtype=torch.float32)
        return num_fea, cat_feas, label

for key, dataset in dataset_mapping.items():
    # 教師データの作成：1着なら1でそれ以外なら0
    dataset.train["object_label"] = dataset.train["label"].isin([
                                                                1]).astype(int)
    dataset.train_dataset = CustomKaibaAIDataset(
        dataset.train[num_feas],
        dataset.train[cat_feas],
        dataset.train["object_label"]
    )
    dataset.valid["object_label"] = dataset.valid["label"].isin([
                                                                1]).astype(int)
    dataset.valid_dataset = CustomKaibaAIDataset(
        dataset.valid[num_feas],
        dataset.valid[cat_feas],
        dataset.valid["object_label"]
    )
    dataset.test["object_label"] = dataset.test["label"].isin([1]).astype(int)
    dataset.test_dataset = CustomKaibaAIDataset(
        dataset.test[num_feas],
        dataset.test[cat_feas],
        dataset.test["object_label"]
    )
    # カテゴリの埋め込みを行うためにカテゴリ数を指定しておく
    dataset.cat_num_list = [dataset.train[cat].nunique()+1 for cat in cat_feas]

import pickle

cache_dir = pathlib.Path("./data")
with open(cache_dir / "dataset_mapping.pkl", "wb") as f:
    pickle.dump(dataset_mapping, f)