# 2017.07.21: Copied from the IJCNN codebase
#

module BayesChpOrig

# using Metadata

export bayes_chp_demo, bayes_detector, pccf_bayes_detector, extract_changes

function constant_hazard(r, lambda)
    p = 1.0 / lambda * ones(size(r))
end

"Non-standardized Student's t-distribution"
function studentpdf(x, mu, var, nu)
    ## As in the Matlab code by author:
    # c = exp(gammaln(nu/2 + 0.5) - gammaln(nu/2)) .* (nu.*pi.*var).^(-0.5);
    # p = c .* (1 + (1./(nu.*var)).*(x-mu).^2).^(-(nu+1)/2);
    c = exp.(lgamma.(nu/2 + 0.5) - lgamma.(nu/2)) .* (nu .* pi .* var).^(-0.5)
    p = c .* (1 + (1.0 ./ (nu .* var)) .* (x-mu).^2).^(-(nu+1)/2)
end

function extract_changes_from_output_naive(seq)
    chps = []
    for i in 2:length(seq)
        if seq[i] < seq[i-1]
            push!(chps, i)
        end
    end
    chps
end


"Function to extract changepoints locations from the bayes_detector() output"
function extract_changes(seq)
    #
    # seq[j] contains current estimate
    # of the location of the last change at the time moment j
    #
    n = length(seq)
    chps = zeros(n)
    for i in 1:n
        chps[i] = i - seq[i]
    end
    out = unique(chps)
    out = map(x->convert(Int,x), out)
end
function test_extract_changes()
    inp1 = [0,1,2,3,4,3]
    rt = extract_changes(inp1)
    print("... test changes extractor ", rt)
end



"""
Dynamic PCCF-based Bayesian detector.

h ~ 1/lambda => larger lambda correspondes to the smaller h

threshold : threshold for PCCF
h_min     : lambda value when PCCF <= threshold;
h_max     : lambda value when PCCF > threshold
h_common  : is either equal to h_min or h_max - the value to be used after max_d moment
max_d     : Max delay, after which we don't use PCCF's predictions

- V.0.1 (04 Nov 2016)
    Added max_d parameter to specify till which
    moment PCCF predictions will be used.

"""
function pccf_bayes_detector(in_sig, lambda_t, h_min, h_max, threshold; max_d = Nullable{Int64}(), h_common = Nullable{Float64}())
    @assert length(in_sig) == length(lambda_t)
    if ! isnull(h_common)
        @assert h_common == h_min || h_common == h_max
    end
    T       = length(in_sig)
    mu0     = 0.0
    kappa0  = 1
    alpha0  = 1
    beta0   = 1
    CP      = [0]
    R       = zeros(T+1, T+1)
    R[1, 1] = 1
    muT     = mu0
    kappaT  = kappa0
    alphaT  = alpha0
    betaT   = beta0
    maxes   = zeros(T+1)

    if isnull(max_d)
        max_d = T
    end

    for t=1:T
        predprobs = studentpdf(in_sig[t], muT, betaT .* (kappaT+1)./ (alphaT.*kappaT), 2*alphaT)

        if t <= max_d
            if lambda_t[t] <= threshold
                lambda = h_min
            else
                lambda = h_max
            end

        else
            lambda = h_common
        end


        H             = constant_hazard([1:t;], lambda)
        R[2:t+1, t+1] = R[1:t, t] .* predprobs .* (1-H)
        R[1, t+1]     = sum(R[1:t, t] .* predprobs .* H)
        R[:, t+1]     = R[:, t+1] / sum(R[:, t+1])
        muT0          = vcat(mu0   , (kappaT .* muT + in_sig[t]) ./ (kappaT + 1))
        kappaT0       = vcat(kappa0, kappaT + 1)
        alphaT0       = vcat(alpha0, alphaT + 0.5)
        betaT0        = vcat(beta0 , kappaT + (kappaT .* (in_sig[t] - muT).^2) ./ (2* (kappaT + 1)) )
        muT           = muT0
        kappaT        = kappaT0
        alphaT        = alphaT0
        betaT         = betaT0
        maxes[t]      = findmax(R[:,t])[2]
    end
    R = -log.(R)
    maxes = map(x -> convert(Int, x), maxes)
    maxes, R
end

#--------------------------------------------------

"Bayesian detector"
function bayes_detector(in_sig; lambda=200)
    T      = length(in_sig)
    # lambda = 200
    mu0    = 0.0
    kappa0 = 1
    alpha0 = 1
    beta0  = 1
    #X = zeros(T)
    CP = [0]
    R = zeros(T+1, T+1)
    R[1, 1] = 1
    #
    muT    = mu0
    kappaT = kappa0
    alphaT = alpha0
    betaT  = beta0
    #
    maxes = zeros(T+1)
    #
    for t=1:T
        # muT - is a vector of all possible hypos
        # We calc. predictive distribution for each hypothesis
        predprobs = studentpdf(in_sig[t], muT, betaT .* (kappaT+1)./ (alphaT.*kappaT), 2*alphaT)

        H = constant_hazard([1:t;], lambda)

        R[2:t+1, t+1] = R[1:t, t] .* predprobs .* (1-H)
        R[1, t+1] = sum(R[1:t, t] .* predprobs .* H)
        R[:, t+1] = R[:, t+1] / sum(R[:, t+1])

        ## Update statistics
        muT0    = vcat(mu0   , (kappaT .* muT + in_sig[t]) ./ (kappaT + 1))
        kappaT0 = vcat(kappa0, kappaT + 1)
        alphaT0 = vcat(alpha0, alphaT + 0.5)
        betaT0  = vcat(beta0 , kappaT + (kappaT .* (in_sig[t] - muT).^2) ./ (2* (kappaT + 1)) )

        muT    = muT0
        kappaT = kappaT0
        alphaT = alphaT0
        betaT  = betaT0

        maxes[t] = findmax(R[:,t])[2]
    end
    R = -log.(R)
    maxes = map(x->convert(Int,x), maxes)
    maxes, R
end

function bayes_chp_demo()
    sig    = readdlm("./data/sig.dat")
    T      = length(sig) # 1000
    lambda = 200
    mu0    = 0.0
    kappa0 = 1
    alpha0 = 1
    beta0  = 1
    #X = zeros(T)
    CP = [0]

    R = zeros(T+1, T+1)
    R[1, 1] = 1

    muT    = mu0
    kappaT = kappa0
    alphaT = alpha0
    betaT  = beta0

    maxes = zeros(T+1)

    for t=1:T
        # muT - is a vector of all possible hypos
        # We calc. predictive distribution for each hypothesis
        predprobs = studentpdf(sig[t], muT, betaT .* (kappaT+1)./ (alphaT.*kappaT), 2*alphaT)

        H = constant_hazard([1:t;], lambda)

        R[2:t+1, t+1] = R[1:t, t] .* predprobs .* (1-H)
        R[1, t+1] = sum(R[1:t, t] .* predprobs .* H)
        R[:, t+1] = R[:, t+1] / sum(R[:, t+1])

        ## Update statistics
        muT0    = vcat(mu0   , (kappaT .* muT + sig[t]) ./ (kappaT + 1))
        kappaT0 = vcat(kappa0, kappaT + 1)
        alphaT0 = vcat(alpha0, alphaT + 0.5)
        betaT0  = vcat(beta0 , kappaT + (kappaT .* (sig[t] - muT).^2) ./ (2* (kappaT + 1)) )

        muT    = muT0
        kappaT = kappaT0
        alphaT = alphaT0
        betaT  = betaT0

        maxes[t] = findmax(R[:,t])[2]
    end
    R = -log.(R)
    writedlm(outdir * "Rmatrix.txt", R, ",")
    writedlm(outdir * "Rmaxes.txt", maxes, ",")
    print(" write into ", outdir)
end

end
