# Helper functions for working with KFAS output

import numpy as np
import pandas as pd

from statsmodels.tools.tools import Bunch


def parse(path, ssm):
    # Dimensions
    n = ssm.nobs
    p = ssm.k_endog
    m = ssm.k_states
    r = ssm.k_posdef

    p2 = p**2
    m2 = m**2
    mp = m * p
    r2 = r**2

    # Extract the different pieces of output from KFAS
    kfas = pd.read_csv(path)
    components = [("r", m), ("r0", m), ("r1", m),
                  ("N", m2), ("N0", m2), ("N1", m2), ("N2", m2),
                  ("m", p), ("v", p), ("F", p), ("Finf", p),
                  ("K", mp), ("Kinf", mp), ("a", m), ("P", m2),
                  ("Pinf", m2), ("att", m), ("Ptt", m2),
                  ("alphahat", m), ("V", m2), ("muhat", p),
                  ("V_mu", p2), ("etahat", r), ("V_eta", r2), ("epshat", p),
                  ("V_eps", p), ("llf", 1)]
    dta = {}
    ix = 0
    for key, length in components:
        dta[key] = kfas.iloc[:, ix:ix + length].fillna(0)
        dta[key].name = None
        ix += length

    # Reformat the KFAS output to compare with statsmodels output
    res = Bunch()
    d = len(kfas["Pinf_1"].dropna())

    # forecasts
    res["forecasts"] = dta["m"].values[:n].T
    res["forecasts_error"] = dta["v"].values[:n].T
    res["forecasts_error_cov"] = np.c_[
        [np.diag(x) for y, x in dta["F"].iloc[:n].iterrows()]].T
    res["forecasts_error_diffuse_cov"] = np.c_[
        [np.diag(x) for y, x in dta["Finf"].iloc[:n].iterrows()]].T

    res["kalman_gain"] = dta["K"].values[:n].reshape(n, m, p, order="F").T
    res["Kinf"] = dta["Kinf"].values[:n].reshape(n, m, p, order="F")

    # filtered
    res["filtered_state"] = dta["att"].values[:n].T
    res["filtered_state_cov"] = dta["Ptt"].values[:n].reshape(
        n, m, m, order="F").T
    # predicted
    res["predicted_state"] = dta["a"].values.T
    res["predicted_state_cov"] = dta["P"].values.reshape(
        n + 1, m, m, order="F").T
    res["predicted_diffuse_state_cov"] = dta["Pinf"].values.reshape(
        n + 1, m, m, order="F").T
    # loglike
    # Note: KFAS only gives the total loglikelihood
    res["llf_obs"] = dta["llf"].values[0, 0]

    # smoothed
    res["smoothed_state"] = dta["alphahat"].values[:n].T
    res["smoothed_state_cov"] = dta["V"].values[:n].reshape(
        n, m, m, order="F").T

    res["smoothed_measurement_disturbance"] = dta["epshat"].values[:n].T
    res["smoothed_measurement_disturbance_cov"] = np.c_[
        [np.diag(x) for y, x in dta["V_eps"].iloc[:n].iterrows()]].T
    res["smoothed_state_disturbance"] = dta["etahat"].values[:n].T
    res["smoothed_state_disturbance_cov"] = dta["V_eta"].values[:n].reshape(
        n, r, r, order="F").T

    # scaled smoothed estimator
    # Note: we store both r and r0 together as "scaled smoothed estimator"
    # while "scaled smoothed diffuse estimator" corresponds to r1
    res["scaled_smoothed_estimator"] = np.c_[dta["r0"][:d].T,
                                             dta["r"][d:].T][..., 1:]
    res["scaled_smoothed_diffuse_estimator"] = dta["r1"].values.T

    N0 = dta["N0"].values[:d].reshape(d, m, m, order="F")
    N = dta["N"].values[d:].reshape(n + 1 - d, m, m, order="F")
    res["scaled_smoothed_estimator_cov"] = np.c_[N0.T, N.T][..., 1:]
    # Have to do a more precise transpose for these since they may not be
    # symmetric
    res["scaled_smoothed_diffuse1_estimator_cov"] = dta["N1"].values.reshape(
        n + 1, m, m, order="F").transpose(1, 2, 0)
    res["scaled_smoothed_diffuse2_estimator_cov"] = dta["N2"].values.reshape(
        n + 1, m, m, order="F").transpose(1, 2, 0)

    # Save the results object for the tests
    return res
