
import unittest
from hypothesis import given, example, settings, assume, note
from hypothesis.strategies import integers, one_of, just, booleans
from utils import one_hot_encode
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

"""
Test hand-made calculations with hypothesis bounds
"""
class TestCrossEntropy(unittest.TestCase):
    @settings(max_examples=1000)
    @given(classes=integers(min_value=2, max_value=10),
           seed=integers(min_value=0, max_value=1000))
    def test_ce_lipschitzianity(self, classes: int, seed: int):
        torch.manual_seed(seed)
        input1 = torch.randn((1, classes))
        input2 = torch.randn((1, classes))
        cls = torch.randint(classes, (1,))
        note(f"Input1: {input1}")
        note(f"Input2: {input2}")
        note(f"Class: {cls}")
        ce1 = F.cross_entropy(input1, cls)
        ce2 = F.cross_entropy(input2, cls)
        note(f"ce1: {ce1}")
        note(f"ce2: {ce2}")
        note(f"input dist: {torch.norm(input1 - input2).item()}")
        note(f"Result: {(ce1 - ce2).abs().item()}")
        self.assertLessEqual((ce1 - ce2).abs().item(), 2.0 * torch.norm(input1 - input2).item())


    @settings(max_examples=1000)
    @given(classes=integers(min_value=2, max_value=10),
           seed=integers(min_value=0, max_value=1000))
    def test_ce_smoothness(self, classes: int, seed: int):
        torch.manual_seed(seed)
        input1 = torch.randn((1, classes))
        input2 = torch.randn((1, classes))
        cls = torch.randint(classes, (1,))
        note(f"Input1: {input1}")
        note(f"Input2: {input2}")
        note(f"Class: {cls}")
        clsce = lambda input: F.cross_entropy(input, cls)
        # We want to get the derivative of the cross_entropy function
        ce1d = torch.autograd.functional.jvp(clsce, input1, v=torch.ones((1, classes)))[0]
        ce2d = torch.autograd.functional.jvp(clsce, input2, v=torch.ones((1, classes)))[0]
        note(f"ce1d: {ce1d}")
        note(f"ce2d: {ce2d}")
        note(f"input dist: {torch.norm(input1 - input2).item()}")
        note(f"Result: {(ce1d - ce2d).abs().item()}")
        self.assertLessEqual((ce1d - ce2d).abs().item(), 2.0 * torch.norm(input1 - input2).item())
        

if __name__ == "__main__":
    unittest.run()
