# module Tmp
using Devectorize
using StatsBase
function pf_bootstrap!(xp,w, y, N, g_density, f)
    T = length(y)
    xp[:,1] = 2*σw*randn(N)
    wT = slice(w,:,1)
    xT = slice(xp,:,1)
    fill!(wT,-log(N))
    g(y[1],xT, wT)
    wT -= logsumexp(wT)
    j = Array(Int64,N)
    #     tw = Float64(N)
    for t = 2:T
        xT1 = xT
        xT = slice(xp,:,t)
        wT1 = wT
        wT = slice(w,:,t)
        # Resample
        if t%2 == 0
            resample_systematic!(j,wT1,N)
            f(xT,xT1[j],t-1)
            fill!(wT,-log(N))
        else # Resample not needed
            f(xT,xT1,t-1)
            copy!(wT,wT1)
        end

        # Time update
        g(y[t],xT, wT)

        # Normalize weights
        wT -= logsumexp(wT)
        #         tw =  1/sum(w(:,t).*2)
    end

    return nothing
end

function pf_bootstrap(y, N, g_density, f)
    T = length(y)
    xp = Array(Float64,(N,T))
    w = Array(Float64,(N,T))
    xp[:,1] = 2*σw*randn(N)
    wT = slice(w,:,1)
    xT = slice(xp,:,1)
    fill!(wT,-log(N))
    g(y[1],xT, wT)
    wT -= logsumexp(wT)
    j = Array(Int64,N)
    #     tw = Float64(N)
    for t = 2:T
        xT1 = xT
        xT = slice(xp,:,t)
        wT1 = wT
        wT = slice(w,:,t)
        # Resample
        if t%2 == 0
            resample_systematic!(j,wT1,N)
            f(xT,xT1[j],t-1)
            fill!(wT,-log(N))
        else # Resample not needed
            f(xT,xT1,t-1)
            copy!(wT,wT1)
        end

        # Time update
        g(y[t],xT, wT)

        # Normalize weights
        wT -= logsumexp(wT)
        #         tw =  1/sum(w(:,t).*2)
    end

    return xp,w
end




function pf_bootstrap_nn(xp, w, wnn, y, N, g_density, f)
    T = length(y)
    xp[:,1] = 2*σw*randn(N)
    wT = slice(w,:,1)
    xT = slice(xp,:,1)
    fill!(wT,-log(N))
    g(y[1],xT, wT)
    wnn[:,t] = wT
    wT -= logsumexp(wT)
    j = Array(Int64,N)
    #     tw = Float64(N)
    for t = 2:T
        xT1 = xT
        xT = slice(xp,:,t)
        wT1 = wT
        wT = slice(w,:,t)
        # Resample
        if t%2 == 0
            resample_systematic!(j,wT1,N)
            f(xT,xT1[j],t-1)
            fill!(wT,-log(N))
        else # Resample not needed
            f(xT,xT1,t-1)
            copy!(wT,wT1)
        end

        # Time update
        g_density(y[t],xp[:,t], wT)
        wnn[:,t] = wT
        # Normalize weights
        #some strange error with this normalization method, I totally have no clue why it's working elsewhere but not here
        wT -= logsumexp(wT)
        #         @assert isapprox(nc,log(sum(exp(wnn[:,t]))))
        #         tw =  1/sum(w(:,t).*2)
    end

    return xp,w, wnn
end


function pf_aux_nn(xp, w, wnn, y, N, g_density, f)
    T = length(y)
    xp[:,1] = 2*σw*randn(N)
    w0 = fill(log(1/N),N)
    xp[:,1] = f(zeros(N),1)
    wnn[:,1] = w0 + g_density(y[1],xp[:,1])
    @devec w[:,1] = wnn[:,1] - log(sum(exp(wnn[:,1])))
    j = Array(Int64,N)
    #     tw = Float64(N)
    for t = 2:T

        xpT = xp[:,t-1]
        f(xpT,t-1)

        wT = copy(w[:,t-1])
        lambda = g_density(y[t],xpT)
        wT += lambda
        wTe = exp(wT)
        wTe /= sum(wTe)

        # Resample
        if t%2 == 0
            resample_systematic_exp!(j,wTe)
            xpT = xp[j,t-1]
            wT = copy(w0)
            lambdaT = lambda[j]
        else # Resample not needed
            xpT = xp[:,t-1]
            #             wT = copy(w[:,t-1])
            lambdaT = lambda
        end

        # Time update
        f(xpT,t-1)
        xp[:,t] = xpT
        wT += g_density(y[t],xp[:,t])
        wnn[:,t] = wT - lambdaT
        # Normalize weights
        #some strange error with this normalization method, I totally have no clue why it's working elsewhere but not here
        # #         offset = maximum(wT)
        # #         normConstant = 0.0
        # #         for i = 1:N
        # #             normConstant += exp(wT[i]-offset)
        # #         end
        # #         w[:,t] = wT - log(normConstant)+offset
        w[:,t] = wnn[:,t] - log(sum(exp(wnn[:,t])))
        #         tw =  1/sum(w(:,t).*2)
    end

    return xp,w, wnn
end


function pf_CSMC!(xp,w, y, N, g_density, f, xc )
    xp[:,1] = 2*σw*randn(N)
    xp[N,1] = xc[1]
    T = length(y)
    wT = slice(w,:,1)
    xT = slice(xp,:,1)
    fill!(wT,-log(N))
    g(y[1],xT, wT)
    wT -= logsumexp(wT)
    j = Array(Int64,N)
    #     tw = Float64(N)
    for t = 2:T
        xT1 = xT
        xT = slice(xp,:,t)
        wT1 = wT
        wT = slice(w,:,t)
        # Resample
        if t%2 == 0
            resample_systematic!(j,wT1,N)
            f(xT,xT1[j],t-1)
            fill!(wT,-log(N))
        else # Resample not needed
            f(xT,xT1,t-1)
            copy!(wT,wT1)
        end
        xp[N,t] = xc[t]
        # Time update
        g(y[t],xT, wT)

        # Normalize weights
        wT -= logsumexp(wT)
        #         tw =  1/sum(w(:,t).*2)
    end

    return nothing
end