include("../src/Util_multi_RJ.jl")
using ReversibleJumpMultipleParticle

function run_rjsmc(N, lbs_set, ubs_set, obj_funcs, ϵ₀, Δϵ_tol, α, Δϵ_window, output;
                   save_interval=0, target=0.234 , parallel=false, N_sample=100, γscale=1., nuscale=1.)
    mkdir(output)
    lap_start = time()
    lap_end = 0
    N_models = length(obj_funcs)
    stepsize_flags = zeros(N_models)
    model_indecies = rand(collect(1:N_models), N)
    particles = generate_init_particles(N, lbs_set, ubs_set, obj_funcs, model_indecies)

    γs = ones(N_models)*γscale
    nus = ones(N_models)*nuscale
    ϵ_old = ϵ₀
    ϵ_new = 0.0
    cnt = 0
    terminate = false
    ESS=N
    Δϵ_mean = zeros(Δϵ_window)
    logfilename = joinpath(output, "log.csv")

    while terminate==false
        lap_end = time()
        lap = lap_end - lap_start
        lap_start = lap_end
        open(logfilename, "a") do fp
            writecsv(fp, Array([cnt, ϵ_old, ESS, lap])')
        end
        cnt += 1

        ϵ_new = calcPA(particles, α)
        Δϵ = ϵ_old - ϵ_new
        Δϵ_mean[2:end] .= Δϵ_mean[1:end-1]
        Δϵ_mean[1] = Δϵ
        if mean(Δϵ_mean) < Δϵ_tol && cnt > (Δϵ_window+100)
            terminate=true
        end

        # resampling
        candidates = deepcopy(particles)
        while ESS < N
            candidates = vcat(candidates, resample(particles, N-ESS))
            ESS = length(candidates)
        end

        F_old = get_obj_func_vals(candidates)
        candidates = reset_weights(candidates)
        model_indecies = get_model_indecies(candidates)
        accept_flags = zeros(length(candidates))
        if parallel == true
            candidates, accept_flags = MCstep_all_parallel_rj(candidates, ϵ_old)
        else
            candidates, accept_flags = MCstep_all_rj(candidates, ϵ_old)
        end

        model_indecies = get_model_indecies(candidates)
        particles = update_weights(candidates, F_old, ϵ_old, ϵ_new)
        particles = drop_zero_weights(particles)
        particles = normalize_weights(particles)


        acceptance = length(particles) / length(candidates)
        MH_acceptances = zeros(length(obj_funcs))
        MH_acceptances .= NaN
        for id in unique(model_indecies)
            flag_temp = accept_flags[model_indecies .== id]
            MH_acceptances[id] = sum(flag_temp)/length(flag_temp)
        end
        open(joinpath(output, "acceptance.csv"), "a") do fp
            writecsv(fp, Array([MH_acceptances' acceptance]))
        end

        model_indecies = get_model_indecies(particles)
        for id in unique(model_indecies)
            if MH_acceptances[id] < target
                stepsize_flags[id] = 1
            end
        end

        for id in unique(model_indecies)
            particles_subset = get_subset(particles, id)
            if stepsize_flags[id] == 0
                particles_subset, γs[id], nus[id] = update_stepsize_inRKHS_all(particles_subset, γs[id], nus[id], target,
                                                                               MH_acceptances[id], parallel=parallel,
                                                                               N=N_sample, adaptive=false)
                particles[model_indecies .== id] = particles_subset
            else
                particles_subset, γs[id], nus[id] = update_stepsize_inRKHS_all(particles_subset, γs[id], nus[id], target,
                                                                               MH_acceptances[id], parallel=parallel,
                                                                               N=N_sample, adaptive=true)
                particles[model_indecies .== id] = particles_subset
            end
        end

        proportion = get_proportion_of_models(particles, length(obj_funcs))
        open(joinpath(output, "proportion.csv"), "a") do fp
            writecsv(fp, vcat(ϵ_new, proportion)')
        end

        if length(particles) == 0
            println("particles were vanished at ϵ=$ϵ_new")
            break
        end

        if save_interval==0
            nothing
        elseif mod(cnt, save_interval) == 0 && terminate == false
            model_indecies = get_model_indecies(particles)
            for id in unique(model_indecies)
                particles_subset = get_subset(particles, id)
                est_params = (lbs_set[id].*(ubs_set[id]./lbs_set[id]).^(hcat(get_values(particles_subset)...)./10))'
                open(joinpath(output, @sprintf("est_param_%d_%05d.csv", id, cnt)), "w") do fp
                    writecsv(fp, est_params)
                end
            end

        elseif terminate == true
            ESS = length(particles)
            lap_end = time()
            lap = lap_end - lap_start
            open(logfilename, "a") do fp
                writecsv(fp, Array([cnt, ϵ_old, ESS, lap])')
            end

            model_indecies = get_model_indecies(particles)
            for id in unique(model_indecies)
                particles_subset = get_subset(particles, id)
                est_params = (lbs_set[id].*(ubs_set[id]./lbs_set[id]).^(hcat(get_values(particles_subset)...)./10))'
                open(joinpath(output, @sprintf("est_param_%d_%05d.csv", id, cnt)), "w") do fp
                    writecsv(fp, est_params)
                end
            end

            proportion = get_proportion_of_models(particles, length(obj_funcs))
            open(joinpath(output, "proportion.csv"), "a") do fp
                writecsv(fp, vcat(ϵ_new, proportion)')
            end
            break
        end

        ϵ_old = ϵ_new
        ESS = length(particles)
    end

end
