import unittest
from config import *

from fastai.vision import *
from fastai.gen_doc.doctest import this_tests
from fastai.callbacks import *
from math import isclose
from fastai.train import ClassificationInterpretation

path = untar_data(URLs.MNIST_TINY)
data = ImageDataBunch.from_folder(path, ds_tfms=(rand_pad(2, 28), []), num_workers=2)
data = data.normalize()

learn = Learner(data, simple_cnn((3,16,16,16,2), bn=True), metrics=[accuracy, error_rate])
learn.fit_one_cycle(3)

class IntegrationTest(unittest.TestCase):

    def test_accuracy(self):
        self.assertGreater(accuracy(*learn.get_preds()), 0.9)
        
    def test_error_rate(self):
        self.assertLess(error_rate(*learn.get_preds()), 0.1)
        
    def test_preds(self):
        pass_tst = False
        for i in range(3):
            img, label = learn.data.valid_ds[i]
            pred_class,pred_idx,outputs = learn.predict(img)
            if outputs[int(label)] > outputs[1-int(label)]:
                pass_tst = True
                break
        self.assertTrue(pass_tst)

class OneCycleTest(unittest.TestCase):
    
    def test_lrs(self):
        lrs = learn.recorder.lrs
        self.assertLess(lrs[0], 0.001)
        self.assertLess(lrs[-1], 0.0001)
        self.assertEqual(np.max(lrs), 3e-3)
        
    def test_moms(self):
        moms = learn.recorder.moms
        self.assertEqual(moms[0],0.95)
        self.assertLess(abs(moms[-1]-0.95), 0.01)
        self.assertEqual(np.min(moms),0.85)