import os

import pandas as pd

import numpy as np

from dataset import LABEL_COLS

if __name__ == '__main__':

    target_cols = LABEL_COLS

    test_csv = "./csv/patient2_kfold/test.csv"

    pred_paths = [

        '/logs/prediction/densenet169-mw-512-resume-0/test_0_ckp_tta.npy',

        '/logs/prediction/densenet169-mw-512-resume-1/test_1_ckp_tta.npy',

        '/logs/prediction/densenet169-mw-512-resume-2/test_2_ckp_tta.npy',

        "/logs/prediction/densenet169-mw-512-resume-3/test_3_ckp_tta.npy",

        "/logs/prediction/densenet169-mw-512-resume-4/test_4_ckp_tta.npy"

        # '/logs/predictions/se_resnext50_32x4d-mw-512-recheck-0/test_0.npy',

    ]

    test_preds = 0

    for pred in pred_paths:

        test_preds += np.load(pred)

    test_preds = test_preds / len(pred_paths)

    test_df = pd.read_csv(test_csv)

    test_df["sop_instance_uid"] = "ID_" + test_df["sop_instance_uid"]

    test_ids = test_df['sop_instance_uid'].values

    ids = []

    labels = []

    for i, id in enumerate(test_ids):

        pred = test_preds[i]

        for j, target in enumerate(target_cols):

            id_target = id + "_" + target

            ids.append(id_target)

            labels.append(pred[j])

        # if not with_any:

        #     id_target = id + "_" + "any"

        #     ids.append(id_target)

        #     labels.append(pred.max())

    submission_df = pd.DataFrame({

        'ID': ids,

        'Label': labels

    })

    os.makedirs(f"/logs/prediction/ensemble/", exist_ok=True)

    submission_df.to_csv(f"/logs/prediction/ensemble/densenet169-mww-512-5folds_tta.csv", index=False)