import numpy as np
import unittest

from pydlm.modeler.trends import trend
from pydlm.modeler.seasonality import seasonality
from pydlm.modeler.builder import builder
from pydlm.base.kalmanFilter import kalmanFilter


class testKalmanFilter(unittest.TestCase):


    def setUp(self):
        self.kf1 = kalmanFilter(discount=[1])
        self.kf0 = kalmanFilter(discount=[1e-10])
        self.kf11 = kalmanFilter(discount=[1, 1])
        self.trend0 = trend(degree=0, discount=1, w=1.0)
        self.trend0_90 = trend(degree=0, discount=0.9, w=1.0)
        self.trend0_98 = trend(degree=0, discount=0.98, w=1.0, name='a')
        self.trend1 = trend(degree=1, discount=1, w=1.0)


    def testForwardFilter(self):
        dlm = builder()
        dlm.add(self.trend0)
        dlm.initialize()
        self.kf1.predict(dlm.model)
        self.assertAlmostEqual(dlm.model.prediction.obs, 0)

        # the prior on the mean is zero, but observe 1, with
        # discount = 1, one should expect the filterd mean to be 0.5
        self.kf1.forwardFilter(dlm.model, 1)
        self.assertAlmostEqual(dlm.model.obs, 0.5)
        self.assertAlmostEqual(dlm.model.prediction.obs, 0)
        self.assertAlmostEqual(dlm.model.sysVar, 0.375)

        self.kf1.predict(dlm.model)
        self.assertAlmostEqual(dlm.model.obs, 0.5)
        self.assertAlmostEqual(dlm.model.prediction.obs, 0.5)

        dlm.initialize()
        self.kf0.predict(dlm.model)
        self.assertAlmostEqual(dlm.model.prediction.obs, 0)

        # the prior on the mean is zero, but observe 1, with discount = 0
        # one should expect the filtered mean close to 1
        self.kf0.forwardFilter(dlm.model, 1)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1)
        self.assertAlmostEqual(dlm.model.prediction.obs[0, 0], 0)
        self.assertAlmostEqual(dlm.model.sysVar[0, 0], 0.5)

        self.kf0.predict(dlm.model)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1)
        self.assertAlmostEqual(dlm.model.prediction.obs[0, 0], 1)


    def testForwardFilterMultiDim(self):
        dlm = builder()
        dlm.add(seasonality(period=2, discount=1, w=1.0))
        dlm.initialize()

        self.kf11.forwardFilter(dlm.model, 1)
        self.assertAlmostEqual(dlm.model.state[0][0, 0], 0.33333333333)
        self.assertAlmostEqual(dlm.model.state[1][0, 0], -0.33333333333)

        self.kf11.forwardFilter(dlm.model, -1)
        self.assertAlmostEqual(dlm.model.state[0][0, 0], -0.5)
        self.assertAlmostEqual(dlm.model.state[1][0, 0], 0.5)


    def testBackwardSmoother(self):
        dlm = builder()
        dlm.add(self.trend0)
        dlm.initialize()

        # with mean being 0 and observe 1 and 0 consectively, one shall
        # expect the smoothed mean at 1 will be 1/3, for discount = 1
        self.kf1.forwardFilter(dlm.model, 1)
        self.kf1.forwardFilter(dlm.model, 0)
        self.kf1.backwardSmoother(dlm.model, \
                                  np.matrix([[0.5]]), \
                                  np.matrix([[0.375]]))
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1.0/3)
        self.assertAlmostEqual(dlm.model.sysVar[0, 0], 0.18518519)


    # second order trend with discount = 1. The smoothed result should be
    # equal to a direct fit on the three data points, 0, 1, -1. Thus, the
    # smoothed observation should be 0.0
    def testBackwardSmootherMultiDim(self):
        dlm = builder()
        dlm.add(self.trend1)
        dlm.initialize()

        self.kf11.forwardFilter(dlm.model, 1)
        state1 = dlm.model.state
        cov1 = dlm.model.sysVar

        self.kf11.forwardFilter(dlm.model, -1)
        self.kf11.backwardSmoother(dlm.model, \
                                   rawState = state1, \
                                   rawSysVar = cov1)

        self.assertAlmostEqual(dlm.model.obs[0, 0], 0.0)


    def testMissingData(self):
        dlm = builder()
        dlm.add(self.trend0)
        dlm.initialize()

        self.kf0.forwardFilter(dlm.model, 1)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1.0)
        self.assertAlmostEqual(dlm.model.obsVar[0, 0], 1.0)

        self.kf0.forwardFilter(dlm.model, None)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1.0)
        self.assertAlmostEqual(dlm.model.obsVar[0, 0]/1e10, 0.5)

        self.kf0.forwardFilter(dlm.model, None)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 1.0)
        self.assertAlmostEqual(dlm.model.obsVar[0, 0]/1e10, 0.5)

        self.kf0.forwardFilter(dlm.model, 0)
        self.assertAlmostEqual(dlm.model.obs[0, 0], 0.0)


    def testMissingEvaluation(self):
        dlm = builder()
        dlm.add(self.trend0)
        dlm.initialize()

        dlm.model.evaluation = np.matrix([[None]])
        self.kf1.forwardFilter(dlm.model, 1.0, dealWithMissingEvaluation = True)
        self.assertAlmostEqual(dlm.model.obs, 0.0)
        self.assertAlmostEqual(dlm.model.transition, 1.0)


    def testEvolveMode(self):
        dlm = builder()
        dlm.add(self.trend0_90)
        dlm.add(self.trend0_98)
        dlm.initialize()

        kf2 = kalmanFilter(discount=[0.9, 0.98],
                           updateInnovation='component',
                           index=dlm.componentIndex)
        kf2.forwardFilter(dlm.model, 1.0)
        self.assertAlmostEqual(dlm.model.innovation[0, 1], 0.0)
        self.assertAlmostEqual(dlm.model.innovation[1, 0], 0.0)


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

[(154, 'unittest.main', 'unittest.main', False, 'import unittest\n'), (14, 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), (15, 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), (16, 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), (17, 'pydlm.modeler.trends.trend', 'trend', False, 'from pydlm.modeler.trends import trend\n'), (18, 'pydlm.modeler.trends.trend', 'trend', False, 'from pydlm.modeler.trends import trend\n'), (19, 'pydlm.modeler.trends.trend', 'trend', False, 'from pydlm.modeler.trends import trend\n'), (20, 'pydlm.modeler.trends.trend', 'trend', False, 'from pydlm.modeler.trends import trend\n'), (24, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (58, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (72, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (91, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (108, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (129, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (133, 'numpy.matrix', 'np.matrix', True, 'import numpy as np\n'), (140, 'pydlm.modeler.builder.builder', 'builder', False, 'from pydlm.modeler.builder import builder\n'), (145, 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), (59, 'pydlm.modeler.seasonality.seasonality', 'seasonality', False, 'from pydlm.modeler.seasonality import seasonality\n'), (81, 'numpy.matrix', 'np.matrix', True, 'import numpy as np\n'), (82, 'numpy.matrix', 'np.matrix', True, 'import numpy as np\n')]