F(t) = hcat(1.0)
G(t) = hcat(1.0)
Σ(t) = hcat(9.0)
Tau(t) = hcat(1.0)
μ0 = vcat(0.0)
Tau0 = hcat(100.0)

Y = [
    5.22896384735402,
    2.059848221254017,
    11.77696919584194,
    5.721925811242492,
    2.647894413351219,
    8.833040488128724,
    10.112474414603064,
    6.208953089526956,
    7.483218024927158,
    1.1450746725213778,
    2.3827281539488214,
    2.5154588929528154,
    7.6938119335857795,
    11.267170019825972,
    11.203900809640977,
    4.641402249798912,
    6.393560809006704,
    5.434277785253692,
    7.754107613915341,
    2.8018519282190573,
    8.844915959029999,
    6.179788030963477,
    8.986422350276268,
    8.129873884597485,
    10.464689898594969,
    7.43846896751153,
    6.195006607054435,
    7.684383628629019,
    6.432819127902091,
    9.784040586681458,
    3.9386281953008275,
    1.0524045623884053,
    5.2287370014479535,
    3.3793567409786296,
    9.109449328047496,
    2.4107833351064665,
    3.1442786924940815,
    8.694980615700167,
    9.34957815303634,
    1.9899759308134595,
    1.7676330607869204,
    6.164419273710474,
    7.378003270378701,
    3.730353213883039,
    8.15997706282842,
    3.912665621137947,
    2.8222087161027987,
    7.89108559137954,
    2.016432756238303,
    -0.7283326731242079,
    5.34191426139306,
    3.7422889485477566,
    6.879539434261266,
    1.697715484841611,
    5.7492808736295835,
    3.4683312150848638,
    7.897465524526299,
    9.358577329317843,
    9.01915967340802,
    5.098410508423983,
    10.161486865442141,
    12.905026331072838,
    11.193444598591357,
    5.284455585283422,
    5.403785641006497,
    2.984985087125228,
    -0.1108114821742765,
    3.8443055205297396,
    1.1335781598255243,
    7.313228382179326,
    1.1291030218984508,
    0.8188795644543472,
    13.044240581728545,
    4.551403331575773,
    5.182374381380807,
    4.064021944452551,
    13.385904690258663,
    2.7383489638189253,
    4.272321073129483,
    1.936449829003382,
    7.579915913135302,
    6.891186573209264,
    5.696714256363617,
    2.9401339639559136,
    4.045792031292042,
    -2.69331630340043,
    4.253305831682493,
    -6.440609695250053,
    1.2224659946193035,
    3.6232261847393623,
    5.212788618742227,
    -0.5932923334934714,
    6.036279124314067,
    6.862343503023183,
    -2.9726591794689665,
    4.290653945351619,
    -2.420184182764869,
    0.6612926367123761,
    0.6312700637171165,
    -0.041609972691161345,
    -0.24839219509891972,
]

Y = hcat(Y)
push_tuple = SMC.push_state_forward(vcat(0.0), G(1), Tau(1), μ0, Tau0)
@test push_tuple[1] ≈ vcat(0.0)
@test push_tuple[2] ≈ Tau(1) + Tau0

ll_tuple = SMC.calc_loglik_t(hcat(0.0), F(1), Σ(1), μ0, Tau0, Y[1, :])
@test size(ll_tuple[1]) == (1, 1)
@test size(ll_tuple[2]) == (1, 1)
@test isa(ll_tuple[3], Float64)

update_tuple = SMC.update_state(F(1), μ0, Tau0, ll_tuple[1], ll_tuple[2])
@test size(update_tuple[1]) == (1, 1)
@test size(update_tuple[2]) == (1, 1)
m, V, loglik_t = SMC.iterate_kalman_filter_mv(
    vcat(0.0),
    F(1),
    vcat(0.0),
    G(1),
    Σ(1),
    Tau(1),
    μ0,
    Tau0,
    @view(Y[1, :])
)
loglik, ms, vs =
    SMC.kalman_filter_mv((t) -> hcat(0.0), F, (t) -> hcat(0.0), G, Σ, Tau, μ0, Tau0, Y)

@assert size(ms) == (1, 101)
@assert size(vs) == (1, 1, 101)
@test loglik ≈ -275.019326809115

@inferred SMC.kalman_filter_mv(
    (t) -> hcat(0.0),
    F,
    (t) -> hcat(0.0),
    G,
    Σ,
    Tau,
    μ0,
    Tau0,
    Y,
)

model_nomutate = SMC.KalmanModel(
    mutate = false,
    F0 = (t) -> hcat(0.0),
    F1 = F,
    G0 = (t) -> hcat(0.0),
    G1 = G,
    Σ = Σ,
    Tau = Tau,
    μ0 = μ0,
    Tau0 = Tau0,
)

@inferred model_nomutate(Y)
ll_nomutate, state_nomutate = model_nomutate(Y)
@test ll_nomutate ≈ -275.019326809115
@test size(state_nomutate.ms) == (1, 101)
@test size(state_nomutate.Vs) == (1, 1, 101)
@test all(state_nomutate.ms ≈ ms)
@test all(state_nomutate.Vs ≈ vs)

# Test is a 3-dimensional concatenation
# Should be 3x the log-likelihood above

using LinearAlgebra
F0(t) = [0.0, 0.0, 0.0]
G0(t) = [0.0, 0.0, 0.0]
F(t) = diagm([1.0, 1.0, 1.0])
G(t) = diagm(([1.0, 1.0, 1.0]))
Σ(t) = 9.0 * diagm(([1.0, 1.0, 1.0]))
Tau(t) = diagm([1.0, 1.0, 1.0])
μ0 = vcat(0.0, 0.0, 0.0)
Tau0 = 100.0 * diagm([1.0, 1.0, 1.0])

y = [
    5.22896384735402,
    2.059848221254017,
    11.77696919584194,
    5.721925811242492,
    2.647894413351219,
    8.833040488128724,
    10.112474414603064,
    6.208953089526956,
    7.483218024927158,
    1.1450746725213778,
    2.3827281539488214,
    2.5154588929528154,
    7.6938119335857795,
    11.267170019825972,
    11.203900809640977,
    4.641402249798912,
    6.393560809006704,
    5.434277785253692,
    7.754107613915341,
    2.8018519282190573,
    8.844915959029999,
    6.179788030963477,
    8.986422350276268,
    8.129873884597485,
    10.464689898594969,
    7.43846896751153,
    6.195006607054435,
    7.684383628629019,
    6.432819127902091,
    9.784040586681458,
    3.9386281953008275,
    1.0524045623884053,
    5.2287370014479535,
    3.3793567409786296,
    9.109449328047496,
    2.4107833351064665,
    3.1442786924940815,
    8.694980615700167,
    9.34957815303634,
    1.9899759308134595,
    1.7676330607869204,
    6.164419273710474,
    7.378003270378701,
    3.730353213883039,
    8.15997706282842,
    3.912665621137947,
    2.8222087161027987,
    7.89108559137954,
    2.016432756238303,
    -0.7283326731242079,
    5.34191426139306,
    3.7422889485477566,
    6.879539434261266,
    1.697715484841611,
    5.7492808736295835,
    3.4683312150848638,
    7.897465524526299,
    9.358577329317843,
    9.01915967340802,
    5.098410508423983,
    10.161486865442141,
    12.905026331072838,
    11.193444598591357,
    5.284455585283422,
    5.403785641006497,
    2.984985087125228,
    -0.1108114821742765,
    3.8443055205297396,
    1.1335781598255243,
    7.313228382179326,
    1.1291030218984508,
    0.8188795644543472,
    13.044240581728545,
    4.551403331575773,
    5.182374381380807,
    4.064021944452551,
    13.385904690258663,
    2.7383489638189253,
    4.272321073129483,
    1.936449829003382,
    7.579915913135302,
    6.891186573209264,
    5.696714256363617,
    2.9401339639559136,
    4.045792031292042,
    -2.69331630340043,
    4.253305831682493,
    -6.440609695250053,
    1.2224659946193035,
    3.6232261847393623,
    5.212788618742227,
    -0.5932923334934714,
    6.036279124314067,
    6.862343503023183,
    -2.9726591794689665,
    4.290653945351619,
    -2.420184182764869,
    0.6612926367123761,
    0.6312700637171165,
    -0.041609972691161345,
    -0.24839219509891972,
];

Y = hcat(y, y, y)
push_tuple = SMC.push_state_forward(G0(1), G(1), Tau(1), μ0, Tau0)
@test push_tuple[1] ≈ vcat(0.0, 0.0, 0.0)
@test push_tuple[2] ≈ Tau(1) + Tau0

ll_tuple = SMC.calc_loglik_t(F0(1), F(1), Σ(1), μ0, Tau0, Y[1, :])
@test size(ll_tuple[1]) == (3,)
@test size(ll_tuple[2]) == (3, 3)
@test isa(ll_tuple[3], Float64)

update_tuple = SMC.update_state(F(1), μ0, Tau0, ll_tuple[1], ll_tuple[2])
@test size(update_tuple[1]) == (3,)
@test size(update_tuple[2]) == (3, 3)

loglik, _, _ = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y)
@test loglik ≈ -275.019326809115 * 3

Tau(t) = [
    1.0 0.5 0.0
    0.5 1.0 0.0
    0.0 0.0 1.0
]

# @test kalman_filter_loglik_mv(vcat(0.0, 0.0, 0.0), F, vcat(0.0, 0.0, 0.0), G, Σ, Tau, μ0, Tau0, Y) ≈ -818.388291668541

## Test filtered distribution
loglik, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y)
@test loglik ≈ -818.388291668541
@test size(ms) == (3, 101)
@test size(Vs) == (3, 3, 101)

Y_missing = Matrix{Union{Missing,Float64}}(deepcopy(Y));
breaks = [2, 5, 18];
Y_missing[breaks, :] .= missing;
loglik, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y_missing)

## The missing likelihood should be the same as if we estimated before/after each break
push_forward(m, v, t) = SMC.push_state_forward(G0(t), G(t), Tau(t), m, v)
ll_pt1, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y[1:1, :])
m_next, V_next = push_forward(push_forward(ms[:, end], Vs[:, :, end], 2)..., 3)

ll_pt2, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, m_next, V_next, Y[3:4, :])
m_next, V_next = push_forward(push_forward(ms[:, end], Vs[:, :, end], 5)..., 6)

ll_pt3, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, m_next, V_next, Y[6:17, :])
m_next, V_next = push_forward(push_forward(ms[:, end], Vs[:, :, end], 18)..., 19)

ll_pt4, ms, Vs = SMC.kalman_filter_mv(F0, F, G0, G, Σ, Tau, m_next, V_next, Y[19:end, :])

@test (ll_pt1 + ll_pt2 + ll_pt3 + ll_pt4) ≈ loglik


## Smoother
ms_smoothed, Vs_smoothed = SMC.kalman_smoother_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y_missing)
@test size(ms_smoothed) == (3, 101)
@test size(Vs_smoothed) == (3, 3, 101)
@inferred SMC.kalman_smoother_mv(F0, F, G0, G, Σ, Tau, μ0, Tau0, Y_missing)

model_nomutate = SMC.KalmanModel(
    mutate = false,
    F0 = F0,
    G0 = G0,
    F1 = F,
    G1 = G,
    Σ = Σ,
    Tau = Tau,
    μ0 = μ0,
    Tau0 = Tau0,
)

state_nomutate = SMC.smooth(model_nomutate, Y_missing);
@test all(ms_smoothed ≈ state_nomutate.ms_smoothed)
@test all(Vs_smoothed ≈ state_nomutate.Vs_smoothed)
@inferred SMC.smooth(model_nomutate, Y_missing);

## Draw posterior pass_threshold
using Random: MersenneTwister
Xs_smoothed =
    @inferred SMC.draw_posterior_path(MersenneTwister(23), ms, Vs, G0, G, Σ, Tau, μ0, Tau0)
