from primitiv import optimizers as O
from primitiv import Optimizer, Parameter, Device, Graph, Shape
from primitiv import initializers as I
from primitiv import devices as D
from primitiv import functions as F
from primitiv import tensor_functions as tF

import unittest
import tempfile

import numpy as np


class TestAdam(Optimizer):
    def __init__(self, alpha, beta1, beta2, eps):
        super().__init__()
        self.alpha_ = np.float32(alpha)
        self.beta1_ = np.float32(beta1)
        self.beta2_ = np.float32(beta2)
        self.eps_ = np.float32(eps)

    def configure_parameter(self, param):
        for name in ("testadam-m1", "testadam-m2"):
            if name not in param.stats:
                param.add_stats(name, param.shape())
                param.stats[name].reset(0)

    def update_parameter(self, scale, param):
        epoch = self.get_epoch() + 1
        g = param.gradient
        param.stats["testadam-m1"] = self.beta1_ * param.stats["testadam-m1"] + (1 - self.beta1_) * g
        param.stats["testadam-m2"] = self.beta2_ * param.stats["testadam-m2"] + (1 - self.beta2_) * g * g
        mm1 = param.stats["testadam-m1"] / (1 - self.beta1_ ** epoch)
        mm2 = param.stats["testadam-m2"] / (1 - self.beta2_ ** epoch)
        param.value -= (scale * self.alpha_) * mm1 / (tF.sqrt(mm2) + self.eps_)

    def get_configs(self):
        uint_configs = {}
        float_configs = {
            "TestAdam.alpha": self.alpha_,
            "TestAdam.beta1": self.beta1_,
            "TestAdam.beta2": self.beta2_,
            "TestAdam.eps": self.eps_,
        }
        return uint_configs, float_configs

    def set_configs(self, uint_configs, float_configs):
        self.alpha_ = float_configs["TestAdam.alpha"]
        self.beta1_ = float_configs["TestAdam.beta1"]
        self.beta2_ = float_configs["TestAdam.beta2"]
        self.eps_ = float_configs["TestAdam.eps"]


class TestException(Exception):
    pass


class ExceptionOptimizer(Optimizer):

    def configure_parameter(self, param):
        raise TestException("configure_parameter")

    def update_parameter(self, scale, param):
        raise TestException("update_parameter")

    def get_configs(self):
        raise TestException("get_configs")

    def set_configs(self, uint_configs, float_configs):
        raise TestException("set_configs")


class IncompleteOptimizer(Optimizer):
    pass


def train_func(optimizer):
    dev = D.Naive(12345)
    Device.set_default(dev)
    g = Graph()
    Graph.set_default(g)

    pw1 = Parameter([8, 2], I.XavierUniform())
    pb1 = Parameter([8], I.Constant(0))
    pw2 = Parameter([1, 8], I.XavierUniform())
    pb2 = Parameter([1], I.Constant(0))

    optimizer.add(pw1, pb1, pw2, pb2)

    input_data = [1, 1, 1, -1, -1, 1, -1, -1]
    output_data = [1, -1, -1, 1]

    for i in range(10):
        g.clear()
        x = F.raw_input(Shape([2], 4), input_data)
        w1 = F.parameter(pw1)
        b1 = F.parameter(pb1)
        w2 = F.parameter(pw2)
        b2 = F.parameter(pb2)
        h = F.tanh(w1 @ x + b1)
        y = w2 @ h + b2

        t = F.raw_input(Shape([], 4), output_data)
        diff = t - y
        loss = F.batch.mean(diff * diff)

        optimizer.reset_gradients()
        loss.backward()
        optimizer.update()

    return [pw1.value.to_list(),
            pb1.value.to_list(),
            pw2.value.to_list(),
            pb2.value.to_list()]


class PythonOptimizerTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass

    def setUp(self):
        self.t = TestAdam(alpha = 0.001, beta1 = 0.9, beta2 = 0.999, eps = 1e-8)

    def tearDown(self):
        pass

    def test_pyoptimizer_get_set_config(self):
        uint_configs, float_configs = Optimizer.get_configs(self.t)
        self.assertAlmostEqual(uint_configs['Optimizer.epoch'], 0)
        self.assertAlmostEqual(float_configs['TestAdam.alpha'], 0.001)
        self.assertAlmostEqual(float_configs['TestAdam.beta1'], 0.9)
        self.assertAlmostEqual(float_configs['TestAdam.beta2'], 0.999)
        self.assertAlmostEqual(float_configs['TestAdam.eps'], 1e-8, places=10)
        float_configs['TestAdam.beta1'] = 200
        Optimizer.set_configs(self.t, uint_configs, float_configs)
        self.assertEqual(self.t.beta1_, 200)

    def test_pyoptimizer_parameter(self):
        dev = D.Naive()
        Device.set_default(dev)
        pw1 = Parameter([8, 2], I.XavierUniform())
        self.t.add(pw1)
        self.assertIn("testadam-m1", pw1.stats)
        self.assertIn("testadam-m2", pw1.stats)

    def test_pyoptimizer_compare_with_cpp(self):
        c_optimizer = O.Adam(alpha = 0.001, beta1 = 0.9, beta2 = 0.999, eps = 1e-8)
        py_params = train_func(self.t)
        c_params = train_func(c_optimizer)
        py_uint_configs, py_float_configs = Optimizer.get_configs(self.t)
        c_uint_configs, c_float_configs = c_optimizer.get_configs()
        self.assertEqual(py_uint_configs["Optimizer.epoch"], c_uint_configs["Optimizer.epoch"])
        self.assertEqual(py_float_configs["TestAdam.alpha"], c_float_configs["Adam.alpha"])
        self.assertEqual(py_float_configs["TestAdam.beta1"], c_float_configs["Adam.beta1"])
        self.assertEqual(py_float_configs["TestAdam.beta2"], c_float_configs["Adam.beta2"])
        self.assertEqual(py_float_configs["TestAdam.eps"], c_float_configs["Adam.eps"])
        self.assertEqual(py_float_configs["Optimizer.clip_threshold"], c_float_configs["Optimizer.clip_threshold"])
        self.assertEqual(py_float_configs["Optimizer.l2_strength"], c_float_configs["Optimizer.l2_strength"])
        self.assertEqual(py_float_configs["Optimizer.lr_scale"], c_float_configs["Optimizer.lr_scale"])
        self.assertTrue(np.isclose(py_params[0], c_params[0]).all())
        self.assertTrue(np.isclose(py_params[1], c_params[1]).all())
        self.assertTrue(np.isclose(py_params[2], c_params[2]).all())
        self.assertTrue(np.isclose(py_params[3], c_params[3]).all())

    def test_pyoptimizer_loadsave(self):
        t_loaded = TestAdam(alpha = 0, beta1 = 0, beta2 = 0, eps = 0)
        self.assertEqual(t_loaded.alpha_, 0)
        self.assertEqual(t_loaded.beta1_, 0)
        self.assertEqual(t_loaded.beta2_, 0)
        self.assertEqual(t_loaded.eps_, 0)
        with tempfile.NamedTemporaryFile() as fp:
            self.t.save(fp.name)
            t_loaded.load(fp.name)
        self.assertAlmostEqual(t_loaded.alpha_, 0.001)
        self.assertAlmostEqual(t_loaded.beta1_, 0.9)
        self.assertAlmostEqual(t_loaded.beta2_, 0.999)
        self.assertAlmostEqual(t_loaded.eps_, 1e-8, places=10)

    def test_pyoptimizer_propagate_exception(self):
        dev = D.Naive()
        Device.set_default(dev)
        optimizer = ExceptionOptimizer()
        p = Parameter()
        with self.assertRaises(TestException) as ctx:
            optimizer.add(p)
        self.assertEqual(str(ctx.exception), "configure_parameter")
        with self.assertRaises(TestException) as ctx:
            optimizer.update()
        self.assertEqual(str(ctx.exception), "update_parameter")
        with self.assertRaises(TestException) as ctx:
            Optimizer.get_configs(optimizer)
        self.assertEqual(str(ctx.exception), "get_configs")
        with self.assertRaises(TestException) as ctx:
            Optimizer.set_configs(optimizer, {'Optimizer.epoch': 1},
                                         {'Optimizer.clip_threshold': 0.0,
                                          'Optimizer.lr_scale': 1.0,
                                          'Optimizer.l2_strength': 0.0})
        self.assertEqual(str(ctx.exception), "set_configs")


    def test_pyoptimizer_not_implemented(self):
        dev = D.Naive()
        Device.set_default(dev)
        optimizer = IncompleteOptimizer()
        p = Parameter()
        with self.assertRaises(NotImplementedError):
            optimizer.add(p)
        with self.assertRaises(NotImplementedError):
            optimizer.update()
        with self.assertRaises(NotImplementedError):
            Optimizer.get_configs(optimizer)
        with self.assertRaises(NotImplementedError):
            Optimizer.set_configs(optimizer, {'Optimizer.epoch': 1},
                                         {'Optimizer.clip_threshold': 0.0,
                                          'Optimizer.lr_scale': 1.0,
                                          'Optimizer.l2_strength': 0.0})
