import unittest
import numpy as np
import torch
from utils.dataloader import Dataloader, tokenize_and_preserve_labels
from utils.metric_tracking import MetricsTracking 
from transformers import BertTokenizer
import re

class DataloaderTest(unittest.TestCase):

    def test_tokenize_sentence(self):
        label_to_ids = {
           'B-MEDCOND': 0,
           'I-MEDCOND': 1,
           'O': 2
        }
        ids_to_label = {
               0:'B-MEDCOND',
               1:'I-MEDCOND',
               2:'O'
        }
        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        max_tokens = 128

        sentence = "Patient presents with glaucoma, characterized by a definitive diagnosis of pigmentary glaucoma. Intraocular pressure measures at 15 mmHg, while the visual field remains normal. Visual acuity is recorded as 20/50. The patient has not undergone prior cataract surgery, but has had LASIK surgery. Additionally, comorbid ocular diseases include macular degeneration."

        tokens = "O O O B-MEDCOND O O O O O O O B-MEDCOND I-MEDCOND O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O O B-MEDCOND I-MEDCOND I-MEDCOND O B-MEDCOND I-MEDCOND O"

        sentence = re.findall(r"\w+|\w+(?='s)|'s|['\".,!?;]", sentence.strip(), re.UNICODE)
        tokens = tokens.split(" ")

        t_sen, t_labl = tokenize_and_preserve_labels(sentence, tokens, tokenizer, label_to_ids, ids_to_label, max_tokens)

        self.assertEqual(len(t_sen), len(t_labl))
        self.assertEqual(t_labl.count("B-MEDCOND"), tokens.count("B-MEDCOND"))

    def test_load_dataset(self):
        label_to_ids = {
           'B-MEDCOND': 0,
           'I-MEDCOND': 1,
           'O': 2
        }
        ids_to_label = {
               0:'B-MEDCOND',
               1:'I-MEDCOND',
               2:'O'
        }

        dataloader = Dataloader(label_to_ids, ids_to_label)

        dataset = dataloader.load_dataset(full = True)

        self.assertEqual(len(dataset), 255)

        sample = dataset.__getitem__(0)

        self.assertEqual(len(sample), 4) #input_ids, attention_mask, token_type_ids, entity



class MetricsTrackingTest(unittest.TestCase):

    def test_avg_metrics(self):
        predictions =  np.array([-100, 0, 0, 0, 1, 1, 1, 2, 2, 2])
        ground_truth = np.array([-100, 0, 1, 0, 1, 2, 1, 2, 0, 2]) #arbitrary, should return 67% for each metric
        
        predictions = torch.from_numpy(predictions)
        ground_truth = torch.from_numpy(ground_truth)

        tracker = MetricsTracking()
        tracker.update(predictions, ground_truth, 0.1)

        metrics = tracker.return_avg_metrics(1) #tracker only updated once

        self.assertEqual(metrics['acc'], 0.667)
        self.assertEqual(metrics['f1'], 0.667)
        self.assertEqual(metrics['precision'], 0.667)
        self.assertEqual(metrics['recall'], 0.667)
        self.assertEqual(metrics['loss'], 0.1)

if __name__ == '__main__':
    unittest.main()