import collections
import logging
import time
from typing import Optional, List, Union, Dict

from torch.utils.data import Dataset
from transformers import Trainer, EvalPrediction, is_torch_tpu_available
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from transformers.trainer_utils import speed_metrics

if is_torch_tpu_available():
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met

logger = logging.getLogger(__name__)


def compute_metrics(pred: EvalPrediction):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }


class SingleTLMTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        """
        If negative sampling enabled. The train dataload should load only positve examples, otherwise it load balanced pos/neg examples.
        :param args:
        :param kwargs:
        """
        super().__init__(*args, **kwargs)



class SiameseTLMTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
