from torch.utils.data import Dataset  # type: ignore
import torch  # type: ignore
import pandas as pd  # type: ignore
import pathlib
import pickle
import warnings
import sys
sys.path.append(".")
sys.path.append("../..")

from src.data_manager.dataset_tools import DatasetDict  # noqa

# 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


cache_dir = pathlib.Path("./data")


with open(cache_dir / "dataset_mapping.pkl", "rb") as f:

    dataset_mapping: dict[str, DatasetDict] = pickle.load(f)

# モデル作成で使用する特徴量
# 量的変数の特徴量
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"
]

import torch.nn as nn  # type: ignore
import torch  # type: ignore
import math


class KeibaAISecondModel(nn.Module):
    def __init__(self, cat_num_list: list[int], numerous_feature_num: int) -> None:
        super(KeibaAISecondModel, self).__init__()

        cat_embed_list = []
        self.embed_num_list = []
        for cat_num in cat_num_list:
            embed_num = round(math.sqrt(cat_num))
            self.embed_num_list += [embed_num]
            embed_layer = nn.Embedding(
                cat_num, embed_num, padding_idx=cat_num-1)
            cat_embed_list += [embed_layer]
        # カテゴリのベクトル埋め込み用レイヤー
        self.cat_embed_list = nn.ModuleList(cat_embed_list)
        self.cat_num_list = cat_num_list

        # 適当に2層ぐらいのNNを作る。
        # 1つ目は入力層 → 隠れ層
        self.linear_hidden_layer = nn.Linear(
            numerous_feature_num+sum(self.embed_num_list),
            1024
            # 128
        )
        # 2つ目は隠れ層 → 隠れ層
        self.linear_hidden_layer2 = nn.Linear(1024, 256)
        # 3つ目は隠れ層 → 隠れ層
        self.linear_hidden_layer3 = nn.Linear(256, 128)
        # 4つ目は隠れ層 → 出力層
        self.linear_hidden_layer4 = nn.Linear(128, 1)
        # 最後に出力層 → sigmoid変換
        self.sigmoid = nn.Sigmoid()

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x_num, x_cat_list):
        x_cat_list = x_cat_list.int()
        cat_embed_list = [self.cat_embed_list[idx](
            x_cat) for idx, x_cat in enumerate(x_cat_list.T)]
        cat_embed_list = torch.cat(
            cat_embed_list, dim=1 if len(x_cat_list.shape) > 1 else 0)
        x = torch.cat([x_num, cat_embed_list],
                      dim=1 if len(x_num.shape) > 1 else 0)

        x = self.linear_hidden_layer(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear_hidden_layer2(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear_hidden_layer3(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear_hidden_layer4(x)
        x = self.sigmoid(x)
        return x

from torchviz import make_dot  # type: ignore
from torch.utils.data import DataLoader  # type: ignore
dataset = dataset_mapping["2019first"]

dataset_loader = DataLoader(dataset.train_dataset, batch_size=1, shuffle=True)
for batch_data1, batch_data2, batch_label in dataset_loader:
    break

model = KeibaAISecondModel(dataset.cat_num_list, len(num_feas))
model.eval()
y = model(batch_data1[0], batch_data2[0])
dot = make_dot(y, params=dict(model.named_parameters()))

# SVG形式でモデル構造の出力
dot.render('KeibaAISecondModel_Graph', format="svg")

# 損失関数
loss_fn = nn.BCELoss()

# オプティマイザ
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

import tqdm  # type: ignore
import numpy as np  # type: ignore

# チュートリアルなので、dataset_mappingの内ひとつだけのdatasetを対象にモデルを学習してみる
target_dataset = dataset_mapping["2022second"]

# デバイスの設定（GPUが利用可能な場合）
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# device = "cpu"

# モデルの初期化
model = KeibaAISecondModel(
    target_dataset.cat_num_list, len(num_feas)).to(device)
# 損失関数
loss_fn = nn.BCELoss()
# オプティマイザ
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=0.01)


# 1.DataLoaderの作成
train_dataloader = DataLoader(
    target_dataset.train_dataset, batch_size=4096, shuffle=True)
valid_dataloader = DataLoader(target_dataset.valid_dataset, batch_size=len(
    target_dataset.valid), shuffle=False)
test_dataloader = DataLoader(
    target_dataset.test_dataset, batch_size=len(target_dataset.test), shuffle=False)


loss_list_all = []
# 学習ループ
num_epochs = 3
for epoch in range(num_epochs):
    model.train()  # 訓練モード
    loss_list = []
    with tqdm.tqdm(total=len(train_dataloader), desc=f"Train Epoch {epoch+1}/{num_epochs}. loss=None") as pbar:
        for num_data, cat_data_list, labels in train_dataloader:
            # データとラベルをGPUに転送
            num_data, cat_data_list, labels = num_data.to(
                device), cat_data_list.to(device), labels.to(device)

            # 順伝播
            optimizer.zero_grad()  # 勾配を初期化
            outputs = model(num_data, cat_data_list)

            # 損失の計算
            loss = loss_fn(outputs.squeeze(), labels)

            # 逆伝播
            loss.backward()

            # 最適化
            optimizer.step()

            # ロスを記録
            loss_list += [loss.item()]
            pbar.update()
            pbar.set_description(
                desc=f"Train Epoch {epoch+1}/{num_epochs}. loss={np.mean(loss_list):.6f}")
    loss_list_all += [loss_list]

    model.eval()  # 推論モード
    # 検証データのロス確認
    for num_data, cat_data_list, labels in valid_dataloader:
        # データとラベルをGPUに転送
        num_data, cat_data_list, labels = num_data.to(
            device), cat_data_list.to(device), labels.to(device)

        # 順伝播
        outputs = model(num_data, cat_data_list)
        # 損失の計算
        loss = loss_fn(outputs.squeeze(), labels)

    # テストデータのロス確認
    for num_data, cat_data_list, labels in test_dataloader:
        # データとラベルをGPUに転送
        num_data, cat_data_list, labels = num_data.to(
            device), cat_data_list.to(device), labels.to(device)

        # 順伝播
        outputs_test = model(num_data, cat_data_list)
        # 損失の計算
        loss_test = loss_fn(outputs_test.squeeze(), labels)

    print(f"valid: {loss.item():.6f}, test: {loss_test.item():.6f}")

import matplotlib.pyplot as plt  # type: ignore
import japanize_matplotlib  # type: ignore

for idx, loss_list in enumerate(loss_list_all):
    y = np.array(loss_list).reshape(-1)
    x = list(range(len(y)*idx, len(y)*(idx+1)))
    plt.plot(x, y, label=f"Epoch {idx+1}")
plt.grid(ls=":")
plt.legend()
plt.show()

dfvalid = target_dataset.valid.copy()
dfvalid["pred_proba"] = outputs.squeeze().to("cpu").detach().numpy()
dfvalid["pred_rank"] = dfvalid.groupby(
    "raceId")["pred_proba"].rank(ascending=False).astype(int)

dftest = target_dataset.test.copy()
dftest["pred_proba"] = outputs_test.squeeze().to("cpu").detach().numpy()
dftest["pred_rank"] = dftest.groupby(
    "raceId")["pred_proba"].rank(ascending=False).astype(int)

print("pred_rank=1に賭けたときの着順の分布")
display(
    pd.concat(
        [
            dfvalid[dfvalid["pred_rank"].isin([1])]["label"].value_counts(
            ).sort_index().to_frame(name="valid").T,
            dfvalid[dfvalid["pred_rank"].isin([1])]["label"].value_counts().sort_index(
            ).to_frame(name="累積分布").T.cumsum(axis=1)/dfvalid["pred_rank"].isin([1]).sum()
        ]
    )
)
display(
    pd.concat(
        [
            dftest[dftest["pred_rank"].isin([1])]["label"].value_counts(
            ).sort_index().to_frame(name="test").T,
            dftest[dftest["pred_rank"].isin([1])]["label"].value_counts().sort_index(
            ).to_frame(name="累積分布").T.cumsum(axis=1)/dftest["pred_rank"].isin([1]).sum()
        ]
    )
)

print("pred_rank=1に賭けたときの人気の分布")
display(
    pd.concat(
        [
            dfvalid[dfvalid["pred_rank"].isin([1])]["favorite"].value_counts(
            ).sort_index().to_frame(name="valid").T,
            dfvalid[dfvalid["pred_rank"].isin([1])]["favorite"].value_counts().sort_index(
            ).to_frame(name="累積分布").T.cumsum(axis=1)/dfvalid["pred_rank"].isin([1]).sum()
        ]
    )
)

display(
    pd.concat(
        [
            dftest[dftest["pred_rank"].isin([1])]["favorite"].value_counts(
            ).sort_index().to_frame(name="test").T,
            dftest[dftest["pred_rank"].isin([1])]["favorite"].value_counts().sort_index(
            ).to_frame(name="累積分布").T.cumsum(axis=1)/dftest["pred_rank"].isin([1]).sum()
        ]
    )
)

idfv = dfvalid[dfvalid["pred_rank"].isin([1]) & dfvalid["raceGrade"].ge(0)]
print(round(idfv[idfv["label"].isin([1])]["odds"].sum(), 1), len(
    idfv), round(idfv[idfv["label"].isin([1])]["odds"].sum(), 1)/len(idfv))
idfv = dftest[dftest["pred_rank"].isin([1]) & dftest["raceGrade"].ge(0)]
print(round(idfv[idfv["label"].isin([1])]["odds"].sum(), 1), len(
    idfv), round(idfv[idfv["label"].isin([1])]["odds"].sum(), 1)/len(idfv))

original_num_feas = [col.split("_dev")[0] for col in num_feas]
original_cat_feas = [col.split(
    "Id_en")[0] + "Name" if "Id_en" in col and ("jockey" not in col and "teacher" not in col) else col.split("_en")[0] for col in cat_feas]

import shap  # type: ignore

shap.initjs()

# backgroud dataの指定：ここは学習データから。100件超えてくると処理が遅くなるので注意
x_inputs = target_dataset.train.sample(150, random_state=777)[
    num_feas+cat_feas]

model.to("cpu")
model.eval()


def func(dfvalid):
    # model.to(device)
    with torch.no_grad():
        x1 = torch.tensor(dfvalid[:, :len(num_feas)]).float()
        x2 = torch.tensor(dfvalid[:, len(num_feas):]).int()
        # x1.to(device)
        # x2.to(device)
        output = model(x1, x2).numpy().reshape(-1)
    return output


explainer = shap.KernelExplainer(
    func, x_inputs.values)

# 分析対象のデータ。ここで選択したデータのSHAP値を計算する。特徴量等が多い場合も恐らく処理に時間がかかります
x_test = target_dataset.valid.iloc[1:].sample(99, random_state=777)
x_test = pd.concat([target_dataset.valid.iloc[[0]], x_test])[num_feas+cat_feas]
shap_values = explainer.shap_values(x_test)

shap_values.shape, x_test.shape

x_test_org = target_dataset.valid.loc[x_test.index][original_num_feas +
                                                    original_cat_feas].copy()

shap.summary_plot(
    shap_values,
    x_test,
    feature_names=original_num_feas+original_cat_feas,
    max_display=10  # 表示されるグラフの大きさを調整します。あまりにも寄与が小さいものは非表示にされます。
)

shap.summary_plot(
    shap_values,
    x_test,
    feature_names=original_num_feas+original_cat_feas,
    plot_type="bar",
    max_display=10
)

# たとえば、インデックス0のデータに対するモデルの予測値 pred_proba は以下
print("モデルの予測値: \t\t\t", dfvalid.loc[x_test.iloc[0].name]["pred_proba"])

# 続いて、背景データのSHAP値の期待値 expected_value は以下
print("背景データのSHAP値の期待値: \t", explainer.expected_value)

# インデックス0のデータに対する全特徴量のSHAP値の和
print("全特徴量のSHAP値の和: \t\t", shap_values[0].sum())

# この３つの値は、以下の関係にある
print("\nモデルの予測値 = 背景データのSHAP値の期待値 + 全特徴量のSHAP値の和")
print("モデルの予測値: \t\t\t\t\t\t", dfvalid.loc[x_test.iloc[0].name]["pred_proba"])
print("背景データのSHAP値の期待値 + 全特徴量のSHAP値の和: \t",
      explainer.expected_value + shap_values[0].sum())

shap.force_plot(
    explainer.expected_value,
    shap_values[0],
    x_test_org.iloc[0],
)

shap.plots._waterfall.waterfall_legacy(
    explainer.expected_value,
    shap_values[0],
    x_test_org.iloc[0],
)

fav1_index = x_test_org.reset_index(names="org_idx")[x_test_org.reset_index()[
    "favorite"].isin([1])].sample(1).index.tolist()[0]
shap.plots._waterfall.waterfall_legacy(
    explainer.expected_value,
    shap_values[fav1_index],
    x_test_org.iloc[fav1_index],
)

max_index = shap_values.sum(axis=1).argmax()
shap.plots._waterfall.waterfall_legacy(
    explainer.expected_value,
    shap_values[max_index],
    x_test_org.iloc[max_index],
)

shap.force_plot(
    explainer.expected_value,
    shap_values,
    x_test_org,
)

model_dir = pathlib.Path("./model")
model_dir.mkdir(exist_ok=True)

# モデルのパラメータを保存
torch.save(model.state_dict(), model_dir / "tutrial_model_weights.pth")