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()

[((5442, 5457), 'unittest.main', 'unittest.main', ([], {}), '()\n', (5455, 5457), False, 'import unittest\n'), ((303, 329), 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', ([], {'discount': '[1]'}), '(discount=[1])\n', (315, 329), False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), ((349, 379), 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', ([], {'discount': '[1e-10]'}), '(discount=[1e-10])\n', (361, 379), False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), ((400, 429), 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', ([], {'discount': '[1, 1]'}), '(discount=[1, 1])\n', (412, 429), False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), ((452, 486), 'pydlm.modeler.trends.trend', 'trend', ([], {'degree': '(0)', 'discount': '(1)', 'w': '(1.0)'}), '(degree=0, discount=1, w=1.0)\n', (457, 486), False, 'from pydlm.modeler.trends import trend\n'), ((512, 548), 'pydlm.modeler.trends.trend', 'trend', ([], {'degree': '(0)', 'discount': '(0.9)', 'w': '(1.0)'}), '(degree=0, discount=0.9, w=1.0)\n', (517, 548), False, 'from pydlm.modeler.trends import trend\n'), ((574, 621), 'pydlm.modeler.trends.trend', 'trend', ([], {'degree': '(0)', 'discount': '(0.98)', 'w': '(1.0)', 'name': '"""a"""'}), "(degree=0, discount=0.98, w=1.0, name='a')\n", (579, 621), False, 'from pydlm.modeler.trends import trend\n'), ((644, 678), 'pydlm.modeler.trends.trend', 'trend', ([], {'degree': '(1)', 'discount': '(1)', 'w': '(1.0)'}), '(degree=1, discount=1, w=1.0)\n', (649, 678), False, 'from pydlm.modeler.trends import trend\n'), ((728, 737), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (735, 737), False, 'from pydlm.modeler.builder import builder\n'), ((2076, 2085), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (2083, 2085), False, 'from pydlm.modeler.builder import builder\n'), ((2586, 2595), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (2593, 2595), False, 'from pydlm.modeler.builder import builder\n'), ((3417, 3426), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (3424, 3426), False, 'from pydlm.modeler.builder import builder\n'), ((3903, 3912), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (3910, 3912), False, 'from pydlm.modeler.builder import builder\n'), ((4628, 4637), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (4635, 4637), False, 'from pydlm.modeler.builder import builder\n'), ((4724, 4743), 'numpy.matrix', 'np.matrix', (['[[None]]'], {}), '([[None]])\n', (4733, 4743), True, 'import numpy as np\n'), ((4980, 4989), 'pydlm.modeler.builder.builder', 'builder', ([], {}), '()\n', (4987, 4989), False, 'from pydlm.modeler.builder import builder\n'), ((5094, 5189), 'pydlm.base.kalmanFilter.kalmanFilter', 'kalmanFilter', ([], {'discount': '[0.9, 0.98]', 'updateInnovation': '"""component"""', 'index': 'dlm.componentIndex'}), "(discount=[0.9, 0.98], updateInnovation='component', index=dlm.\n    componentIndex)\n", (5106, 5189), False, 'from pydlm.base.kalmanFilter import kalmanFilter\n'), ((2102, 2142), 'pydlm.modeler.seasonality.seasonality', 'seasonality', ([], {'period': '(2)', 'discount': '(1)', 'w': '(1.0)'}), '(period=2, discount=1, w=1.0)\n', (2113, 2142), False, 'from pydlm.modeler.seasonality import seasonality\n'), ((2964, 2982), 'numpy.matrix', 'np.matrix', (['[[0.5]]'], {}), '([[0.5]])\n', (2973, 2982), True, 'import numpy as np\n'), ((3020, 3040), 'numpy.matrix', 'np.matrix', (['[[0.375]]'], {}), '([[0.375]])\n', (3029, 3040), True, 'import numpy as np\n')]