import GPflow
import tensorflow as tf
import numpy as np
import unittest


class DumbModel(GPflow.model.Model):
    def __init__(self):
        GPflow.model.Model.__init__(self)
        self.a = GPflow.param.Param(3.)

    def build_likelihood(self):
        return -tf.square(self.a)


class NoArgsModel(DumbModel):
    @GPflow.model.AutoFlow()
    def function(self):
        return self.a


class TestNoArgs(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.m = NoArgsModel()
        self.m._compile()

    def test_return(self):
        self.assertTrue(np.allclose(self.m.function(), 3.))

    def test_kill(self):
        # make sure autoflow dicts are removed when _needs_recompile is set.
        keys = [k for k in self.m.__dict__.keys() if k[-11:] == '_AF_storage']
        self.assertTrue(len(keys) == 0, msg="no AF storage should be present to start.")

        self.m.function()

        keys = [k for k in self.m.__dict__.keys() if k[-11:] == '_AF_storage']
        self.assertTrue(len(keys) == 1, msg="AF storage should be present after function call.")

        self.m._needs_recompile = True

        keys = [k for k in self.m.__dict__.keys() if k[-11:] == '_AF_storage']
        self.assertTrue(len(keys) == 0, msg="no AF storage should be present after recompile switch set.")

        self.m.function()
        self.m._needs_recompile = False

        keys = [k for k in self.m.__dict__.keys() if k[-11:] == '_AF_storage']
        self.assertTrue(len(keys) == 1, msg="AF storage should be persistent when setting _needs_recompile to False.")


class AddModel(DumbModel):
    @GPflow.model.AutoFlow((tf.float64,), (tf.float64,))
    def add(self, x, y):
        return tf.add(x, y)


class TestShareArgs(unittest.TestCase):
    """
    This is designed to replicate bug #85, where having two models caused
    autoflow functions to fail because the tf_args were shared over the
    instances.
    """
    def setUp(self):
        tf.reset_default_graph()
        self.m1 = AddModel()
        self.m1._compile()
        self.m2 = AddModel()
        self.m2._compile()
        rng = np.random.RandomState(0)
        self.x = rng.randn(10, 20)
        self.y = rng.randn(10, 20)

    def test_share_args(self):
        self.m1.add(self.x, self.y)
        self.m2.add(self.x, self.y)
        self.m1.add(self.x, self.y)


class TestAdd(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.m = AddModel()
        self.m._compile()
        rng = np.random.RandomState(0)
        self.x = rng.randn(10, 20)
        self.y = rng.randn(10, 20)

    def test_add(self):
        self.assertTrue(np.allclose(self.x + self.y, self.m.add(self.x, self.y)))


class IncrementModel(DumbModel):
    def __init__(self):
        DumbModel.__init__(self)
        self.a = GPflow.param.DataHolder(np.array([3.]))

    @GPflow.model.AutoFlow((tf.float64,))
    def inc(self, x):
        return x + self.a


class TestDataHolder(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.m = IncrementModel()
        rng = np.random.RandomState(0)
        self.x = rng.randn(10, 20)

    def test_add(self):
        self.assertTrue(np.allclose(self.x + 3, self.m.inc(self.x)))


class TestGPmodel(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        rng = np.random.RandomState(0)
        X, Y = rng.randn(2, 10, 1)
        self.m = GPflow.svgp.SVGP(X, Y, kern=GPflow.kernels.Matern32(1),
                                  likelihood=GPflow.likelihoods.StudentT(),
                                  Z=X[::2].copy())
        self.Xtest = np.random.randn(100, 1)
        self.Ytest = np.random.randn(100, 1)

    def test_predict_f(self):
        mu, var = self.m.predict_f(self.Xtest)

    def test_predict_y(self):
        mu, var = self.m.predict_y(self.Xtest)

    def test_predict_density(self):
        self.m.predict_density(self.Xtest, self.Ytest)

    def test_multiple_AFs(self):
        self.m.compute_log_likelihood()
        self.m.compute_log_prior()
        self.m.compute_log_likelihood()


class TestResetGraph(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        k = GPflow.kernels.Matern32(1)
        X, Y = np.random.randn(2, 10, 1)
        self.Xnew = np.random.randn(5, 1)
        self.m = GPflow.gpr.GPR(X, Y, kern=k)

    def test(self):
        mu, var = self.m.predict_f(self.Xnew)
        tf.reset_default_graph()
        mu1, var1 = self.m.predict_f(self.Xnew)


class TestFixAndPredict(unittest.TestCase):
    """
    Bug #54 says that if a model parameter is fixed  between calls to predict
    (an autoflow fn) then the second call fails. This test ensures replicates
    that and ensures that the bugfix remains in furure.
    """

    def setUp(self):
        rng = np.random.RandomState(0)
        X, Y = rng.randn(2, 10, 1)
        self.m = GPflow.svgp.SVGP(X, Y, kern=GPflow.kernels.Matern32(1),
                                  likelihood=GPflow.likelihoods.StudentT(),
                                  Z=X[::2].copy())
        self.Xtest = np.random.randn(100, 1)
        self.Ytest = np.random.randn(100, 1)

    def test(self):
        self.m._compile()
        self.m.kern.variance.fixed = True
        _, _ = self.m.predict_f(self.m.X.value)


class TestSVGP(unittest.TestCase):
    """
    This replicates Alex's code from bug #99
    """
    def test(self):
        rng = np.random.RandomState(1)
        X = rng.randn(10, 1)
        Y = rng.randn(10, 1)
        Z = rng.randn(3, 1)
        model = GPflow.svgp.SVGP(X=X, Y=Y, kern=GPflow.kernels.RBF(1), likelihood=GPflow.likelihoods.Gaussian(), Z=Z)
        model.compute_log_likelihood()


class NestedModels(GPflow.param.Parameterized):

    def __init__(self):
        super(NestedModels, self).__init__()
        self.m1 = NoArgsModel()
        self.m2 = NoArgsModel()
        self.m1._compile()
        self.m2._compile()

    @GPflow.model.AutoFlow()
    def function(self):
        return self.m1.a + self.m2.a


class TestNestedModels(unittest.TestCase):

    def setUp(self):
        self.m = NestedModels()

    def assertAutoFlowCaches(self, p, number, msg=""):
        keys = [k for k in p.__dict__.keys() if k[-11:] == '_AF_storage']
        self.assertTrue(len(keys) == number, msg=msg)

    def test_kill(self):
        self.assertAutoFlowCaches(self.m, 0, "no AF storage should be present to start.")
        self.assertAutoFlowCaches(self.m.m1, 0, "no AF storage should be present to start.")
        self.assertAutoFlowCaches(self.m.m2, 0, "no AF storage should be present to start.")

        self.m.function()
        self.assertAutoFlowCaches(self.m, 1, "AF storage should be present after function call.")
        self.assertAutoFlowCaches(self.m.m1, 0)
        self.assertAutoFlowCaches(self.m.m2, 0)

        self.m.m1._needs_recompile = False
        self.assertAutoFlowCaches(self.m, 1, "AF storage should be persistent after setting needs_recompile to False.")
        self.m.m1._needs_recompile = True
        self.assertAutoFlowCaches(self.m, 0, "Setting needs_recompile to True should erase parent AF storage.")

    def test_kill_sibling_persistence(self):
        self.m.m1.function()
        self.m.m2.function()
        self.m.function()
        self.assertAutoFlowCaches(self.m, 1, "AF storage should be present after function call.")
        self.assertAutoFlowCaches(self.m.m1, 1, "AF storage should be present after function call.")
        self.assertAutoFlowCaches(self.m.m2, 1, "AF storage should be present after function call.")

        self.m.m1._needs_recompile = True
        self.assertAutoFlowCaches(self.m, 0, "Setting needs_recompile to True should erase parent AF storage.")
        self.assertAutoFlowCaches(self.m.m1, 0, "Setting needs_recompile to True should erase AF storage.")
        self.assertAutoFlowCaches(self.m.m2, 1, "AF storage should be persistent, as model is in different subtree")

        self.m.m1._needs_recompile = False
        self.m.m1.function()
        self.assertAutoFlowCaches(self.m, 0, "no AF storage should be present without running function.")
        self.m.m2._needs_recompile = True
        self.assertAutoFlowCaches(self.m.m2, 0, "Setting needs_recompile to True should erase AF storage.")
        self.assertAutoFlowCaches(self.m.m1, 1, "AF storage should be persistent, as model is in different subtree")

    def test_kill_replacements(self):
        self.m.function()
        self.assertAutoFlowCaches(self.m, 1, "AF storage should be present after function call.")
        self.m.m1 = NoArgsModel()
        self.assertAutoFlowCaches(self.m, 0, "AF storage should be erased after replacing a Param.")


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