# -*- coding: utf-8 -*-
# @Date    : 2020/12/16 19:52
# @Author  : stellahong (stellahong@fuzhi.ai)
# @Desc    :
import logging
import os

import torch
import numpy as np
from tqdm import tqdm
import csv
from typing import List
from sentence_transformers import InputExample
from sklearn.metrics import roc_auc_score
from torch.utils.data.dataloader import DataLoader


class CEAucEvaluator():
    """
    This evaluator can be used with the CrossEncoder class.

    It is designed for CrossEncoders with 2 or more outputs. It measure the
    accuracy of the predict class vs. the gold labels.
    """
    
    def __init__(self, sentence_pairs: List[List[str]], labels: List[int], name: str = '', num_class: int = 2):
        self.sentence_pairs = sentence_pairs
        self.labels = labels
        self.name = name
        
        self.csv_file = "CEAucEvaluator" + ("_" + name if name else '') + "_results.csv"
        self.csv_headers = ["epoch", "steps", "Accuracy"]
    
    @classmethod
    def from_input_examples(cls, examples: List[InputExample], **kwargs):
        sentence_pairs = []
        labels = []
        num_class = kwargs["num_class"]
        for example in examples:
            sentence_pairs.append(example.texts)
            labels.append(example.label)
        labels = np.asarray(labels).reshape(-1, num_class)
        return cls(sentence_pairs, labels, **kwargs)
    
    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        if epoch != -1:
            if steps == -1:
                out_txt = " after epoch {}:".format(epoch)
            else:
                out_txt = " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"
        
        logging.info("CEAucEvaluator: Evaluating the model on " + self.name + " dataset" + out_txt)
        pred_scores = model.predict(self.sentence_pairs, convert_to_numpy=True, apply_softmax=True,
                                    show_progress_bar=True)
        
        
        auc = roc_auc_score(self.labels, pred_scores, average='micro')
        
        logging.info("Auc: {:.2f}".format(auc * 1))
        
        if output_path is not None:
            csv_path = os.path.join(output_path, self.csv_file)
            
            output_file_exists = os.path.isfile(csv_path)
            with open(csv_path, mode="a" if output_file_exists else 'w', encoding="utf-8") as f:
                writer = csv.writer(f)
                if not output_file_exists:
                    writer.writerow(self.csv_headers)

                writer.writerow([epoch, steps, auc])
        
        return auc
