import os, sys
sys.path.append(os.path.abspath('../../main/python'))

import datetime as dt

import numpy as np
import numpy.testing as npt
import matplotlib.pyplot as plt

from thalesians.tsa.distrs import NormalDistr as N
import thalesians.tsa.filtering as filtering
import thalesians.tsa.filtering.kalman as kalman
import thalesians.tsa.filtering.particle as particle
import thalesians.tsa.numpyutils as npu
import thalesians.tsa.processes as proc

import importlib
importlib.reload(particle)
importlib.reload(proc)

"""
Explanation: Particle filtering
Particle filtering is not working yet - WORK IN PROGRESS!
End of explanation
"""


process = proc.WienerProcess.create_from_cov(mean=3., cov=0.0001)

"""
Explanation: A single-process, univariate example
First we need a process model. In this case it will be a single stochastic process,
End of explanation
"""


t0 = dt.datetime(2017, 5, 12, 16, 18, 25, 204000)
pf = particle.ParticleFilter(t0, state_distr=N(mean=100., cov=0.0000000000001), process=process)

"""
Explanation: This we pass to a newly created particle filter, along with the initial time and initial state. The latter takes the form of a normal distribution. We have chosen to use Python datetimes as our data type for time, but we could have chosen ints or something else.
End of explanation
"""


observable = pf.create_observable(kalman.LinearGaussianObsModel.create(1.), process)

"""
Explanation: Next we create an observable, which incorporates a particular observation model. In this case, the observation model is particularly simple, since we are observing the entire state of the particle filter. Our observation model is a 1x1 identity:
End of explanation
"""


t1 = t0 + dt.timedelta(hours=1)

"""
Explanation: Let's roll forward the time by one hour:
End of explanation
"""


prior_predicted_obs1 = observable.predict(t1)
prior_predicted_obs1

"""
Explanation: What is our predicted observation at this time? Since we haven't observed any actual information, this is our prior observation estimate:
End of explanation
"""


np.mean(pf._prior_particles), 100. + 3./24.

prior_predicted_obs1

prior_predicted_obs1 = observable.predict(t1)
npt.assert_almost_equal(prior_predicted_obs1.distr.mean, 100. + 3./24.)
npt.assert_almost_equal(prior_predicted_obs1.distr.cov, 250. + 25./24.)
npt.assert_almost_equal(prior_predicted_obs1.cross_cov, prior_predicted_obs1.distr.cov)

"""
Explanation: We confirm that this is consistent with how our (linear-Gaussian) process model scales over time:
End of explanation
"""


observable.observe(time=t1, obs=N(mean=100.35, cov=100.0))

"""
Explanation: Let us now actually observe our observation. Say, the observation is 100.35 and the observation noise covariance is 100.0:
End of explanation
"""


posterior_predicted_obs1 = observable.predict(t1); posterior_predicted_obs1

"""
Explanation: Having seen an actual observation, let us obtain the posterior observation estimate:
End of explanation
"""


t2 = t1 + dt.timedelta(hours=2)
        
prior_predicted_obs2 = observable.predict(t2)
npt.assert_almost_equal(prior_predicted_obs2.distr.mean, 100.28590504 + 2.*3./24.)
npt.assert_almost_equal(prior_predicted_obs2.distr.cov, 71.513353115 + 2.*25./24.)
npt.assert_almost_equal(prior_predicted_obs2.cross_cov, prior_predicted_obs2.distr.cov)
        
observable.observe(time=t2, obs=N(mean=100.35, cov=100.0))

posterior_predicted_obs2 = observable.predict(t2)
npt.assert_almost_equal(posterior_predicted_obs2.distr.mean, 100.45709020)
npt.assert_almost_equal(posterior_predicted_obs2.distr.cov, 42.395213845)
npt.assert_almost_equal(posterior_predicted_obs2.cross_cov, posterior_predicted_obs2.distr.cov)


"""
Explanation: We can now fast-forward the time, by two hours, say, and repeat the process:
End of explanation
"""


process1 = proc.WienerProcess.create_from_cov(mean=3., cov=25.)
process2 = proc.WienerProcess.create_from_cov(mean=[1., 4.], cov=[[36.0, -9.0], [-9.0, 25.0]])

"""
Explanation: A multi-process, multivariate example
The real power of our particle filter interface is demonstrated for process models consisting of several (independent) stochastic processes:
End of explanation
"""


t0 = dt.datetime(2017, 5, 12, 16, 18, 25, 204000)
kf = kalman.KalmanFilter(
    t0,
    state_distr=N(
        mean=[100.0, 120.0, 130.0],
        cov=[[250.0, 0.0, 0.0],
             [0.0, 360.0, 0.0],
             [0.0, 0.0, 250.0]]),
    process=(process1, process2))

"""
Explanation: Such models are common in finance, where, for example, the dynamics of a yield curve may be represented by a (multivariate) stochastic process, whereas the idiosyncratic spread for each bond may be an independent stochastic process.
Let us pass process1 and process2 as a (compound) process model to our particle filter, along with the initial time and state:
End of explanation
"""


state_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(1.0, np.eye(2)),
    process1, process2)

"""
Explanation: We shall now create several observables, each corresponding to a distinct observation model. The first one will observe the entire state:
End of explanation
"""


coord0_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(1.),
    process1)

"""
Explanation: The second observable will observe the first coordinate of the first process:
End of explanation
"""


coord1_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(npu.row(1., 0.)),
    process2)

"""
Explanation: The third, the first coordinate of the second process:
End of explanation
"""


coord2_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(npu.row(0., 1.)),
    process2)

"""
Explanation: The fourth, the second coordinate of the second process:
End of explanation
"""


sum_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(npu.row(1., 1., 1.)),
    process1, process2)

"""
Explanation: The fifth will observe the sum of the entire state (across the two processes):
End of explanation
"""


lin_comb_observable = kf.create_observable(
    kalman.KalmanFilterObsModel.create(npu.row(2., 0., -3.)),
    process1, process2)

"""
Explanation: And the sixth a certain linear combination thereof:
End of explanation
"""


t1 = t0 + dt.timedelta(hours=1)

"""
Explanation: Fast-forward the time by one hour:
End of explanation
"""


predicted_obs1_prior = state_observable.predict(t1)
predicted_obs1_prior

"""
Explanation: Let's predict the state at this time...
End of explanation
"""


npt.assert_almost_equal(predicted_obs1_prior.distr.mean,
                        npu.col(100.0 + 3.0/24.0, 120.0 + 1.0/24.0, 130.0 + 4.0/24.0))
npt.assert_almost_equal(predicted_obs1_prior.distr.cov,
                        [[250.0 + 25.0/24.0, 0.0, 0.0],
                         [0.0, 360.0 + 36.0/24.0, -9.0/24.0],
                         [0.0, -9.0/24.0, 250 + 25.0/24.0]])
npt.assert_almost_equal(predicted_obs1_prior.cross_cov, predicted_obs1_prior.distr.cov)

"""
Explanation: And check that it is consistent with the scaling of the (multivariate) Wiener process with time:
End of explanation
"""


state_observable.observe(time=t1, obs=N(mean=[100.35, 121.0, 135.0],
                                        cov=[[100.0, 0.0, 0.0],
                                             [0.0, 400.0, 0.0],
                                             [0.0, 0.0, 100.0]]));

"""
Explanation: Suppose that a new observation arrives, and we observe each of the three coordinates individually:
End of explanation
"""


state_observable.predict(t1)

"""
Explanation: Let's look at our (posterior) predicted state:
End of explanation
"""


coord0_observable.predict(t1)

coord1_observable.predict(t1)

coord2_observable.predict(t1)

"""
Explanation: Let's also look at the predictions for the individual coordinates:
End of explanation
"""


sum_observable.predict(t1)

"""
Explanation: The predicted sum:
End of explanation
"""


lin_comb_observable.predict(t1)

"""
Explanation: And the predicted linear combination:
End of explanation
"""


t2 = t1 + dt.timedelta(minutes=30)

"""
Explanation: Let's now go 30 minutes into the future:
End of explanation
"""


coord1_observable.observe(time=t2, obs=N(mean=125.25, cov=4.))

"""
Explanation: And observe only the first coordinate of the second process, with a pretty high confidence:
End of explanation
"""


state_observable.predict(t2)

"""
Explanation: How does our predicted state change?
End of explanation
"""


t3 = t2 + dt.timedelta(minutes=30)

"""
Explanation: Thirty minutes later...
End of explanation
"""


sum_observable.observe(time=t3, obs=N(mean=365.00, cov=9.))

"""
Explanation: We observe the sum of the three coordinates, rather than the individual coordinates:
End of explanation
"""


state_observable.predict(t3)

"""
Explanation: How has our prediction of the state changed?
End of explanation
"""


sum_observable.predict(t3)

"""
Explanation: And what is its predicted sum?
End of explanation
"""


