import numpy as np
from timemachines.skatertools.utilities.conventions import Y_TYPE, A_TYPE, R_TYPE, E_TYPE, T_TYPE
from timemachines.skatertools.components.parade import parade
from timemachines.skatertools.utilities.nonemath import nonecast, nonecenter
from timemachines.skatertools.utilities.conventions import wrap


def batch_skater_factory(y: Y_TYPE, s, k: int, a: A_TYPE = None, t: T_TYPE = None, e: E_TYPE = None, r: R_TYPE = None,
                         iskater=None, iskater_kwargs: dict=None, emp_mass:float=0.0, emp_std_mass:float=1.0,
                         min_e = None, n_warm=10):
    """
           An easy way to expose "offline" batch-style forecasters as skaters
           This accumulates data points one at a time, and feed them en-masse to the batch forecaster
           It also tracks empirical errors
       
          :param  iskater      -  function wrapping the batch-style forecasting capability
          :param  min_e        -  threshold value triggering a call of the batch forecaster (not optional!)

                                     min_e = -np.inf    ... always call the batch forecaster after it is warm unless e=-np.inf
                                     min_e = 0          ... call if e > 0

          :param  emp_mass     -  how much to empirical mean (i.e. bias) versus forecast
          :param  emp_std_mass -  how much to weigh empirical std (by default it only uses empirical std, not the std
                                   that is generated by the batch forecaster)

    """
    assert iskater is not None,'You must supply iskater'
    if iskater_kwargs is None:
        iskater_kwargs = {}
    if min_e is None:
        print('You should probably supply min_e when using batch_skater_factory, since I don''t know how fast your forecaster is. Setting to 0')
        min_e=0
    assert 0 <= emp_mass <= 1
    assert 0 <= emp_std_mass <= 1

    y = wrap(y)
    a = wrap(a)

    using_a = a is not None
    using_t = t is not None
    if e is None:
        e = 0

    if not s.get('y'):
        # On the first invocation we initialize, and freeze the pattern of usage
        s = {'p': {},  # parade
             'dim_y': len(y),
             'dim_a': len(a) if a is not None else None,
             'using_a': using_a,
             'using_t': using_t,
             'k': k,
             'y': list(),  # historical y
             'a': list(),  # list of a known k steps in advance
             't': list(),
             'warm': False}  # indicates whether iskater has been called yet
    else:
        # Assert immutability of k, dimensions of y,a
        if s['y']:
            assert len(y) == s['dim_y']
            assert k == s['k']
        if s['a']:
            assert s['using_a'], 'skater was supplied a one time but not another'
            assert len(a) == s['dim_a'], 'skater was supplied different length a from one invocation to next'

    if y is None:
        # A convention for puking the state 
        return None, s, None
    else:
        s['y'].append(y)
        if a is not None:
            s['a'].append(a)
        if t is not None:
            assert isinstance(t, float), 'epoch time please'
            s['t'].append(t)

        # Decide whether to call batch
        warm_enough = len(s['y']) >= n_warm
        enough_a = (a is None) or (len(s['a']) > k)
        caller_wants = e >= min_e
        do_call = caller_wants and enough_a and warm_enough

        if do_call:
            # Contemporaneous y, t, a are supplied to the batch forecast, as this is usually the preference
            # Note that a will be k longer than y and t, since it is required for prediction k-steps ahead
            t_arg = s['t'][k:] if t is not None else None
            a_arg = s['a']
            y_arg = s['y'][k:]
            x, x_std = iskater(y=y_arg, k=k, a=a_arg, t=t_arg, **iskater_kwargs)
            s['warm'] = True  # Flag indicating a model has been fit
        else:
            # Default to naive forecast (or a smoother depending on emp_mass, emp_std_mass settings)
            x = [y[0]] * k
            x_std = [1] * k

        # Get running mean prediction errors from the prediction parade
        x_resid, x_resid_std, s['p'] = parade(p=s['p'], x=x, y=y[0])
        x_resid_std = nonecast(x_resid_std, 1.0)

        # Compute center of mass between bias-corrected and uncorrected predictions
        x_corrected = np.array(x_resid) + np.array(x)
        x_center = nonecenter(m=[emp_mass, 1 - emp_mass], x=[x_corrected, x])
        x_std_center = nonecenter(m=[emp_std_mass, 1 - emp_std_mass], x=[x_resid_std, x_std])

        return x_center, x_std_center, s
