include("Util_single.jl")
include("Util_RKHS.jl")


@everywhere module MultipleParticle
using StatsBase
using KernelDensity
using SingleParticle
using KernelTrick

export generate_init_particles, generate_init_particles2, resample, get_values, get_weights, get_obj_func_vals, get_stepsizes,
       normalize_weights, reset_weights, drop_zero_weights, update_weights, calc_ESS, calcPA,
       MCstep_all, MCstep_all2, MCstep_all_parallel, MCstep_all_parallel2, update_stepsize_inRKHS_all


function generate_init_particles(N, lbs, ubs, obj_func)
    """
    N: number of samplers
    """
    particles = [init_particle(lbs, ubs, obj_func, 1 / N) for i in 1:N]
    return particles
end

function resample(particles)
    N = length(particles)
    wv = [particle.weight for particle in particles]
    wv = Weights(wv)
    samples = similar(particles)
    for i in 1:N
        samples[i] = deepcopy(sample(particles, wv))
    end
    return samples
end

function resample(particles, N)
    wv = [particle.weight for particle in particles]
    wv = Weights(wv)
    samples = Array{eltype(particles), 1}(N)
    for i in 1:N
        samples[i] = deepcopy(sample(particles, wv))
    end
    return samples
end

function get_values(particles)
    return [particle.x_present for particle in particles]
end

function get_weights(particles)
    return [particle.weight for particle in particles]
end

function get_obj_func_vals(particles)
    return [particle.obj_func_val for particle in particles]
end

function get_stepsizes(particles)
    return [particle.stepsize for particle in particles]
end

function normalize_weights(particles)
    candidates = deepcopy(particles)
    wv = [candidate.weight for candidate in candidates]
    weight_sum = sum(wv)
    for candidate in candidates
        candidate.weight = candidate.weight / weight_sum
    end
    return candidates
end

function reset_weights(particles)
    N = length(particles)
    for particle in particles
        particle.weight = 1 / N
    end
    return particles
end

function drop_zero_weights(particles)
    wv = get_weights(particles)
    particles_new = eltype(particles)[]
    for (i, w) in enumerate(wv)
        if w > 0
            push!(particles_new, particles[i])
        end
    end
    return particles_new
end

function update_weights(particles,
                        obj_func_val_old::Array{Float64,1},
                        ϵ_old::Float64,
                        ϵ_new::Float64)
    for (i, particle) in enumerate(particles)
        W = particle.weight
        F_new = particle.obj_func_val < ϵ_new
        F_old = obj_func_val_old[i] < ϵ_old
        particle.weight = W * (F_new / F_old)
    end
    return particles
end

function calc_ESS(particles)
    """
    NOTE: weights must be normalized before.
    """
    wv = get_weights(particles)
    wv = wv / sum(wv)
    ESS = sum(wv.^2)^(-1)
    return ESS
end

function calcPA(particles, α)
    F = get_obj_func_vals(particles)
    return quantile(F, α)
end

function MCstep_all(particles, ϵ)

    # kernel density estimation of prior
    params = hcat(get_values(particles)...)'
    prior = [kde(params[:, i], npoints=16) for i in 1:size(params)[2]]

    N = length(particles)
    flags = zeros(length(particles))
    for i in 1:N
        particles[i], flags[i] = MCstep(particles[i], prior, ϵ)
    end
    return particles, flags
end

function MCstep_all_parallel(particles, ϵ)
    """
    to fetch results of @parallel for macro from several worker process,
    SharedArray must be required.
    """
    N = length(particles)

    # kernel density estimation of prior
    params = hcat(get_values(particles)...)'
    proposed = SharedArray{Float64}(size(params))
    obj_func_vals = SharedArray{Float64}(N)
    flags = SharedArray{Float64}(N)
    indecies = SharedArray{Int}(N)

    prior = [kde(params[:, i], npoints=16) for i in 1:size(params)[2]]

    @sync @parallel for i in 1:N
        indecies[i] = i
        proposed[i,:], obj_func_vals[i], flags[i] = MCstep_parallel(particles[i], prior, ϵ)
    end

    for i in indecies
        particles[i].x_present = proposed[i,:]
        particles[i].obj_func_val = obj_func_vals[i]
    end

    return particles, flags
end

function update_stepsize_inRKHS_all(particles, γ, nu, acceptance_target, acceptance; parallel=false, N=100, adaptive=false)
    if length(particles) < 3
        nothing
    else
        N = minimum([length(particles), N])
        sampled = sample(particles, N, replace=false)
        σ = bandwidth(sampled)
        if σ < 0
            nothing
        else
            num_param = length(particles[1].x_present)
            C = eye(N)-1/N
            for (i, particle) in enumerate(particles)
                sampled = sample(particles, N, replace=false)
                M = GaussianKernelInRKHS_M(sampled, particle, σ)
                Σ = κ_KAMH(sampled, particle, C, M, γ=γ, nu=nu)
                particle.cov = Σ
            end
            if adaptive==true
                γ = γ * (1. + rand()*(acceptance - acceptance_target))
                nu = nu * (1. + rand()*(acceptance - acceptance_target))
            end
        end
    end
    return particles, γ, nu
end

end
