# Path: pyparticleest/models/hierarchial.py
# class HierarchicalRSBase(HierarchicalBase, FFBSiRS):
#     def __init__(self, **kwargs):
#         super(HierarchicalRSBase, self).__init__(**kwargs)
# 
#     def logp_xnext_max(self, particles, u, t):
#         """
#         Calculate maximum value of the logp_xnext function, used for
#         rejection sampling
# 
# 
#         Args:
#          - particles  (array-like): Model specific representation
#            of all particles, with first dimension = N (number of particles)
#          - u (array-like): input signal
#          - t (float): time stamp
# 
#          Returns:
#           (float) max value over all particle and possible future states
#         """
#         N = len(particles)
#         lpxi = self.logp_xnext_xi_max(particles, u, t)
#         (Az, _fz, Qz, _, _, _) = self.get_lin_pred_dynamics_int(particles, u, t)
#         lpz = numpy.empty_like(lpxi)
#         (_xil, _zl, Pl) = self.get_states(particles)
#         nx = len(Qz[0])
#         for i in range(N):
#             # Predict z_{t+1}
#             Pn = Az[i].dot(Pl[i]).dot(Az[i].T) + Qz[i]
#             lpz[i] = -0.5 * nx * math.log(2 * math.pi) + numpy.linalg.slogdet(Pn)[1]
#         lpmax = numpy.max(lpxi + lpz)
#         return lpmax
# 
#     @abc.abstractmethod
#     def logp_xnext_xi_max(self, particles, u, t):
#         """
#         Maximum for nonlinear part of the logp_xnext, called from
#         logp_xnext_max
# 
#         Args:
#          - particles  (array-like): Model specific representation
#            of all particles, with first dimension = N (number of particles)
#          - u (array-like): input signal
#          - t (float): time stamp
# 
#          Returns:
#           (array-like) max value for each particle
#         """
#         pass


import numpy
import math
import pyparticleest.simulator
import pyparticleest.utils.kalman as kalman
import matplotlib.pyplot as plt
import scipy.stats
from pyparticleest.models.hierarchial import HierarchicalRSBase
from builtins import range



def generate_dataset(steps, P0_xi, P0_z, Q_xi, Q_z, R_xi, R_z):
    xi = numpy.zeros((1, steps + 1))
    z = numpy.zeros((2, steps + 1))
    y = numpy.zeros((steps, 2))
    xi[:, 0] = numpy.random.normal(0.0, math.sqrt(P0_xi))
    z[:, 0] = numpy.random.multivariate_normal(numpy.zeros((2,)), P0_z)
    for k in range(1, steps + 1):
        xi[:, k] = xi[:, k - 1] + numpy.random.normal(0.0, Q_xi)
        Ak = numpy.asarray(((math.cos(xi[:, k - 1]), math.sin(xi[:, k - 1])),
                            (-math.sin(xi[:, k - 1]), math.cos(xi[:, k - 1]))
                            ))
        z[:, k] = Ak.dot(z[:, k - 1]) + \
            numpy.random.multivariate_normal(numpy.zeros((2,)), Q_z)
        C = numpy.asarray(((math.cos(xi[:, k - 1]), math.sin(xi[:, k - 1]))))
        y[k - 1, 0] = xi[:, k] + numpy.random.normal(0.0, math.sqrt(R_xi))
        y[k - 1, 1] = C.dot(z[:, k]) + numpy.random.normal(0.0, math.sqrt(R_z))

    x = numpy.vstack((xi, z))
    return (x, y)



# Path: pyparticleest/models/hierarchial.py
# class HierarchicalRSBase(HierarchicalBase, FFBSiRS):
#     def __init__(self, **kwargs):
#         super(HierarchicalRSBase, self).__init__(**kwargs)
# 
#     def logp_xnext_max(self, particles, u, t):
#         """
#         Calculate maximum value of the logp_xnext function, used for
#         rejection sampling
# 
# 
#         Args:
#          - particles  (array-like): Model specific representation
#            of all particles, with first dimension = N (number of particles)
#          - u (array-like): input signal
#          - t (float): time stamp
# 
#          Returns:
#           (float) max value over all particle and possible future states
#         """
#         N = len(particles)
#         lpxi = self.logp_xnext_xi_max(particles, u, t)
#         (Az, _fz, Qz, _, _, _) = self.get_lin_pred_dynamics_int(particles, u, t)
#         lpz = numpy.empty_like(lpxi)
#         (_xil, _zl, Pl) = self.get_states(particles)
#         nx = len(Qz[0])
#         for i in range(N):
#             # Predict z_{t+1}
#             Pn = Az[i].dot(Pl[i]).dot(Az[i].T) + Qz[i]
#             lpz[i] = -0.5 * nx * math.log(2 * math.pi) + numpy.linalg.slogdet(Pn)[1]
#         lpmax = numpy.max(lpxi + lpz)
#         return lpmax
# 
#     @abc.abstractmethod
#     def logp_xnext_xi_max(self, particles, u, t):
#         """
#         Maximum for nonlinear part of the logp_xnext, called from
#         logp_xnext_max
# 
#         Args:
#          - particles  (array-like): Model specific representation
#            of all particles, with first dimension = N (number of particles)
#          - u (array-like): input signal
#          - t (float): time stamp
# 
#          Returns:
#           (array-like) max value for each particle
#         """
#         pass


# Path: test/manual/filtering/hierarchical_model.py
import numpy
import math
import pyparticleest.simulator
import pyparticleest.utils.kalman as kalman
import matplotlib.pyplot as plt
import scipy.stats
from pyparticleest.models.hierarchial import HierarchicalRSBase
from builtins import range




def generate_dataset(steps, P0_xi, P0_z, Q_xi, Q_z, R_xi, R_z):
    xi = numpy.zeros((1, steps + 1))
    z = numpy.zeros((2, steps + 1))
    y = numpy.zeros((steps, 2))
    xi[:, 0] = numpy.random.normal(0.0, math.sqrt(P0_xi))
    z[:, 0] = numpy.random.multivariate_normal(numpy.zeros((2,)), P0_z)
    for k in range(1, steps + 1):
        xi[:, k] = xi[:, k - 1] + numpy.random.normal(0.0, Q_xi)
        Ak = numpy.asarray(((math.cos(xi[:, k - 1]), math.sin(xi[:, k - 1])),
                            (-math.sin(xi[:, k - 1]), math.cos(xi[:, k - 1]))
                            ))
        z[:, k] = Ak.dot(z[:, k - 1]) + \
            numpy.random.multivariate_normal(numpy.zeros((2,)), Q_z)
        C = numpy.asarray(((math.cos(xi[:, k - 1]), math.sin(xi[:, k - 1]))))
        y[k - 1, 0] = xi[:, k] + numpy.random.normal(0.0, math.sqrt(R_xi))
        y[k - 1, 1] = C.dot(z[:, k]) + numpy.random.normal(0.0, math.sqrt(R_z))

    x = numpy.vstack((xi, z))
    return (x, y)



class Model(HierarchicalRSBase):