import numpy as np
import random
from torch.utils.data import Dataset
import torch
import pandas as pd
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)


def set_seed(seed):
    random.seed(seed)  # Python標準の乱数シード
    np.random.seed(seed)  # NumPyの乱数シード
    torch.manual_seed(seed)  # PyTorchの乱数シード（CPU用）
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)  # CUDAの乱数シード（単一GPU用）
        torch.cuda.manual_seed_all(seed)  # 複数GPU用
    # 再現性のためにPyTorchの動作設定を変更
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# シードを固定
set_seed(42)

# モデル作成で使用する特徴量
# 量的変数の特徴量
num_feas = [
    'distance_dev',
    'number_dev',
    'boxNum_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"
]

for key, dataset in dataset_mapping.items():
    for mode in ["train", "valid", "test"]:
        idf: pd.DataFrame = dataset.__dict__[mode]
        raceId_map = idf.sort_values(
            ["raceDate", "raceId"])[["raceId"]].drop_duplicates(
                ignore_index=True).reset_index(
                    names="rst_idx").set_index("raceId")["rst_idx"].to_dict()
        idf.index = idf["raceId"].map(raceId_map).tolist()
        idf["odds_rate"] = 0.8/idf["odds"]
        idf["odds_rate"] /= idf["raceId"].map(
            idf[["raceId", "odds_rate"]].groupby("raceId")["odds_rate"].sum().to_dict())
        dataset_mapping[key].__dict__[mode] = idf.copy()

from typing import Literal


class CustomKaibaAIDatasetForMultiTask(Dataset):
    def __init__(self, idf: pd.DataFrame, num_feas: list[str], cat_feas: list[str]) -> None:
        self.static_length = 18
        idxlist = []
        for k, v in (self.static_length-idf.index.value_counts()).sort_index().to_dict().items():
            idxlist += [k]*v
        dfnum_pad = pd.DataFrame(0, index=idxlist, columns=num_feas)
        dfcat_pad = pd.DataFrame([dataset.cat_num_list] *
                                 len(idxlist), index=idxlist, columns=cat_feas)-1
        dfpad = pd.concat([dfnum_pad, dfcat_pad], axis=1)
        idf2 = pd.concat([idf, dfpad]).loc[list(
            range(idf.index.nunique()))]

        self.numerous = idf2[num_feas]
        self.cat = idf2[cat_feas]
        self.label = idf2["label"].fillna(19)
        self.odds_rate = idf2["odds_rate"].fillna(0)

    def __len__(self):
        return self.label.index.nunique()

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

        label = torch.tensor(self.label[index].values, dtype=torch.float32)
        odds_rate = torch.tensor(
            self.odds_rate[index].values, dtype=torch.float32)

        mask = torch.ones_like(label).masked_fill(label == 19, 0)
        return {
            "num_feas": num_feas,
            "cat_feas": cat_feas,
            "label": label,
            "odds_rate": odds_rate,
            "mask": mask
        }

idf = dataset.train
tensor_dataset = CustomKaibaAIDatasetForMultiTask(idf, num_feas, cat_feas)

from torch.utils.data import DataLoader
data_loader = DataLoader(tensor_dataset, batch_size=10, shuffle=False)
datas = next(iter(data_loader))

print(datas["num_feas"][0])

print("shape: \t\t\t[batch_size, horse_num, features_size]")
print("num_feas shape: \t", datas["num_feas"].shape)
print("cat_feas shape: \t", datas["cat_feas"].shape)
print("label shape: \t\t", datas["label"].shape)
print("odds_rate shape: \t", datas["odds_rate"].shape)

for key, dataset in dataset_mapping.items():
    dataset_mapping[key].train_dataset = CustomKaibaAIDatasetForMultiTask(
        dataset.train, num_feas, cat_feas)
    dataset_mapping[key].valid_dataset = CustomKaibaAIDatasetForMultiTask(
        dataset.valid, num_feas, cat_feas)
    dataset_mapping[key].test_dataset = CustomKaibaAIDatasetForMultiTask(
        dataset.test, num_feas, cat_feas)

import torch.nn as nn

if 1:
    class MultiTaskLoss(nn.Module):
        def __init__(self, l1=1/3, l2=1/3, l3=1/3, topn=1):
            super(MultiTaskLoss, self).__init__()
            self.l1 = l1
            self.l2 = l2
            self.l3 = l3
            self.topn = topn
            self.KLloss1 = nn.KLDivLoss(reduction="batchmean")
            self.KLloss2 = nn.KLDivLoss(reduction="batchmean")
            # self.MAE = nn.MSELoss()
            self.MAE = nn.L1Loss()
            self.loss1 = 0
            self.loss2 = 0
            self.loss3 = 0

        def forward(self, pred1: torch.Tensor, pred2: torch.Tensor, label: torch.Tensor, odds_rate: torch.Tensor,  mask: torch.Tensor) -> torch.Tensor:
            if len(pred1.shape) == 1:
                batch_size = 1
                pred1, pred2, label, odds_rate, mask = pred1.unsqueeze(0), pred2.unsqueeze(
                    0), label.unsqueeze(0), odds_rate.unsqueeze(0), mask.unsqueeze(0)
            else:
                batch_size = pred1.shape[0]

            # loss1
            _, indices = torch.topk(label, k=self.topn, largest=False)
            _, indices2 = torch.topk(label, k=18-self.topn)
            ans_proba = torch.cat([torch.gather(odds_rate, dim=1, index=indices), torch.gather(
                odds_rate, dim=1, index=indices2).sum(dim=1).reshape(batch_size, 1)], dim=1)

            _, indices_pred1 = torch.topk(pred1, k=self.topn)
            _, indices_pred2 = torch.topk(pred1, k=18-self.topn, largest=False)
            pred_proba = torch.cat([torch.gather(pred2, dim=1, index=indices_pred1), torch.gather(
                pred2, dim=1, index=indices_pred2).sum(dim=1).reshape(batch_size, 1)], dim=1)

            loss1 = self.KLloss1(nn.functional.log_softmax(
                pred_proba, dim=1), ans_proba)
            self.loss1 = loss1.item()

            # loss2
            # loss2 = self.KLloss2(self.log_softmax(pred2), odds_rate)
            # loss2 = self.MAE(
            #     nn.functional.softmax(pred2, dim=1), odds_rate)

            # pred2 = 0.8/torch.clip(nn.functional.softmax(pred2, dim=1), 0.8/500)
            # odds_rate = 0.8/torch.clip(odds_rate, 0.8/500)

            pred2 = torch.log(torch.clip(
                nn.functional.softmax(pred2, dim=1), 0.8/500))
            odds_rate = torch.log(torch.clip(odds_rate, 0.8/500))

            loss2 = self.MAE(pred2, odds_rate)
            self.loss2 = loss2.item()

            # loss3
            ans_rank = (19 - label)
            pair_diff = pred1.unsqueeze(2) - pred1.unsqueeze(1)
            # pair_labels = (ans_rank.unsqueeze(2) - ans_rank.unsqueeze(1)).sign()
            pair_labels = 5*(ans_rank.unsqueeze(2) - ans_rank.unsqueeze(1))
            # 数値安定性を考慮したロジスティック損失
            log_loss = nn.functional.softplus(-pair_labels * pair_diff)
            # 平均損失を返す
            loss3 = log_loss.mean()
            self.loss3 = loss3.item()

            # loss3 = self.KLloss2(
            #     nn.functional.log_softmax(pred2, dim=1), odds_rate)

            return self.l1*loss1 + self.l2*loss2 + self.l3*loss3
            # return self.l2*log_loss + self.l3*loss3

datas1 = dataset_mapping["2019first"].valid_dataset[0]
datas1["label"], datas1["odds_rate"]

# 着順予測：スコア値、確信度的なものなので、適当に正規分布からランダムサンプリング
output1 = torch.cat([torch.randn_like(datas1["odds_rate"][datas1["label"] < 19]),
                    torch.zeros_like(datas1["label"][datas1["label"] >= 19])])
# オッズの勝率分布推定：パディングを除いて正規化された確率値なので、すべて1にしてsoftmaxで正規化する。つまりすべてのオッズ値が等しい状態にする
output2 = torch.cat([torch.log_softmax(torch.ones_like(datas1["odds_rate"][datas1["label"] < 19]), dim=0),
                    torch.zeros_like(datas1["label"][datas1["label"] >= 19])-100])

# 微分可能かの確認用に勾配計算の対象とする
output1 = torch.tensor(output1, requires_grad=True)
output2 = torch.tensor(output2, requires_grad=True)
output1, output2

criterion = MultiTaskLoss(topn=17)
loss = criterion(output1, output2,
                 datas1["label"], datas1["odds_rate"], datas1["mask"])
loss.backward(retain_graph=True)  # 微分計算
loss

output1.grad, output2.grad

import math
from typing import Literal
import torch.nn.functional as F


class KeibaAIThirdModelForMultiTask(nn.Module):
    def __init__(self, cat_num_list: list[int], numerous_feature_num: int) -> None:
        super(KeibaAIThirdModelForMultiTask, 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

        # 畳み込み層
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=1)

        # 畳み込み層（共有）
        self.shared_conv = nn.Conv2d(
            in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1)

        self.task1_input0 = nn.Linear(
            8 * (18//4) * (((numerous_feature_num+sum(self.embed_num_list))//2)//2), 2048)
        self.task2_input0 = nn.Linear(
            8 * (18//4) * (((numerous_feature_num+sum(self.embed_num_list))//2)//2), 2048)

        self.task1_input1 = nn.Linear(2048, 512)
        self.task2_input1 = nn.Linear(2048, 512)

        self.task1_input2 = nn.Linear(512, 256)
        self.task2_input2 = nn.Linear(512, 256)

        self.task1_input3 = nn.Linear(256, 128)
        self.task2_input3 = nn.Linear(256, 128)

        self.task1_input4 = nn.Linear(128, 64)
        self.task2_input4 = nn.Linear(128, 64)

        # タスク1: 着順の推定
        self.fc_task1 = nn.Linear(64, 18)
        # タスク2: オッズの分布推定
        self.fc_task2 = nn.Linear(64, 18)

        self.drop_rate = 0.25
        self.dropout = nn.Dropout(p=self.drop_rate)
        self.relu = nn.ReLU()

        # self.log_softmax = nn.LogSoftmax(dim=0)
        self.mode = torch.tensor(0)
        self.max_pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.max_pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

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

        # 畳み込み層 + プーリング
        # input(batch_size, 1, 18, 特徴量数) -> (batch_size, 16, 18, 特徴量数)
        x = self.relu(self.conv1(x))
        # (batch_size, 16, 18, 特徴量数) -> (batch_size, 16, 9, 特徴量数//2)
        x = self.max_pool1(x)

        # 畳み込み層（共有）
        # (batch_size, 16, 9, 特徴量数//2) -> (batch_size, 32, 9, 特徴量数//2)
        x = self.relu(self.shared_conv(x))
        # (batch_size, 32, 9, 特徴量数//2) -> (batch_size, 32, 4, 特徴量数//4)
        x = self.max_pool2(x)

        # フラット化
        if len(x_num.shape) > 2:
            x = x.view(x.size(0), -1)
        else:
            x = x.view(-1)

        # タスク1: 着順の推定
        x_task1 = self.dropout(self.relu(self.task1_input0(x)))
        x_task1 = self.dropout(self.relu(self.task1_input1(x_task1)))
        x_task1 = self.dropout(self.relu(self.task1_input2(x_task1)))
        x_task1 = self.dropout(self.relu(self.task1_input3(x_task1)))
        x_task1 = self.dropout(self.relu(self.task1_input4(x_task1)))
        task1_output: torch.Tensor = self.fc_task1(x_task1)

        # タスク2: オッズの分布推定
        x_task2 = self.dropout(self.relu(self.task2_input0(x)))
        x_task2 = self.dropout(self.relu(self.task2_input1(x_task2)))
        x_task2 = self.dropout(self.relu(self.task2_input2(x_task2)))
        x_task2 = self.dropout(self.relu(self.task2_input3(x_task2)))
        x_task2 = self.dropout(self.relu(self.task2_input4(x_task2)))
        task2_output: torch.Tensor = self.fc_task2(x_task2)

        # task2_output = self.log_softmax(task2_output)
        if self.mode < 1:
            return task1_output, task2_output
        elif self.mode < 2:
            return task1_output
        else:
            return task2_output

    def predict(self, x_num: torch.Tensor, x_cat_list: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        output1, output2 = self.forward(x_num, x_cat_list)
        return output1[mask > 0], output2[mask > 0]

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)
datas1 = next(iter(dataset_loader))

model = KeibaAIThirdModelForMultiTask(dataset.cat_num_list, len(num_feas))
model.eval()
y = model(datas1["num_feas"][0], datas1["cat_feas"][0])
dot = make_dot(y, params=dict(model.named_parameters()))

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

y

# Early Stopping クラス
class EarlyStopping:
    def __init__(self, patience=5, min_delta=0.0, lr_threshold=1e-9, path='best_model.pth'):
        """
        patience: 改善が見られないエポック数の許容値  
        min_delta: 改善とみなす最小の変化量  
        path: 最良モデルを保存するファイルパス
        """
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float('inf')
        self.early_stop = False
        self.path = path
        self.lr_threshold = lr_threshold
        self.best_epoch = 0

    def __call__(self, val_loss, model, epoch, current_lr):
        # 学習率が閾値を下回った場合、Early Stoppingを有効化
        if current_lr <= self.lr_threshold:
            print(
                f"Learning rate ({current_lr:.2e}) is below the threshold ({self.lr_threshold:.2e}). Stopping training.")
            self.early_stop = True
            return

        if self.best_loss is None or val_loss < self.best_loss - self.min_delta:
            # 改善が見られた場合、モデルを保存
            self.best_loss = val_loss
            self.counter = 0
            torch.save(model.state_dict(), self.path)  # 最良モデルを保存
            self.best_epoch = epoch+1
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True

    def load_best_model(self, model, device):
        """保存した最良モデルをロードする"""
        model.load_state_dict(torch.load(self.path, map_location=device))

import tqdm
import numpy as np
from torch.utils.tensorboard import SummaryWriter
import datetime

for _ in range(20):
    del model
    torch.cuda.empty_cache()

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

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

    # モデルの初期化
    model = KeibaAIThirdModelForMultiTask(
        target_dataset.cat_num_list, len(num_feas)).to(device)
    # 損失関数
    param = np.array([1., 1., 1.])
    param /= param.sum()
    loss_fn = MultiTaskLoss(*param.tolist(), topn=3)
    # オプティマイザ
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=5e-4, weight_decay=0.5)

    # スケジューラー (ReduceLROnPlateau)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=2, verbose=True)

    # Early Stopping インスタンス
    best_model_path = "runs/third_model_monitor/" + \
        datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    early_stopping = EarlyStopping(
        patience=10, min_delta=0.0000, lr_threshold=1e-7, path=best_model_path+"/best_model.pth")

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

    writer = SummaryWriter(log_dir=best_model_path)

    loss_list_all = []
    # 学習ループ
    num_epochs = 100
    valid_loss, test_loss, best_loss, global_step = 0, 0, 0, 0
    torch.cuda.empty_cache()

    with tqdm.tqdm(total=len(train_dataloader)*num_epochs, desc=f"Epoch 1/{num_epochs}. loss: train=None, valid=None, test=None, Best=None (0)") as pbar:
        for epoch in range(num_epochs):
            model.train()  # 訓練モード
            loss_list = []
            for datas in train_dataloader:
                global_step += 1
                num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                    "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
                # データとラベルをGPUに転送
                num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                    device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

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

                # 損失の計算
                loss = loss_fn(outputs[0], outputs[1], labels, odds_rate, mask)

                # 逆伝播
                loss.backward()

                # 最適化
                optimizer.step()
                current_lr = optimizer.param_groups[0]['lr']
                writer.add_scalar('Loss/Loss1 (TOPn オッズ分布)',
                                  loss_fn.loss1, global_step)
                writer.add_scalar('Loss/Loss2 (オッズ推定MAE)',
                                  loss_fn.loss2, global_step)
                writer.add_scalar('Loss/Loss3 (着順Pair Wise)',
                                  loss_fn.loss3, global_step)
                writer.add_scalar('Loss/Total', loss.item(), global_step)

                # ロスを記録
                loss_list += [loss.item()]
                pbar.update()
                pbar.set_description(
                    desc=f"Epoch {epoch+1}/{num_epochs}. loss: train={np.mean(loss_list):.4f}, valid={valid_loss:.4f}, test={test_loss:.4f}, Best={best_loss:.4f} ({early_stopping.best_epoch})")
                del num_data, cat_data_list, labels, odds_rate, mask, datas

            loss_list_all += [loss_list]

            model.eval()  # 推論モード
            with torch.no_grad():
                # 検証データのロス確認
                for datas in valid_dataloader:
                    num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                        "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
                    # データとラベルをGPUに転送
                    num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                        device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

                    # 順伝播
                    outputs = model(num_data, cat_data_list)
                    # 損失の計算
                    loss = loss_fn(outputs[0], outputs[1],
                                   labels, odds_rate, mask)
                    del num_data, cat_data_list, labels, odds_rate, mask, datas

                # テストデータのロス確認
                for datas in test_dataloader:
                    num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                        "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
                    # データとラベルをGPUに転送
                    num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                        device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

                    # 順伝播
                    outputs_test = model(num_data, cat_data_list)
                    # 損失の計算
                    loss_test = loss_fn(
                        outputs_test[0], outputs_test[1], labels, odds_rate, mask)
                    del num_data, cat_data_list, labels, odds_rate, mask, datas

                valid_loss = loss.item()
                test_loss = loss_test.item()
                writer.add_scalar('Epoch Loss/Total-Train',
                                  np.mean(loss_list), global_step)
                writer.add_scalar('Epoch Loss/Total-Valid',
                                  valid_loss, global_step)
                writer.add_scalar('Epoch Loss/Total-Test',
                                  test_loss, global_step)

            # 学習率スケジューラーの更新
            scheduler.step(valid_loss, epoch)
            writer.add_scalar('Learning Rate/current', current_lr, epoch)

            # Early Stopping のチェック
            early_stopping(valid_loss, model, epoch, current_lr)
            best_loss = early_stopping.best_loss
            pbar.set_description(
                desc=f"Epoch {epoch+1}/{num_epochs}. loss: train={np.mean(loss_list):.4f}, valid={valid_loss:.4f}, test={test_loss:.4f}, Best={best_loss:.4f} ({early_stopping.best_epoch})")

            if early_stopping.early_stop:
                print(f"Early stopping at epoch {epoch+1}")
                break

            torch.cuda.empty_cache()
    writer.close()

    # 最良モデルに戻す
    early_stopping.load_best_model(model, device)
    print(
        f"Restored the best model parameters! Best Epoch: {early_stopping.best_epoch}, Best Loss: {early_stopping.best_loss}")
    torch.cuda.empty_cache()

import tqdm
import numpy as np

target_dataset = dataset_mapping["2023second"]
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

dataMap = {}
root_model_dir = pathlib.Path("./runs/third_model_monitor")
for path in tqdm.tqdm(list(root_model_dir.glob("**/*.pth"))):
    # モデルの初期化
    model = KeibaAIThirdModelForMultiTask(
        target_dataset.cat_num_list, len(num_feas)).to(device)
    model.load_state_dict(torch.load(path, map_location=device))
    with torch.no_grad():
        model.to(device)
        model.eval()
        train_dataloader = DataLoader(
            target_dataset.train_dataset, batch_size=1024, shuffle=False)
        valid_dataloader = DataLoader(target_dataset.valid_dataset, batch_size=len(
            target_dataset.valid_dataset), shuffle=False)
        test_dataloader = DataLoader(
            target_dataset.test_dataset, batch_size=len(target_dataset.test_dataset), shuffle=False)
        all_outputs0 = []
        all_outputs1 = []
        for datas in train_dataloader:
            num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
            # データとラベルをGPUに転送
            num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

            # 順伝播
            outputs = model.predict(num_data, cat_data_list, mask)
            all_outputs0 += outputs[0].to("cpu").detach().numpy().tolist()
            all_outputs1 += outputs[1].to("cpu").detach().numpy().tolist()
            del num_data, cat_data_list, labels, odds_rate, mask

        dftrain = target_dataset.train.copy()
        dftrain["pred_odds_rate"] = np.exp(np.array(all_outputs1))
        dftrain["pred_odds_rate"] /= dftrain["raceId"].map(
            dftrain[["raceId", "pred_odds_rate"]].groupby("raceId")["pred_odds_rate"].sum())
        dftrain["pred_odds"] = 0.8/dftrain["pred_odds_rate"]
        dftrain["pred_odds_rank"] = dftrain.groupby(
            "raceId")["pred_odds_rate"].rank(ascending=False).astype(int)
        dftrain["pred_proba"] = np.array(all_outputs0)
        dftrain["pred_rank"] = dftrain.groupby(
            "raceId")["pred_proba"].rank(ascending=False).astype(int)
        del all_outputs0, all_outputs1

        for datas in valid_dataloader:
            num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
            # データとラベルをGPUに転送
            num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

            # 順伝播
            outputs = model.predict(num_data, cat_data_list, mask)
            del num_data, cat_data_list, labels, odds_rate, mask

        dfvalid = target_dataset.valid.copy()
        dfvalid["pred_odds_rate"] = np.exp(
            outputs[1].to("cpu").detach().numpy())
        dfvalid["pred_odds_rate"] /= dfvalid["raceId"].map(
            dfvalid[["raceId", "pred_odds_rate"]].groupby("raceId")["pred_odds_rate"].sum())
        dfvalid["pred_odds"] = 0.8/dfvalid["pred_odds_rate"]
        dfvalid["pred_odds_rank"] = dfvalid.groupby(
            "raceId")["pred_odds_rate"].rank(ascending=False).astype(int)
        dfvalid["pred_proba"] = outputs[0].to("cpu").detach().numpy()
        dfvalid["pred_rank"] = dfvalid.groupby(
            "raceId")["pred_proba"].rank(ascending=False).astype(int)
        del outputs

        for datas in test_dataloader:
            num_data, cat_data_list, labels, odds_rate, mask = datas["num_feas"], datas[
                "cat_feas"], datas["label"], datas["odds_rate"], datas["mask"]
            # データとラベルをGPUに転送
            num_data, cat_data_list, labels, odds_rate, mask = num_data.to(
                device), cat_data_list.to(device), labels.to(device), odds_rate.to(device), mask.to(device)

            # 順伝播
            outputs_test = model.predict(num_data, cat_data_list, mask)
            del num_data, cat_data_list, labels, odds_rate, mask

        dftest = target_dataset.test.copy()
        dftest["pred_odds_rate"] = np.exp(
            outputs_test[1].to("cpu").detach().numpy())
        dftest["pred_odds_rate"] /= dftest["raceId"].map(
            dftest[["raceId", "pred_odds_rate"]].groupby("raceId")["pred_odds_rate"].sum())
        dftest["pred_odds"] = 0.8/dftest["pred_odds_rate"]
        dftest["pred_odds_rank"] = dftest.groupby(
            "raceId")["pred_odds_rate"].rank(ascending=False).astype(int)
        dftest["pred_proba"] = outputs_test[0].to("cpu").detach().numpy()
        dftest["pred_rank"] = dftest.groupby(
            "raceId")["pred_proba"].rank(ascending=False).astype(int)
        del outputs_test
    del model
    torch.cuda.empty_cache()
    dataMap[path.parent.name] = {
        "train": dftrain,
        "valid": dfvalid,
        "test": dftest
    }