import pathlib
import warnings
import sys
sys.path.append(".")
sys.path.append("..")
from src.model_manager.lgbm_manager import LightGBMModelManager  # noqa
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
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 scipy.special import rel_entr
import pandas as pd


def volatility_score(dfrace: pd.DataFrame, topN: int = 3):
    dfrace["odds_rate"] = 0.8/dfrace["odds"]
    dfrace["odds_rate"] /= dfrace["odds_rate"].sum()
    odds_fav_topN = dfrace.iloc[dfrace["favorite"].values.argsort(
    )]["odds_rate"]
    odds_fav_topN_adjust = odds_fav_topN.head(
        topN).tolist() + [1-odds_fav_topN.head(topN).sum()]

    odds_label_topN = dfrace.iloc[dfrace["label"].values.argsort(
    )]["odds_rate"]
    odds_label_topN_adjust = odds_label_topN.head(
        topN).tolist() + [1-odds_label_topN.head(topN).sum()]

    return rel_entr(odds_label_topN_adjust, odds_fav_topN_adjust).sum()

# 2023年高松宮記念
raceId = "202307020611"
dfp = df[df["raceId"].isin([raceId])]
dfp["odds_rate"] = 0.8/dfp["odds"]
display(dfp.set_index("label")[
        ["odds_rate", "odds", "favorite"]].sort_index().T)
for num in range(1, 4):
    print(f"{num}着まで: {volatility_score(dfp,num) = }")

# 2023年秋華賞
raceId2 = "202308020511"
dfp = df[df["raceId"].isin([raceId2])]
dfp["odds_rate"] = 0.8/dfp["odds"]
display(dfp.set_index("label")[
        ["odds_rate", "odds", "favorite"]].sort_index().T)
for num in range(1, 4):
    print(f"{num}着まで: {volatility_score(dfp,num) = :.17f}")

import tqdm

score_dict = {
    g: {
        f"label{num}": volatility_score(dfg, num)
        for num in range(1, 4)
    }
    for g, dfg in tqdm.tqdm(df[["raceId", "favorite", "label", "odds"]].groupby("raceId"))
}

idf = pd.DataFrame.from_dict(
    score_dict, orient="index").reset_index(names="raceId")
idf = idf[~idf["raceId"].str[:4].isin(
    [str(s) for s in range(2000, 2010)])].reset_index(drop=True)

idf.describe()

dfrefund = data_loader.load_racerefund()
bets = ["tan", "uren", "utan", "sanfuku", "santan"]
dflist = {
    g: dfg.set_index("bet")["refund"].T.to_dict()
    for g, dfg in tqdm.tqdm(dfrefund[dfrefund["bet"].isin(bets)].groupby("raceId"))
}
dfr = pd.DataFrame.from_dict(dflist, orient="index")
for b in bets:
    dfr[b] = pd.to_numeric(dfr[b], downcast="integer")

idf2 = pd.merge(idf, dfr.reset_index(names="raceId"), on="raceId", how="left")
idf2.sort_values("santan")

target_label = "label3"
idf2[idf2[target_label] >= idf2[target_label].quantile(
    0.9)][bets+[target_label]].describe().T.convert_dtypes()

target_label = "label3"
idf2[idf2[target_label] <= idf2[target_label].quantile(
    0.1)][bets+[target_label]].describe().T.convert_dtypes()

import numpy as np
import lightgbm as lgbm
df["odds_rate"] = 0.8/df["odds"]
df["odds_rate"] /= df["raceId"].map(df[["raceId", "odds_rate"]
                                       ].groupby("raceId")["odds_rate"].sum().to_dict())

import numpy as np
import pandas as pd
import lightgbm as lgbm


def kl_divergence_metric(y_pred, train_data):
    dflabels = train_data.get_label()
    labels = pd.DataFrame()
    labels["label"] = dflabels.astype(int)
    labels["odds_rate"] = dflabels-labels["label"]
    # y_pred = np.clip(y_pred, 1e-15, 1 - 1e-15)

    group_info = train_data.get_group()

    topN = 1
    start = 0
    gIdList = []
    for num, group_size in enumerate(group_info):
        # end = start + group_size
        # group_preds = y_pred[start:end]/0.1
        # y_pred[start:end] = np.exp(group_preds) / np.sum(np.exp(group_preds))
        # start = end
        gIdList += [num]*group_size

    labels["raceId"] = gIdList

    labels["preds"] = y_pred
    labels["rank"] = labels.groupby("raceId")["preds"].rank(
        ascending=False, method="first")
    labels["label"] = labels.groupby("raceId")["label"].rank(method="first")

    target = "odds_rate"
    labels["p_sum"] = labels["raceId"].map(labels[~labels["label"].isin(
        list(range(topN+1)))].groupby("raceId")[target].sum())
    labels["q_sum"] = labels["raceId"].map(labels[~labels["rank"].isin(
        list(range(topN+1)))].groupby("raceId")[target].sum())

    datalist = labels.set_index(["raceId", "label"])[target].to_dict()
    labels["proba"] = labels.apply(
        lambda row: datalist[(row["raceId"], row["rank"])], axis=1)

    kl_div = np.sum(labels["proba"] *
                    np.log(labels["proba"] / labels["odds_rate"]))

    return 'kl_divergence', kl_div, False


def kl_divergence_objective(y_pred, dataset: lgbm.Dataset):
    dflabels = dataset.get_label()
    labels = pd.DataFrame()
    labels["label"] = dflabels.astype(int)
    labels["odds_rate"] = dflabels-labels["label"]
    # y_pred = np.clip(y_pred, 1e-15, 1 - 1e-15)

    group_info = dataset.get_group()

    topN = 3
    start = 0
    gIdList = []
    for num, group_size in enumerate(group_info):
        # end = start + group_size
        # group_preds = y_pred[start:end]/0.1
        # y_pred[start:end] = np.exp(group_preds) / np.sum(np.exp(group_preds))
        # start = end
        gIdList += [num]*group_size

    labels["raceId"] = gIdList

    labels["preds"] = y_pred
    labels["rank"] = labels.groupby("raceId")["preds"].rank(
        ascending=False, method="first")
    labels["label"] = labels.groupby("raceId")["label"].rank(method="first")

    target = "odds_rate"
    labels["p_sum"] = labels["raceId"].map(labels[~labels["label"].isin(
        list(range(topN+1)))].groupby("raceId")[target].sum())
    labels["q_sum"] = labels["raceId"].map(labels[~labels["rank"].isin(
        list(range(topN+1)))].groupby("raceId")[target].sum())
    labels["p_sum"] = labels["p_sum"].fillna(0)
    labels["q_sum"] = labels["q_sum"].fillna(1)

    datalist = labels.set_index(["raceId", "label"])[target].to_dict()
    labels["proba"] = labels.apply(
        lambda row: datalist[(row["raceId"], row["rank"])], axis=1)

    labels["threash"] = labels["rank"] <= topN

    labels["grad"] = labels["threash"]*(-labels["proba"]/labels[target]) + \
        (1-labels["threash"])*(-labels["p_sum"]/labels["q_sum"])

    labels["hessian"] = labels["threash"]*(labels["proba"]/labels[target].pow(2)) + \
        (1-labels["threash"])*(labels["p_sum"]/labels["q_sum"].pow(2))

    return labels["grad"].values, labels["hessian"].values

# 説明変数にするカラム
feature_columns = [
    'distance',
    'number',
    'boxNum',
    'age',
    'jweight',
    'weight',
    'gl',
    'race_span',
    "raceGrade",  # グレード情報を追加
] + dataPreP.encoding_columns

# 血統情報を追加
feature_columns += ["stallionId", "breedId", "bStallionId", "b2StallionId"]

# 目的変数用のカラム
label_column = "label_col"
df[label_column] = df["label"] + df["odds_rate"]

dftrain, dfvalid, dftest = df[df["raceId"].str[:4].isin([str(y) for y in range(2014, 2020)])], df[df["raceId"].str[:4].isin(
    ["2020"])], df[df["raceId"].str[:4].isin(["2021"])]
train_data = lgbm.Dataset(
    dftrain[feature_columns],
    label=dftrain[label_column],
    group=dftrain["raceId"].drop_duplicates().map(
        dftrain.groupby("raceId")["raceId"].count().to_dict()).values,
    # init_score=np.random.rand(len(dftrain))
)
valid_data = lgbm.Dataset(
    dfvalid[feature_columns],
    label=dfvalid[label_column],
    group=dfvalid["raceId"].drop_duplicates().map(
        dfvalid.groupby("raceId")["raceId"].count().to_dict()).values,
    # init_score=np.random.rand(len(dfvalid))
)
test_data = lgbm.Dataset(
    dftest[feature_columns],
    label=dftest[label_column],
    group=dftest["raceId"].drop_duplicates().map(
        dftest.groupby("raceId")["raceId"].count().to_dict()).values,
    # init_score=np.random.rand(len(dftest))
)

# 学習用パラメータ（ここでは適当に設定しておく）
params = {
    'boosting_type': 'gbdt',
    # 二値分類
    'objective': kl_divergence_objective,
    'verbose': 1,
    'seed': 77777,
    'learning_rate': 0.1,
    # "n_estimators": 100
}

# モデル学習
model = lgbm.train(
    params,
    train_data,
    num_boost_round=1000,
    feval=kl_divergence_metric,
    valid_sets=[train_data, valid_data],
    callbacks=[
        lgbm.early_stopping(stopping_rounds=25, verbose=True,),
        lgbm.log_evaluation(10 if True else 0)
    ],
)

df["pred_proba"] = model.predict(df[feature_columns])
df["pred_rank"] = df[["raceId", "pred_proba"]].groupby(
    "raceId")["pred_proba"].rank(ascending=False)

raceId = "202107010308"
df[df["raceId"].isin([raceId])][
    [
        "raceId", "label", "favorite",
        "odds", "pred_rank", "pred_proba"
    ]
].sort_values("pred_rank")

dftrain["pred_proba"] = model.predict(dftrain[feature_columns])
dftrain["pred_rank"] = dftrain.groupby(
    "raceId")["pred_proba"].rank(ascending=False)

dfvalid["pred_proba"] = model.predict(dfvalid[feature_columns])
dfvalid["pred_rank"] = dfvalid.groupby(
    "raceId")["pred_proba"].rank(ascending=False)

dftest["pred_proba"] = model.predict(dftest[feature_columns])
dftest["pred_rank"] = dftest.groupby(
    "raceId")["pred_proba"].rank(ascending=False)

q = dftrain.groupby("raceId")["pred_proba"].max().quantile(0.0)

display(
    pd.concat(
        [
            dftest[dftest["pred_rank"].isin([p]) & (
                dftest["pred_proba"] > q)]["label"].value_counts().sort_index().to_frame(name=p).T
            for p in range(1, 6)
        ]
    )
)

display(
    pd.concat(
        [
            dftest[dftest["pred_rank"].isin([p]) & (dftest["pred_proba"] > q)]["label"].value_counts().sort_index().to_frame(
                name=p).T/(dftest["pred_rank"].isin([p]) & (dftest["pred_proba"] > q)).sum()
            for p in range(1, 6)
        ]
    )
)

for mode, dfp in zip(["valid", "test"], [dfvalid, dftest]):
    dffilter = dfp["pred_rank"].isin([1]) & dfp["label"].isin(
        [1]) & (dfp["pred_proba"] > q)
    profit = dfp[dffilter]["odds"].sum()
    print(mode, profit, profit/(dfp["pred_rank"].isin([1]) & (dfp["pred_proba"] > q)).sum(),
          dffilter.sum(), dffilter.sum()/(dfp["pred_rank"].isin([1]) & (dfp["pred_proba"] > q)).sum())