import unittest
import random
import torch
from sklearn.metrics import r2_score
from lightwood.encoders.numeric import NumericEncoder
from lightwood.encoders.text import DistilBertEncoder
from lightwood import COLUMN_DATA_TYPES


class TestDistilBERT(unittest.TestCase):
    def test_encode_and_decode(self):
        pass
        # random.seed(2)
        # priming_data = []
        # primting_target = []
        # test_data = []
        # test_target = []
        # for i in range(0, 300):
        #     if random.randint(1, 5) == 3:
        #         test_data.append(str(i) + ''.join(['n'] * i))
        #         # test_data.append(str(i))
        #         test_target.append(i)
        #     # else:
        #     priming_data.append(str(i) + ''.join(['n'] * i))
        #     # priming_data.append(str(i))
        #     primting_target.append(i)

        # output_1_encoder = NumericEncoder(is_target=True)
        # output_1_encoder.prepare_encoder(primting_target)

        # encoded_data_1 = output_1_encoder.encode(primting_target)
        # encoded_data_1 = encoded_data_1.tolist()

        # enc = DistilBertEncoder()

        # enc.prepare_encoder(priming_data,
        #                     training_data={'targets': [
        #                         {'output_type': COLUMN_DATA_TYPES.NUMERIC,'encoded_output': encoded_data_1},
        #                         {'output_type': COLUMN_DATA_TYPES.NUMERIC, 'encoded_output': encoded_data_1}
        #                     ]})

        # encoded_predicted_target = enc.encode(test_data).tolist()

        # predicted_targets_1 = output_1_encoder.decode(torch.tensor([x[:3] for x in encoded_predicted_target]))
        # predicted_targets_2 = output_1_encoder.decode(torch.tensor([x[3:] for x in encoded_predicted_target]))

        # for predicted_targets in [predicted_targets_1, predicted_targets_2]:
        #     real = list(test_target)
        #     pred = list(predicted_targets)

        #     # handle nan
        #     for i in range(len(pred)):
        #         try:
        #             float(pred[i])
        #         except:
        #             pred[i] = 0

        #     print(real[0:25], '\n', pred[0:25])
        #     encoder_accuracy = r2_score(real, pred)

        #     print(f'Categorial encoder accuracy for: {encoder_accuracy} on testing dataset')
        #     # assert(encoder_accuracy > 0.5)
