'''
SQ-UKF implementations for Parameter Estimation

w/ additive parameters & observation noises
'''

import numpy as np
from numpy import ma

from ..SKLearn_utils import array1d, array2d, get_params, preprocess_arguments, check_random_state
from ..KF_utils import _determine_dimensionality, _last_dims, _arg_or_default
from .SQ_UKF_P_utilities import noise_ff, _reconstruct_covariances, Moments, moments2points, \
    unscented_filter_predict_P, unscented_filter_correct, _additive_unscented_filter_P


class UnscentedMixin(object):
    '''Methods for SQ-UKF_P implementations'''
    
    def __init__(self, observation_functions=None, observation_covariance=None,
            initial_state_mean=None, initial_state_covariance=None,
            n_dim_state=None, n_dim_obs=None, random_state=None):

        # size of state & obs spaces
        n_dim_state = _determine_dimensionality(
            [(initial_state_covariance, array2d, -2),
             (initial_state_mean, array1d, -1)],
            n_dim_state
        )
        n_dim_obs = _determine_dimensionality(
            [(observation_covariance, array2d, -2)],
            n_dim_obs
        )

        # set params
        self.observation_functions = observation_functions
        self.observation_covariance = observation_covariance
        self.initial_state_mean = initial_state_mean
        self.initial_state_covariance = initial_state_covariance
        self.n_dim_state = n_dim_state
        self.n_dim_obs = n_dim_obs
        self.random_state = random_state


    def _initialize_parameters(self):
        '''Retrieve params if they exist, else replace w/ defaults'''

        arguments = get_params(self)
        defaults = self._default_parameters()
        converters = self._converters()

        processed = preprocess_arguments([arguments, defaults], converters)
        return (
            processed['observation_functions'],
            processed['observation_covariance'],
            processed['initial_state_mean'],
            processed['initial_state_covariance']
        )


    def _parse_observations(self, obs):
        '''Convert obs to their expected format'''
        
        obs = ma.atleast_2d(obs)
        if obs.shape[0] == 1 and obs.shape[1] > 1:
            obs = obs.T
        return obs


    def _converters(self):
        return {
            'observation_functions': array1d,
            'observation_covariance': array2d,
            'initial_state_mean': array1d,
            'initial_state_covariance': array2d,
            'n_dim_state': int,
            'n_dim_obs': int,
            'random_state': check_random_state,
        }



class AdditiveUnscentedKalmanFilter_P(UnscentedMixin):
    '''Implements SQ-UKF_P w/ additive noise

        w_0       w/   (mu_0, sigma_0)
        w_{t+1}   =    w_t + (0, Q)
        y_{t}     =    g_t(x_t, w_t) + (0, R)

    where :
        w : parameters to update, viewed as state update
        x : observation as noise point or exogeneous input
        y : observation

    The UKF Additive version is imposed by parameters estimation conditions.
    Complexity of O(Tn^3) w/ :
        T number of timesteps - or n_timesteps,
        n state space size.
        
    Parameters noise Q is processed either by forgetting factor (RLS-like) or Robbins-Monro process.
    Observations noise R is fixed.