using SMC:
    dt_smc2_estimation,
    smc_pmcmc_proposal_logdens,
    smc_pmcmc_proposal,
    density_tempered_pmcmc
using SMC.Models: LinearGaussianInfer, rprior, dprior, loglik_fn
using Base.Iterators


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)


m = LinearGaussianInfer(10, Y, 101)
ps = @inferred rprior(m, 10)
d1 = @inferred dprior(m, ps[1])
ds = @inferred dprior(m, ps)

l = @inferred loglik_fn(m, ps[2])
ls = @inferred broadcast(m, ps)

prior = p -> dprior(m, p)

lower_bound = (x, n) -> 0.0
upper_bound = (x, n) -> Inf
C = (2.38^2) / 2.0

rprop_pmh =
    (theta, mean, cov) -> begin
        thetavec = smc_pmcmc_proposal(
            vcat(theta...),
            mean,
            cov;
            lower_bound = lower_bound,
            upper_bound = upper_bound,
            C = C,
        )
        NamedTuple{keys(theta)}(Tuple(thetavec))
    end

dprop_pmh =
    (theta, theta_new, mean, cov) -> smc_pmcmc_proposal_logdens(
        vcat(theta...),
        vcat(theta_new...),
        mean,
        cov;
        lower_bound = lower_bound,
        upper_bound = upper_bound,
        C = C,
    )

# Currently type-unstable because of pmap
res = dt_smc2_estimation(ps, m, prior; rprop_pmh = rprop_pmh, dprop_pmh = dprop_pmh)
@test typeof(res[1]) == typeof(ps)
@test size(res[1]) == size(ps)
@test res[2] isa Vector{Float64}
@test size(res[2]) == (size(ps, 1),)
@test res[3] isa Array{Bool,3}
@test size(res[3], 1) == size(res[3], 2) == size(ps, 1)

# Should be type-stable
ξ = 0.5
pmcmc_out = @inferred density_tempered_pmcmc(
    ps,
    ls,
    m,
    prior,
    ξ,
    ξ;
    rprop_pmh = rprop_pmh,
    dprop_pmh = dprop_pmh,
)
