#
# Jeffrey D. Scargle et.al., 2013,
# "Studies In Astronomical Time Series Analysis. VI. Bayesian Block Representations",
# http://iopscience.iop.org/article/10.1088/0004-637X/764/2/167/meta
#
# To be used for adaptive ADWIN
#
# Bayesian blocks method joins bins in the regular histogram in an optimal way.
#
using Base.Test


"""
Eq. 19 from Scargle 2012
:param n_k: number of elements in each block
:param t_k: width of the block
:return: fitness value.

NOTE: Looks similar to the entropy function.
"""
function fitness_eq19(n_k, t_k)
    n_k * (log(n_k) - log(t_k))
end

function fit_prior(n, p0=0.05)
    4 - log(75.33 * p0 * n ^(-0.478))
end


"""
In this method the regular histogram is computed first and after that its bins
are joined using dynamic programming.

It is like a normal changedetection but along the Y-axis.
That is why last[] and best[] have lengths of the number of unique elements
(defining edges boundaries).

About the statement
 n_k = n_k = cumsum(x[end:-1:1])[end:-1:1]
 Here the twice reversing is used because e.g for the input
 [1, 1, 1, 2, 1]
the output should be
 [6, 5, 4, 3, 1]
which is a possible number of elements in the last block prior to r.
"""
function bayesian_blocks(t_input, p0_prior=0.05, delta=0.0)
    uniq_t = sort(unique(t_input))
    #... Start Calculate histogram
    edges = [uniq_t[1]-delta;
             0.5 * (uniq_t[2:end] + uniq_t[1:end-1]);
             uniq_t[end]+delta]

    # Count elements in blocks.
    x::Array{Int,1} = []
    for e in uniq_t
        push!(x, count(x -> x == e, t_input))
    end
    @assert length(x) == length(edges) - 1
    #... End Calculate histogram
    # println("...1 edges, counts:", edges, x)
    block_lengths = edges[end] - edges[1:end-1]
    last::Array{Int,1} = zeros(length(uniq_t))
    best::Array{Float64,1} = zeros(length(uniq_t))
    n = length(uniq_t)
    ## Construct a vector of possible block lengths
    ## and count the number of elements in each block.
    for r = 1:n
        if r==1
            last[r] = 1
            best[r] = 0.0
        else
            ## Possible blocks widths prior to r
            t_k = block_lengths[1:r-1] - block_lengths[r] # ??? + delta
            ## Possible numbers of elements in blocks prior to r.
            n_k = cumsum(x[1:r-1][end:-1:1])[end:-1:1]
            # if r==5 @assert length(t_k) == length(n_k) end
            fit_vec = map(x -> fitness_eq19(x[1], x[2]), zip(n_k, t_k))
            a_r = zeros(r) ## a_r[1]==0 automatically; to make sure that the size is correct
            a_r[2:r] = fit_vec - fit_prior(n, p0_prior)
            a_r[2:r] += best[1:r-1]
            v_max, i_max = findmax(a_r)
            last[r] = i_max
            best[r] = v_max
        end
    end
    # last, best
    chps = extract_changepoints(last)
    edges_c = [edges[c] for c in chps]
    if edges_c[end]!=edges[end]
        edges_c = [edges_c; edges[end]]
    end
    edges_c
end



"""
    In this method an input is a is a pre-calculated histogram.
"""
function bayesian_blocks(edges::Array{Float64,1},
                         counts::Array{Int,1},
                         p0_prior=0.05)
    x = counts
    # println("... 2 edges, counts:", edges, x)
    @assert length(edges) - 1 == length(x)
    block_lengths = edges[end] - edges[1:end-1]
    uniq_t = zeros(length(edges)-1) ## to be readable along other method
    last::Array{Int,1} = zeros(length(uniq_t))
    best::Array{Float64,1} = zeros(length(uniq_t))
    n=length(uniq_t)
    # println("...2: n=$n")
    ## Construct a vector of possible block lengths
    ## and count the number of elements in each block.
    for r = 1:n
        if r==1
            last[r] = 1
            best[r] = 0.0
        else
            ## Possible blocks widths prior to r
            t_k = block_lengths[1:r-1] - block_lengths[r] # ??? + delta
            ## Possible numbers of elements in blocks prior to r.
            n_k = cumsum(x[1:r-1][end:-1:1])[end:-1:1]
            if r==5 @assert length(t_k) == length(n_k) end
            fit_vec = map(x -> fitness_eq19(x[1], x[2]), zip(n_k, t_k))
            a_r = zeros(r) ## a_r[1]==0 automatically; to make sure that the size is correct
            a_r[2:r] = fit_vec - fit_prior(n, p0_prior)
            a_r[2:r] += best[1:r-1]
            v_max, i_max = findmax(a_r)
            last[r] = i_max
            best[r] = v_max
        end
    end
    chps=extract_changepoints(last)
    edges_c = [edges[c] for c in chps]
    if edges_c[end]!=edges[end]
        edges_c = [edges_c; edges[end]]
    end
    edges_c
    # last
end


function extract_changepoints(last::Array{Int,1})
    # cp1 = last[end]
    # cp2 = last[cp1]
    chps::Array{Int,1} = []
    next_chps = last[end]
    while next_chps != 1
        push!(chps, next_chps)
        if next_chps != last[next_chps]
            next_chps = last[next_chps]
        else
            next_chps = last[next_chps-1]
        end
    end
    push!(chps,1)
    sort(chps)
end


function runtests()
    println("... Start tests()")
    ret1 = bayesian_blocks([1, 2, 3, 4, 5, 5, 1.1, 1.1], 0.05)

    ret2 = bayesian_blocks([1.0, 1.05, 1.55, 2.5, 3.5, 4.5, 5.0],
                           [1, 2, 1, 1, 1, 2], 0.05)
    @test ret1 == ret2
    # test extract_changepoints
    @test extract_changepoints([1,1,1,2,2,2,5]) == [1,2,5]
    @test extract_changepoints([1,1,1,2,2,2,5,3]) == [1,3]
    println("... End tests()")
end
# runtests()
