"""Public development checks, not an independent learner assessment."""
import copy
import importlib.util
import json
import math
from pathlib import Path
import tempfile
import unittest

HERE = Path(__file__).resolve().parent
spec = importlib.util.spec_from_file_location('baseline', HERE/'perception-baseline.py')
b = importlib.util.module_from_spec(spec)
spec.loader.exec_module(b)
DATA = HERE/'subset.json'
if not DATA.exists():
    DATA = HERE.parent/'data/pets/subset.json'


class BaselineChecks(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.data = b.load_data(DATA)

    def test_known_metrics_and_ignored_pixels(self):
        value = b.score([1, 1, 2, 2, 3], [True, False, True, False, True])
        self.assertEqual(value['petIoU'], 1/3)
        self.assertEqual(value['accuracy'], .5)
        self.assertEqual(value['excludedPixels'], 1)
        self.assertIsNone(b.score([3], [True])['accuracy'])
        self.assertIsNone(b.score([2], [False])['petIoU'])

    def test_all_six_gradients_and_simultaneous_update(self):
        rows = [([1, -.3, .8, -.4, .2, 1], 1), ([-.7, .4, .1, .9, -.2, 1], 0)]
        weights = [.1, -.2, .3, .4, -.5, .2]
        loss, gradient = b.objective(weights, rows)
        expected = sum(-y*math.log(b.sigmoid(sum(w*x for w, x in zip(weights, xs))))
                       -(1-y)*math.log(1-b.sigmoid(sum(w*x for w, x in zip(weights, xs)))) for xs, y in rows)/2
        self.assertAlmostEqual(loss, expected, places=12)
        for j in range(6):
            plus, minus = weights[:], weights[:]
            plus[j] += 1e-5
            minus[j] -= 1e-5
            numerical = (b.objective(plus, rows)[0]-b.objective(minus, rows)[0])/2e-5
            self.assertAlmostEqual(gradient[j], numerical, places=8)
        records = self.data['records']
        gradient = b.objective([0.]*6, b.prepare(records, 'color-position', 'training'))[1]
        trained, _ = b.train(records, 'color-position', .3, 1)
        self.assertEqual(trained['weights'], [-.3*g for g in gradient])

    def test_development_and_ignored_pixels_never_train(self):
        records = copy.deepcopy(self.data['records'])
        for r in records:
            for i, label in enumerate(r['trimap']):
                if r['split'] == 'development' or label == 3:
                    r['rgb'][3*i:3*i+3] = [255, 0, 255]
            if r['split'] == 'development':
                r['trimap'] = [1]*len(r['trimap'])
        for mode in ('color', 'color-position'):
            expected = b.train(self.data['records'], mode, .3, 2)
            self.assertEqual(expected, b.train(records, mode, .3, 2))

    def test_source_identity_duplicates_and_partition(self):
        records = copy.deepcopy(self.data['records'])
        records[-1]['sourceImage']['sha256'] = records[0]['sourceImage']['sha256']
        with self.assertRaises(ValueError):
            b.audit(records)
        records = copy.deepcopy(self.data['records'])
        records[0]['split'] = 'development'
        with self.assertRaises(ValueError):
            b.audit(records)
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp)/'subset.json'
            path.write_bytes(DATA.read_bytes().replace(b'boxer', b'BOXER', 1))
            with self.assertRaises(ValueError):
                b.load_data(path)

    def test_replay_ignores_forged_scores_and_preserves_weights(self):
        model, _ = b.train(self.data['records'], 'color', .3, 2)
        artifact = {'version': b.VERSION, 'sourceHash': b.SOURCE_HASH, 'model': model,
                    'threshold': .5, 'images': 'forged', 'trainingLossRecomputed': -100}
        copy_before = copy.deepcopy(artifact)
        result = b.replay(self.data, artifact)
        self.assertEqual(artifact, copy_before)
        self.assertEqual(result['model'], model)
        self.assertEqual(len(result['images']), 8)
        self.assertGreater(result['trainingLossRecomputed'], 0)
        self.assertEqual(result['freshSourceEvaluation'], 'pending')
        self.assertEqual(b.replay(self.data, json.loads(json.dumps(result))), result)

    def test_bounds(self):
        for count in (0, 1001, True, 1.5):
            with self.assertRaises(ValueError):
                b.train(self.data['records'], updates=count)
        for rate in (0, -1, float('nan'), float('inf'), True):
            with self.assertRaises(ValueError):
                b.train(self.data['records'], rate=rate)
        for weights in ([0]*5, [float('nan')]*6, [51]*6, [True]*6):
            with self.assertRaises(ValueError):
                b.validate_model({'mode': 'color-position', 'weights': weights, 'updates': 1})

    def test_learning_and_determinism(self):
        first, trace = b.train(self.data['records'], updates=10)
        self.assertLess(trace['finalTrainingLoss'], trace['initialTrainingLoss'])
        self.assertEqual(first, b.train(self.data['records'], updates=10)[0])
        self.assertEqual(first['weights'][3:5], [0, 0])


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