from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import unittest
from . import adskalman
import numpy
import pkg_resources
import warnings

def assert_3d_vs_kpm_close(A,B,debug=False):
    """compare my matrices (where T dimension is first) vs. KPM's (where it's last)"""
    for i in range(A.shape[0]):
        if debug:
            print()
            print('i=%d'%i)
            print('A[i],',A[i].shape,A[i])
            print("B[:,:,i]",B[:,:,i].shape,B[:,:,i])
            diff = A[i]-B[:,:,i]
            print('diff',diff)
        try:
            assert numpy.allclose( A[i], B[:,:,i])
            if debug:
                print('(same)')
        except Exception as err:
            if debug:
                print('(different)')
            raise

class TestStats(unittest.TestCase):
    def test_likelihood1(self):
        pi = numpy.pi
        exp = numpy.exp

        x=numpy.array([0,0])
        m=numpy.array([[0,0]])
        C=numpy.array([[1.0,0],[0,1]])
        lik = adskalman.gaussian_prob(x,m,C)
        lik_should = 1/(2*pi) * 1 * exp( 0 )
        assert numpy.allclose(lik,lik_should)

    def test_rand_mvn1(self):
        mu = numpy.array([1,2,3,4])
        sigma = numpy.eye(4)
        sigma[:2,:2] = [[2.0, 0.1],[0.1,0.2]]
        N = 1000
        Y = adskalman.rand_mvn(mu,sigma,N)
        assert Y.shape==(N,4)
        mu2 = numpy.mean(Y,axis=0)

        eps = .2
        assert numpy.sqrt(numpy.sum((mu-mu2)**2)) < eps # expect occasional failure here

        sigma2 = adskalman.covar(Y)
        eps = .3
        assert numpy.sqrt(numpy.sum((sigma-sigma2)**2)) < eps # expect occasional failure here

    def test_rand_mvn2(self):
        mu = numpy.array([1])
        sigma = 2*numpy.eye(1)
        N = 10000
        Y = adskalman.rand_mvn(mu,sigma,N)
        Y2 = numpy.random.standard_normal((N,1))*sigma[0]+mu[0]
        assert Y.shape == (N,1)
        assert Y2.shape == (N,1)

        Y = Y[:,0] # drop an axis
        Y2 = Y2[:,0] # drop an axis

        meanY = numpy.mean(Y)
        #print 'meanY',meanY
        meanY2 = numpy.mean(Y2)
        #print 'meanY2',meanY2

        # according to z-test on wikipedia
        SE = sigma[0]/numpy.sqrt(N)
        SE2 = sigma[0]/numpy.sqrt(N)

        z = (meanY-mu[0])/SE
        z2 = (meanY2-mu[0])/SE2

        if abs(z) > 3:
            raise ValueError('Sample mean was three sigma away from true mean. This error is expected rarely.')
        if abs(z2) > 3:
            raise ValueError('Sample mean was three sigma away from true mean. This error is expected rarely.')
        #print 'z',z
        #print 'z2',z2

##         assert abs(meanY) < 1e-5

##         est_std_diff = abs(numpy.std(Y)-1)
##         #print 'err',est_std_diff
##         assert est_std_diff < tol

class TestAutogeneratedData(unittest.TestCase):
    # parametric.ParametricTestCase
    #: Prefix for tests with independent state.  These methods will be run with
    #: a separate setUp/tearDown call for each test in the group.
    #_indepParTestPrefix = 'test_kalman_murphy_smooth'

    def setUp(self):
        # process model
        A = numpy.array([[1, 0],
                         [0, .5]],
                        dtype=numpy.float64)

        # observation model
        C = numpy.array([[1, 0],
                         [0, 1]],
                        dtype=numpy.float64)

        # process covariance
        Q = numpy.eye(2)

        # measurement covariance
        R = numpy.eye(2)

        Q = 0.001*Q
        R = 0.001*R

        N = 1000
        numpy.random.seed(4)
        x0 = numpy.random.randn( 2 )
        P0 = numpy.random.randn( 2,2 )

        process_noise = adskalman.rand_mvn( 0, Q, N )
        observation_noise = adskalman.rand_mvn( 0, R, N )

        X = []
        Y = []
        x = x0
        for i in range(N):
            X.append( x )
            Y.append( numpy.dot(C,x) + observation_noise[i] )

            x = numpy.dot(A,x) + process_noise[i] # for next time step

        X = numpy.array(X)
        Y = numpy.array(Y)

        Abad = numpy.eye(2)
        Abad[0,0] = .1
        Abad[1,1] = 4

        self.X = X
        self.Y = Y
        self.A = A
        self.Abad = Abad
        self.C = C
        self.Q = Q
        self.R = R
        self.x0 = x0
        self.P0 = P0

    def test_kalman_murphy_smooth(self):
        test = self._test_kalman_murphy_smooth
        for missing_type in ['nan','large']:
            #yield (test,missing_type)
            test(missing_type)

    def _test_kalman_murphy_smooth(self,missing_type):
        Ymissing = numpy.array(self.Y, copy=True)

        if missing_type=='nan':
            valid_idx = None
        elif missing_type=='large':
            valid_idx = list(range(len(Ymissing)))
        else:
            raise ValueError('unknown missing_type "%s"'%missing_type)

        bad_idx = list(range(5,len(Ymissing),10))
        for i in bad_idx:
            if missing_type=='nan':
                Ymissing[i] = numpy.nan*Ymissing[i]
            elif missing_type=='large':
                valid_idx.remove(i)
                Ymissing[i] = 9e999*Ymissing[i]

        assert len(bad_idx)>5 # make sure some data are actually missing!
        xsmooth, Vsmooth = adskalman.kalman_smoother(Ymissing,self.A,self.C,self.Q,self.R,self.x0,self.P0,
                                                     valid_data_idx=valid_idx)
        dist = numpy.sqrt(numpy.sum((xsmooth-self.X)**2,axis=1))
        mean_dist = numpy.mean(dist)
        #print 'mean_dist',missing_type,mean_dist
        assert mean_dist < 0.1 # Not sure of exact value, but shouldn't be too large

    def _test_kalman_murphy_EM(self): # temporarily disabled
        def returnC(orig):
            return self.C
        def returnQ(orig):
            return self.Q
        def returnR(orig):
            return self.R
        constr_fun_dict = {
            'C':returnC,
            #'Q':returnQ,
            #'R':returnR,
            }
        if 1:
            print('-'*80,'Murphy')
        Alearn, Clearn, Qlearn, Rlearn, initx2, initV2, LL = adskalman.learn_kalman(
            self.Y, self.Abad, self.C, self.Q, self.R, self.x0,
            self.P0, max_iter=500, verbose=2,
            constr_fun_dict=constr_fun_dict,
            thresh=1e-5,
            )

        xsmooth, Vsmooth, VVsmooth, loglik = adskalman.kalman_smoother(
            self.Y,self.A,self.C,self.Q,self.R,self.x0,self.P0,
            full_output=True)
        if 1:
            print('best LL',loglik)

            print()
            print('*'*80,'Murphy')
            print('A')
            print(self.A)
            print('Alearn')
            print(Alearn)
            print('Clearn')
            print(Clearn)
            print('Qlearn')
            print(Qlearn)
            print('Rlearn')
            print(Rlearn)
            print('*'*80,'Murphy')

class TestLotsOfAutogeneratedData(unittest.TestCase):
    def setUp(self):
        dt = 0.5
        # process model
        A = numpy.array([[1, 0, dt, 0],
                         [0, 1, 0, dt],
                         [0, 0, .5,  0],
                         [0, 0, 0,  .5]],
                        dtype=numpy.float64)
        # observation model
        C = numpy.array([[1, 0, 0, 0],
                         [0, 1, 0, 0]],
                        dtype=numpy.float64)

        # process covariance
        Q = numpy.eye(4)
        Q[:2,:2] = [[ 2, .1],[.1,2]]

        # measurement covariance
        R = numpy.array([[ 1, 0.2], [0.2, 1]],
                        dtype=numpy.float64)

        Q = 0.001*Q
        R = 0.001*R

        N = 1000
        numpy.random.seed(4)
        x0 = numpy.random.randn( 4 )
        P0 = numpy.random.randn( 4,4 )

        process_noise = adskalman.rand_mvn( 0, Q, N )
        observation_noise = adskalman.rand_mvn( 0, R, N )

        X = []
        Y = []
        x = x0
        for i in range(N):
            X.append( x )
            Y.append( numpy.dot(C,x) + observation_noise[i] )

            x = numpy.dot(A,x) + process_noise[i] # for next time step

        X = numpy.array(X)
        Y = numpy.array(Y)

        Abad = numpy.eye(4)
        #Abad[0,0] = 0
        #Abad[1,1] = 100
        #Abad[0,2] = .0
        #Abad[1,3] = .0

        self.Y = Y
        self.A = A
        self.Abad = Abad
        self.C = C
        self.Q = Q
        self.R = R
        self.x0 = x0
        self.P0 = P0

    def _test_kalman_ads1(self):
        print('-'*80,'DRO')
        xlearn, Vlearn, Alearn, Clearn, Qlearn, Rlearn = adskalman.DROsmooth(self.Y,
                                                                             self.Abad,
                                                                             self.C,
                                                                             self.Q,
                                                                             self.R,
                                                                             self.x0,
                                                                             self.P0,
                                                                             mode='EM')
        print()
        print('*'*80,'DRO')
        print('Alearn',Alearn)
        print('Clearn',Clearn)
        print('Qlearn',Qlearn)
        print('Rlearn',Rlearn)
        print('*'*80,'DRO')


class TestKalman(unittest.TestCase):
    def test_kalman1(self,time_steps=100,Qsigma=0.1,Rsigma=0.5):
        dt = 0.1
        # process model
        A = numpy.array([[1, 0, dt, 0],
                         [0, 1, 0, dt],
                         [0, 0, 1,  0],
                         [0, 0, 0,  1]],
                        dtype=numpy.float64)
        # observation model
        C = numpy.array([[1, 0, 0, 0],
                         [0, 1, 0, 0]],
                        dtype=numpy.float64)
        # process covariance
        Q = Qsigma*numpy.eye(4)
        # measurement covariance
        R = Rsigma*numpy.eye(2)

        x = numpy.array([0,0,0,0],dtype=numpy.float)
        x += Qsigma*numpy.random.standard_normal(size=x.shape)

        kf = adskalman.KalmanFilter(A,C,Q,R,x,Q)
        y = numpy.dot(C,x)
        y += Rsigma*numpy.random.standard_normal(y.shape)

        xs = []
        xhats = []
        for i in range(time_steps):
            if i==0: isinitial=True
            else: isinitial = False

            xhat,P = kf.step(y=y, isinitial=isinitial)

            # calculate new state
            x = numpy.dot(A,x) + Qsigma*numpy.random.standard_normal(x.shape)
            # and new observation
            y = numpy.dot(C,x) + Rsigma*numpy.random.standard_normal(y.shape)

            xs.append(x)
            xhats.append( xhat )
        xs = numpy.array(xs)
        xhats = numpy.array(xhats)
        # XXX some comparison between xs and xhats

def get_test_suite():
    ts=unittest.TestSuite([unittest.makeSuite(TestKalman),
                           unittest.makeSuite(TestAutogeneratedData),
                           #unittest.makeSuite(TestLotsOfAutogeneratedData),
                           unittest.makeSuite(TestStats),
                           ])
    return ts

if __name__=='__main__':
    numpy.set_printoptions(suppress=True)
    if 0:
        suite = get_test_suite()
        suite.debug()
    elif 0:
        suite = unittest.makeSuite(TestLotsOfAutogeneratedData)
        suite.debug()
    elif 1:
        suite = unittest.makeSuite(TestAutogeneratedData)
        suite.debug()
