import os

import torch

from dataclasses import dataclass
from enum import Enum

from src.envs import CACHE_PATH


@dataclass
class Task:
    benchmark: str
    metric: str
    col_name: str
    num_fewshot: int


class Tasks(Enum):
    # task_key in the json file, metric_key in the json file, name to display in the leaderboard
    # task0 = Task("anli_r1", "acc", "ANLI")
    # task1 = Task("logiqa", "acc_norm", "LogiQA")
    task0 = Task("nq_open", "em", "NQ Open", 64)  # 64, as in the ATLAS paper
    task1 = Task("triviaqa", "em", "TriviaQA", 64)  # 64, as in the ATLAS paper
    task2 = Task("truthfulqa:mc", "mc2", "TruthfulQA", 0)  # TruthfulQA is intended as a zero-shot benchmark [5, 47]. https://owainevans.github.io/pdfs/truthfulQA_lin_evans.pdf


# NUM_FEWSHOT = 64  # Change with your few shot

EVAL_REQUESTS_PATH_BACKEND = os.path.join(CACHE_PATH, "eval-queue-bk")
EVAL_RESULTS_PATH_BACKEND = os.path.join(CACHE_PATH, "eval-results-bk")

DEVICE = "cuda" if torch.cuda.is_available() else 'cpu'

LIMIT = None  # Testing; needs to be None