text:
using JLD
using PyPlot
include("../src/util.jl")
using StatsBase
using VBMatrixFactorization

"""
    getBag(data, field, id)

Extracts a field from a MIL data variable given an id of a bag.
"""
function getBag(data::Dict{String,Any}, field::String, id::Int)
    s = size(data[field])
    inds = (data["bagids"] .== id)
    if length(s) == 1
        return data[field][inds]
    else
        return data[field][:, inds]
    end
end

"""
    getY(data, id)

Extracts a Y matrix from a MIL data variable given the id of a bag.
"""
function getY(data::Dict{String,Any}, id::Int)
    res = getBag(data, "fMat", id)
    return convert(Array{Float64, 2}, res)
end

"""
    getLabel(data, id)

Extracts label of a bag with index id.
"""
function getLabel(data::Dict{String,Any}, id::Int)
    return getBag(data, "y", id)[1]
end

"""
    get_matrices(data, bag_ids)

For given bag indices, extracts negative and positive bags and returns them as a matrix.
"""
function get_matrices(data::Dict{String,Any}, bag_ids)
    Y0 = 0
    Y1 = 0
    n0 = 0
    n1 = 0
    inds0 = 0
    inds1 = 0
    for i in bag_ids
        label = getLabel(data,i);
        mat = getY(data, i);
        if label == 1
            n1 += 1
            if Y1 == 0
                Y1 = mat;
                inds1 = repmat([i], size(mat)[2]);
            else
                Y1 = cat(2, Y1, mat);
                inds1 = cat(1, inds1, repmat([i], size(mat)[2]));
            end
        else
            n0 += 1
            if Y0 == 0
                Y0 = mat;
                inds0 = repmat([i], size(mat)[2]);
            else
                Y0 = cat(2, Y0, mat)
                inds0 = cat(1, inds0, repmat([i], size(mat)[2]));

            end
        end
    end

    # returns the concatenated negative and positive matrices, their numbers and indices of
    # the bags for slicing
    return Y0, Y1, n0, n1, inds0, inds1
end

"""
    train(data, bag_ids, solver, H, niter; verb::Bool = true)

Using data and a vector of bag ids, trains the classifier using a given solver.
The inner dimension of the factorization is H.

solver = basic - calls vbmf()
solver = sparse - calls vbmf_sparse()

"""
function train(data::Dict{String,Any}, bag_ids, solver::String, H::Int, niter::Int; eps::Float64 = 1e-6,
    verb::Bool = true, diag_var::Bool = false)
    Y0, Y1, n0, n1, inds0, inds1 = get_matrices(data, bag_ids)

    if n0 == 0 || n1 ==0
        warn("One of the input matrices is empty, ending.")
        return 0, 0
    end

    if verb
        println("Extracted $n0 negative and $n1 positive bags.")
    end

    L, M0 = size(Y0)
    L, M1 = size(Y1)
    max_restarts = 10 # how many times should at most be vbmf restarted

    if solver == "basic"
        res0 = VBMatrixFactorization.vbmf_init(Y0, H)
        VBMatrixFactorization.vbmf!(Y0, res0, niter, eps = eps, est_covs = true, est_var = true, verb = verb)
        res1 = VBMatrixFactorization.vbmf_init(Y1, H)
        VBMatrixFactorization.vbmf!(Y1, res1, niter, eps = eps, est_covs = true, est_var = true, verb = verb)
    elseif solver == "sparse"
        # this is to decide whether to compute full covariance of A or just the diagonal
        if (H*M0 > 800) || (H*M1 > 800)
            full_cov = false
        else
            full_cov = true
        end
        full_cov = false

        # do more random restarts if bad convergence
        nres = 0
        delta = 2*eps + 1.0
        while (nres < max_restarts) && (delta > 2*eps)
            res0 = VBMatrixFactorization.vbmf_sparse_init(Y0, H)
            delta = VBMatrixFactorization.vbmf_sparse!(Y0, res0, niter, eps = eps, diag_var = diag_var, verb = verb, full_cov = full_cov)
            if isnan(delta)
                delta = 2*eps + 1.0
            end
            nres += 1
        end

        nres = 0
        delta = 2*eps + 1.0
        while (nres < max_restarts) && (delta > 2*eps)
            res1 = VBMatrixFactorization.vbmf_sparse_init(Y1, H)
            delta = VBMatrixFactorization.vbmf_sparse!(Y1, res1, niter, eps = eps, diag_var = diag_var, verb = verb, full_cov = full_cov)
            if isnan(delta)
                delta = 2*eps + 1.0
            end
            nres += 1
        end
    else
        error("Unknown type of solver. Use 'basic' or 'sparse'.")
        return
    end

    return res0, res1
end

"""
    ols(Y::Array{Float64, 2}, B::Array{Float64, 2})

Solves min(||Y - B*X||^2) for unknown X.
"""
function ols(Y::Array{Float64, 2}, B::Array{Float64, 2})
    return inv(B'*B)*B'*Y;
end

"""
    rls(Y::Array{Float64, 2}, B::Array{Float64, 2}, lambda::Float64)

Solves min(||Y - B*X||^2 + lambda*||X||^2) for unknown X.
"""
function rls(Y::Array{Float64, 2}, B::Array{Float64, 2}, lambda::Float64)
    m, n = size(B)
    return inv(B'*B + lambda*eye(n))*B'*Y;
end

"""
    vbls(Y::Array{Float64, 2}, params, niter::Int;
    diag_var::Bool = false, full_cov::Bool = false)

Solves the factorization Y = BA^T + E for fixed B and CB that is stored in params.
"""
function vbls!(Y::Array{Float64, 2}, params, niter::Int;
    diag_var::Bool = false, full_cov::Bool = false)
    for iter in 1:niter
        if typeof(params) == VBMatrixFactorization.vbmf_parameters
            VBMatrixFactorization.updateA!(Y, params)
            VBMatrixFactorization.updateCA!(params)
            VBMatrixFactorization.updateSigma2!(Y, params)
        elseif typeof(params) == VBMatrixFactorization.vbmf_sparse_parameters
            VBMatrixFactorization.updateA!(Y, params, full_cov = full_cov, diag_var = diag_var)
            VBMatrixFactorization.updateCA!(params)
            VBMatrixFactorization.updateSigma!(Y, params, diag_var = diag_var)
        elseif typeof(params) == VBMatrixFactorization.vbmf_dual_parameters
            VBMatrixFactorization.updateA!(Y, params, full_cov = full_cov, diag_var = diag_var)
            VBMatrixFactorization.updateCA!(params)
            VBMatrixFactorization.updateSigma!(Y, params, diag_var = diag_var)
        elseif typeof(params) == VBMatrixFactorization.vbmf_trial_parameters
            VBMatrixFactorization.updateA!(Y, params, full_cov = full_cov, diag_var = diag_var)
            VBMatrixFactorization.updateCA!(params)
            VBMatrixFactorization.updateSigma!(Y, params, diag_var = diag_var)
        end
    end
    # finally, compute the estimate of Y
    VBMatrixFactorization.updateYHat!(params)
    return params.AHat
end

"""
    copy_vbmf_params(Y::Array{Float64, 2}, old_params)

Creates a new instance of vbmf_parameters/vbmf_sparse_parameters from an old one
to be used by vbls classification. This cannot be done by simple copy()
since the dimension of the problem M is not the same for classification.
"""
function copy_vbmf_params(Y::Array{Float64, 2}, old_params)
    if typeof(old_params) == VBMatrixFactorization.vbmf_parameters
        # init a new structure
        # dont copy the labels and H1, that probably doesnt make any sense now!
        params = VBMatrixFactorization.vbmf_init(Y, old_params.H,
         sigma2 = old_params.sigma2)
        # copy the parameters that wont change
        params.BHat = copy(old_params.BHat)
        params.SigmaB = copy(old_params.SigmaB)
        params.CB = copy(old_params.CB)
        params.invCB = copy(old_params.invCB)
        return params
    elseif typeof(old_params) == VBMatrixFactorization.vbmf_sparse_parameters
        # init a new structure
        params = VBMatrixFactorization.vbmf_sparse_init(Y, old_params.H, alpha0 = old_params.alpha0,
                beta0 = old_params.beta0, gamma0 = old_params.gamma0, delta0 = old_params.delta0,
                eta0 = old_params.eta0, zeta0 = old_params.zeta0)
        # copy the parameters that wont change
        params.BHat = copy(old_params.BHat)
        params.SigmaB = copy(old_params.SigmaB)
        params.CB = copy(old_params.CB)
        params.gamma = old_params.gamma
        params.delta = copy(old_params.delta)
        params.invCB = copy(old_params.invCB)
        return params
    elseif typeof(old_params) == VBMatrixFactorization.vbmf_dual_parameters
        params = VBMatrixFactorization.vbmf_dual_init(Y, old_params.H, old_params.H0,
            gamma0 = old_params.gamma0, delta0 = old_params.delta0,
            eta0 = old_params.eta0, zeta0 = old_params.zeta0)
        # copy the parameters that wont change
        params.BHat = copy(old_params.BHat)
        params.SigmaB = copy(old_params.SigmaB)
        params.CB = copy(old_params.CB)
        params.gamma = old_params.gamma
        params.delta = copy(old_params.delta)
        # also the priors
        params.alpha00 = old_params.alpha00
        params.beta00 = old_params.beta00
        params.alpha01 = old_params.alpha01
        params.beta01 = old_params.beta01
        return params
    elseif typeof(old_params) == VBMatrixFactorization.vbmf_trial_parameters
        L, M = size(Y)
        # here we have just one of the special basis matrices
        params0 = VBMatrixFactorization.vbmf_trial_init(Y, old_params.H, old_params.H0, M,
            gamma0 = old_params.gamma0, delta0 = old_params.delta0,
            eta0 = old_params.eta0, zeta0 = old_params.zeta0)
        # copy the parameters that wont change
        params0.BHat = copy(old_params.BHat)
        params0.SigmaB = copy(old_params.SigmaB)
        params0.CB = copy(old_params.CB)
        params0.gamma = old_params.gamma
        params0.delta = copy(old_params.delta)
        # also the priors
        params0.alpha01 = old_params.alpha01
        params0.beta01 = old_params.beta01
        params0.alpha02 = old_params.alpha02
        params0.beta02 = old_params.beta02
        params0.alpha03 = 1e-10
        params0.beta03 = 1e-10

        params1 = VBMatrixFactorization.vbmf_trial_init(Y, old_params.H, old_params.H0, M,
            gamma0 = old_params.gamma0, delta0 = old_params.delta0,
            eta0 = old_params.eta0, zeta0 = old_params.zeta0)
        # copy the parameters that wont change
        params1.BHat = copy(old_params.BHat)
        params1.SigmaB = copy(old_params.SigmaB)
        params1.CB = copy(old_params.CB)
        params1.gamma = old_params.gamma
        params1.delta = copy(old_params.delta)
        # also the priors
        params1.alpha01 = old_params.alpha01
        params1.beta01 = old_params.beta01
        params1.alpha02 = old_params.alpha03
        params1.beta02 = old_params.beta03
        params1.alpha03 = 1e-10
        params1.beta03 = 1e-10

        return params0, params1
    end

end

### this is for the new classification

"""
    train_local(data, train_inds, H, H1, niter, eps = eps, verb = verb, diag_var = diag_var)

Trains the localized version of VBMF.
"""
function train_local(data, train_inds, H, H1, niter; eps = 1e-4, verb = false, diag_var = false)
    Y0_train, Y1_train = get_matrices(data, train_inds);
    Y_train = cat(2, Y0_train, Y1_train)
    L, M0 = size(Y0_train)
    L, M1 = size(Y1_train)
    L, M = size(Y_train)
    params = VBMatrixFactorization.vbmf_sparse_init(Y_train, H, H1 = H1, labels = Array(1:M0));

    max_restarts = 10
    nres = 0
    delta = norm(params.AHat) + norm(params.BHat)
    while (nres < max_restarts) && ((norm(params.AHat)) < delta) && ((norm(params.BHat)) < delta)
        d = VBMatrixFactorization.vbmf_sparse!(Y_train, params, niter, eps = eps, verb = verb, diag_var = diag_var);
        delta = 1e-2
        nres+=1
    end

    return params
end

"""
    train_dual(data, train_inds, H, H0, niter; eps = 1e-4, verb = false, diag_var = false)

Trains the classifier using the dual basis decomposition.
"""
function train_dual(data, train_inds, H, H0, niter; eps = 1e-4, verb = false, diag_var = false)
    Y0_train, Y1_train = get_matrices(data, train_inds)
    L, M0 = size(Y0_train)
    L, M1 = size(Y1_train)

    max_restarts = 10
    nres = 0
    delta = 1e-3
    if M0*H > 3200
        full_cov = false
    else
        full_cov = true
    end

    # use this on local machines
    maxM = Int(floor(3200/H))
    inds0 = sample(1:M0, min(maxM, M0), replace = false)
    Y0_train = Y0_train[:, inds0]
    full_cov = true

    params0 = VBMatrixFactorization.vbmf_dual_init(Y0_train, H, H0);
    while (nres < max_restarts) && (delta < 1e-2)
        params0 = VBMatrixFactorization.vbmf_dual_init(Y0_train, H, H0);
        d = VBMatrixFactorization.vbmf_dual!(Y0_train, params0, niter, eps = eps, verb = verb,
            diag_var = diag_var, full_cov = full_cov);
        delta = norm(params0.AHat) + norm(params0.BHat)
        nres+=1
    end
    if (delta < 1e-2)
        println("Sensible result in negative samples factorization not achieved after $(nres) retries!")
    end

    nres = 0
    delta = 1e-3
    if M1*H > 3200
        full_cov = false
    else
        full_cov = true
    end
#    if full_cov
#        println("computing full covariance")
#   end
    inds1 = sample(1:M1, min(maxM, M1), replace = false)
    Y1_train = Y1_train[:, inds1]
    full_cov = true
    params1 = VBMatrixFactorization.vbmf_dual_init(Y1_train, H, H0);
    while (nres < max_restarts) && (delta < 1e-2)
        params1 = VBMatrixFactorization.vbmf_dual_init(Y1_train, H, H0);
        d = VBMatrixFactorization.vbmf_dual!(Y1_train, params1, niter, eps = eps, verb = verb,
            diag_var = diag_var, full_cov = full_cov);
        delta = norm(params1.AHat) + norm(params1.BHat)
        nres+=1
    end
    if (delta < 1e-2)
        println("Sensible result in negative samples factorization not achieved after $(nres) retries!")
    end

    return params0, params1
end

"""
    factorize_bag(Y::Array{Float64,2}, params)

Computes the two factorizations Y = B0*A0^T and Y = [B0 B1]*A1^T.
Used instead of the train() function.
"""
function factorize_bag(Y::Array{Float64,2}, params; niter = 20)
    H = params.H
    H1 = params.H1

    # first get the A matrix for the decomposition Y = B0*A0'
    params0 = VBMatrixFactorization.vbmf_sparse_init(Y, H-H1);
    # copy the parameters that wont change
    params0.BHat = params.BHat[:,1:(H-H1)];
    params0.SigmaB = params.SigmaB[1:(H-H1),1:(H-H1)];
    params0.CB = params.CB[1:(H-H1)]
    params0.gamma = params.gamma
    params0.delta = params.delta[1:(H-H1)]
    if params0.M*(H-H1) < 1600
        full_cov = true
    else
        full_cov = false
    end
    A0 = vbls!(Y, params0, niter, full_cov = full_cov)

    # now do the same with full B matrix for the decomposition Y = B0*A10' + B1*A11'
    params1 = copy_vbmf_params(Y, params)
    A1 = vbls!(Y, params1, niter, full_cov = full_cov)

    return params0, params1
end

"""
    factorization_error(Y::Array{Float64,2}, params0, params1)

What are the errors produced by factorizations from factorize_bag()?
"""
function factorization_error(Y::Array{Float64,2}, params0, params1)
    H0 = params0.H

    #extract B0 and B1
    B0 = params0.BHat
    B1 =  params1.BHat[:,(H0+1):end]
    # also extract the submatrices A10 and A11
    A10 = params1.AHat[:, 1:H0];
    A11 = params1.AHat[:, (H0+1):end];

    # now compute the deltas
    delta0 = norm(Y - params0.YHat)
    delta1 = norm(Y - params1.YHat)
    delta10 = norm(Y - B0*A10')
    delta11 = norm(Y - B1*A11')

    return delta0, delta1, delta10, delta11
end

#####################

"""
    classify_one(res0, res1, Y; class_alg = "ols")

Using training data res0 and res1, classifies the specimen Y.
Argument class_alg specifies the way in which the A matrix
for a the new specimen is computed. Either "ols", "rls" or "vbls"
is used.
"""
function classify(res0, res1, Y::Array{Float64, 2}; threshold = 1e-1, class_alg::String = "ols")
    if class_alg in ["ols", "rls", "vbls"]
        # compute the ols estimate of YHat and choose the label
        # depending on the distance to the real Y matrix.
        if class_alg == "ols"
            B0 = res0.BHat
            AT0 =  ols(Y, B0)
            B1 = res1.BHat
            AT1 =  ols(Y, B1)
        elseif class_alg == "rls"
            # rls is used mainly for stabilization of the inversion
            # therefore lambda does not have to be very large
            B0 = res0.BHat
            AT0 =  rls(Y, B0, 1e-2)
            B1 = res1.BHat
            AT1 =  rls(Y, B1, 1e-2)
        elseif class_alg == "vbls"
            # this should be more optimal due to using the estimated
            # covariances
            # init a new instance of params
            vbparams = copy_vbmf_params(Y, res0)
            vbls!(Y, vbparams, 150)
            AT0 = vbparams.AHat'

            vbparams = copy_vbmf_params(Y, res1)
            vbls!(Y, vbparams, 150)
            AT1 = vbparams.AHat'
        end

        # compute the errors
        err0 = norm(Y - res0.BHat*AT0)
        err1 = norm(Y - res1.BHat*AT1)

        # produce label based on the smaller error
        if err0 > err1
            label = 1
        else
            label = 0
        end
    else # use a different approach
        if class_alg == "min_err"
            params0, params1 = factorize_bag(Y, res0)

            err0, err1, err10, err11 = factorization_error(Y, params0, params1)
            if abs((err0-err1)/err0) < threshold
                label = 0
            else
                label = 1
            end
        elseif class_alg == "lower_bound"
            params0, params1 = factorize_bag(Y, res0)
            L0 = VBMatrixFactorization.lowerBound(Y, params0)
            L1 = VBMatrixFactorization.lowerBoundTrimmed(Y, params1, threshold)

            if L1 > L0
                label = 1
            else
                label = 0
            end

            err0 = L0
            err1 = L1
        elseif class_alg == "dual"
            niter = 20
            L, M = size(Y)
            params0 = copy_vbmf_params(Y, res0);
            params1 = copy_vbmf_params(Y, res1);
            vbls!(Y, params0, niter, full_cov = true);
            vbls!(Y, params1, niter, full_cov = true);

            err0 = norm(params0.YHat - Y)/(L*M)
            err1 = norm(params1.YHat - Y)/(L*M)

            if err0 < err1
                label = 0
            else
                label = 1
            end
        end
    end

    return label, err0, err1
end

"""
    test_one(res0, res1, Y, bag_id, data)

Using training data res0 and res1, test the classification of Y.
Returns one of the set {-1,0,1} = {false positive, match, false negative}.
"""
function test_one(res0, res1, bag_id::Int, data::Dict{String,Any}; class_alg::String = "ols", threshold = 1e-1)
    Y = getY(data, bag_id)

    label = getLabel(data, bag_id)

    est_label, err0, err1 = classify(res0, res1, Y, class_alg = class_alg, threshold = threshold)
    return label - est_label
end

"""
    test(res0, res1, data, bag_ids)

For given bag_ids, it tests them all against a traning dataset. Returns
mean error rate, equal error rate and false positives and negatives count.
"""
function test_classification(res0, res1, data::Dict{String,Any}, bag_ids; class_alg::String = "ols", threshold = 1e-1)
    n = size(bag_ids)[1]
    n0 = 0 # number of negative/positive bags tested
    n1 = 0

    fp = 0 # number of false positives
    fn = 0 # number of false negatives
    for id in bag_ids
        res = test_one(res0, res1, id, data, class_alg = class_alg, threshold = threshold)

        if res == 1
            fn += 1
        elseif res == -1
            fp += 1
        end

        # also we count number of positive and negative try bags
        label = getLabel(data, id)
        if label == 0
            n0 +=1
        else
            n1 += 1
        end
    end

    mer = (fp+fn)/n
    eer = (fp/n0 + fn/n1)/2

    return mer, eer, fp, fn, n0, n1
end

"""
    validate(p_known, data, niter, solver, H; eps = 1e-6)

For a dataset and a percentage of known labels, asses the classification.
"""
function validate(p_known::Float64, data::Dict{String,Any}, niter::Int, solver::String, H::Int; H1::Int = 1,
 eps::Float64 = 1e-6, threshold = 1e-1,
    verb::Bool = true, diag_var::Bool = false, class_alg::String = "ols")
    nBags = data["bagids"][end]

    rand_inds = sample(1:nBags, nBags, replace = false);
    train_inds = rand_inds[1:Int(floor(p_known*nBags))];
    test_inds = rand_inds[Int(floor(p_known*nBags))+1:end];

    # training
    if class_alg in ["ols", "rls", "vbls"]
        res0, res1 = train(data, train_inds, solver, H, niter, eps = eps, verb = verb, diag_var = diag_var);
    elseif class_alg == "dual"
        res0, res1 = train_dual(data, train_inds, H, H-H1, niter, eps = eps, verb = verb, diag_var = diag_var);
    else # use the better approach
        res0 = train_local(data, train_inds, H, H1, niter, eps = eps, verb = verb, diag_var = diag_var)
        res1 = 0 # this does not actually make a lot of sense
    end
    if res0 == 0
        return -1.0, -1.0, -1.0, -1.0, -1.0, -1.0
    end

    # validation
    mer, eer, fp, fn, n0, n1 = test_classification(res0, res1, data, test_inds, class_alg = class_alg, threshold = threshold)

    return mer, eer, fp, fn, n0, n1
end

"""
    validate_with_cvs(data::Dict{String,Any}, which::Int, niter::Int, solver::String, H::Int; eps::Float64 = 1e-6,

For a dataset and a cv_index array index "which", asses the classification.
"""
function validate_with_cvs(data::Dict{String,Any}, test_inds, train_inds, niter::Int, solver::String, H::Int; H1::Int = 1,
 eps::Float64 = 1e-6, threshold = 1e-1,
    verb::Bool = true, diag_var::Bool = false, class_alg::String = "ols")
    # training
    if class_alg in ["ols", "rls", "vbls"]
        res0, res1 = train(data, train_inds, solver, H, niter, eps = eps, verb = verb, diag_var = diag_var);
    elseif class_alg == "dual"
        res0, res1 = train_dual(data, train_inds, H, H-H1, niter, eps = eps, verb = verb, diag_var = diag_var);
    else # use the better approach
        res0 = train_local(data, train_inds, H, H1, niter, eps = eps, verb = verb, diag_var = diag_var)
        res1 = 0 # this does not actually make a lot of sense
    end
    if res0 == 0
        return -1.0, -1.0, -1.0, -1.0, -1.0, -1.0
    end

    # validation
    mer, eer, fp, fn, n0, n1 = test_classification(res0, res1, data, test_inds, class_alg = class_alg, threshold = threshold)

    return mer, eer, fp, fn, n0, n1
end

"""
    validate_dataset(data::Dict{String,Any}, inputs::Dict{Any, Any}; verb::Bool = true)

Validates classification using vbmf on a whole MIL dataset using vector of
percentages of known labels. Inputs contain the vector of  percentages of known
labels, number of iterations over percentages and in vbmf algorithm etc.
Example of input:

inputs = Dict()
inputs["p_vec"] =  [0.01, 0.02, 0.05, 0.1, 0.33, 0.5, 0.75, 0.9]
inputs["nclass_iter"] = 1 # how many times is iterated over one p
inputs["niter"] = 200 # vbmf algorithm iterations
inputs["eps"] = 1e-3 # convergence criterion for vbmf
inputs["solver"] = "basic" # basic/sparse vbmf
inputs["H"] = 1 # inner dimension of factorization
inputs["scale_y"] = true # should Y be scaled to standard distribution?
inputs["use_cvs"] = true # should cv_indexes be used instead of p_vec?
inputs["diag_var"] = true # should homo- or heteroscedastic noise model be used?
inputs["class_alg"] = "ols" #/"rls"/"vbls"/"lower bound"/"dual" - how should the classification be computed
"""
function validate_dataset(data::Dict{String,Any}, inputs::Dict{Any, Any}, fname; verb::Bool = true)
    p_vec = inputs["p_vec"]
    nclass_iter = inputs["nclass_iter"]
    np = size(p_vec)[1]
    full_cvs = inputs["full_cvs"]

    res_mat = Array{Any,2}(nclass_iter*np, 7) # matrix of resulting error numbers

    # always compute the cv indexes partitioning first
    if inputs["use_cvs"]
        println("using cv_indexes")
        nrows, nfolds  = size(data["cvindexes"])
        ncvs = nrows*nfolds
        cv_res_mat = Array{Any,2}(ncvs, 7) # matrix of resulting error numbers for cv_indexes
        cv_res_mat[:, 1] = "cvs"
        cv_res_mat[:, 2:end] = -1.0

        n = 1
        for row in 1:1
            println("row = $(row)")
            print("fold = ")

            cv_indices = data["cvindexes"][row,:];
            # from the cv indices, choose one of the folds as validation and the rest as training
            for fold in 1:nfolds
                print("$fold ")
                # select test indices
                if full_cvs
                    test_inds = cv_indices[fold]
                    train_inds = Array{Int64,1}()
                    # create the train indices array
                    for j in 1:nfolds
                        if j!=fold
                            train_inds = cat(1, train_inds, cv_indices[j])
                        end
                    end
                else # for training, use just one fold and test on the rest
                    train_inds = cv_indices[fold]
                    test_inds = Array{Int64,1}()
                    for j in 1:nfolds
                        if j!=fold
                            test_inds = cat(1, test_inds, cv_indices[j])
                        end
                    end
                end

                try
                    mer, eer, fp, fn, n0, n1 = validate_with_cvs(data, test_inds, train_inds, inputs["niter"],
                        inputs["solver"],
                        inputs["H"], H1 = inputs["H1"], eps = inputs["eps"], verb = verb, diag_var = inputs["diag_var"],
                        class_alg = inputs["class_alg"], threshold = inputs["threshold"])
                    cv_res_mat[n,2:end] = [mer, eer, fp, fn, n0, n1]
                catch y
                    warn("Something went wrong during vbmf, no output produced.")
                    println(y)
                    cv_res_mat[n,2:end] = [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0]
                end

                # save the line to a temp file
                if isfile(fname)
                    contents = load(fname)
                    mat = contents["res_mat"]
                    mat = cat(1, mat, cv_res_mat[n,:]')
                    save(fname, "res_mat", mat, "inputs", inputs)
                else
                    save(fname, "res_mat", cv_res_mat[n,:]', "inputs", inputs)
                end

                n += 1
            end
            println("")
        end
    end
    println("")

    # then the partitioning given by a percentage of known bags p
    for ip in 1:np
        p = p_vec[ip]
        println("p = $(p) ")
        print("n = ")

        for n in 1:nclass_iter
            print("$(n) ")
            res_mat[(ip-1)*nclass_iter+n,1] = p
            try
                mer, eer, fp, fn, n0, n1 = validate(p, data, inputs["niter"], inputs["solver"],
                    inputs["H"], H1 = inputs["H1"],
                     eps = inputs["eps"], verb = verb, diag_var = inputs["diag_var"],
                    class_alg = inputs["class_alg"], threshold = inputs["threshold"])
                res_mat[(ip-1)*nclass_iter+n,2:end] = [mer, eer, fp, fn, n0, n1]
            catch y
                warn("Something went wrong during vbmf, no output produced.")
                println(y)
                res_mat[(ip-1)*nclass_iter+n,2:end] = [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0]
            end

            # save the line to a temp file
            if isfile(fname)
                contents = load(fname)
                mat = contents["res_mat"]
                mat = cat(1, mat, res_mat[(ip-1)*nclass_iter+n,:]')
                save(fname, "res_mat", mat, "inputs", inputs)
            else
                save(fname, "res_mat", res_mat[(ip-1)*nclass_iter+n,:]', "inputs", inputs)
            end
        end
    end
    println("")

    if inputs["use_cvs"]
        res = cat(1, cv_res_mat, res_mat)
    else
        res = res_mat
    end

    rm(fname)

    return res
end

"""
    validate_datasets(inputs::Dict{Any, Any}, file_inds::UnitRange{Int64}, input_path::String, output_path::String)

Wrapper for validate_dataset() that takes a whole folder of inputs, some settings and computes data for evaluation of the vbmf
classification.
"""
function validate_datasets(inputs::Dict{Any, Any}, file_inds, input_path::String, output_path::String;
    verb::Bool = true)

    files = readdir(input_path)
    #if verb
    #    println("The directory $input_path contains the following files:")
    #    for file in files
    #        println(file)
    #    end
    #    println("")
    #end

    mkpath(output_path)

    # loop through all the files, train them using vbmf, then validate the classification using a testing dataset
    # then save the results

    for file in files[file_inds]
        dataset_name, suf = split(file, ".")
        if suf != "jld" # if the file is not a .jld file, move on
            continue
        end

        if verb
            println("Processing file $file...")
        end

        # load the data
        data = load(string(mil_path, "/", file));

        # decide whether to ignore some rows
        # now replace the original Y matrix with a new one
        if inputs["scale_y"]
            data["fMat"] = VBMatrixFactorization.preprocess(data["fMat"], inputs["scale"], verb = true)
        end

        # perform testing of the classification on the dataset
        inputs["dataset_name"] = dataset_name
        fname = string(dataset_name, "_", inputs["solver"], "_", inputs["H"], "_", inputs["nclass_iter"])

        res_mat = validate_dataset(data, inputs, "$output_path/$fname.temp.jld", verb = false)

        # save the outputs and inputs
        save("$output_path/$fname.jld", "res_mat", res_mat, "inputs", inputs)

        if verb
            println("Done.")
            println()
        end
    end
end

"""
    warmup()

Serves to precompile the validate_datasets() function.
"""
function warmup(mil_path::String)
    print("Precompiling the validation function... ")
    inputs = Dict()
    inputs["p_vec"] =  [0.01]
    inputs["nclass_iter"] = 1 # how many times should be bags randomly assigned and classification tested over one percentage of known labels
    inputs["niter"] = 1 # iterations for vbmf solver
    inputs["eps"] = 1e-3 # the convergence limit for vbmf
    inputs["solver"] = "sparse" # basic/sparse for non/full ARD on A matrix in vbmf
    inputs["H"] = 1 # inner dimension of the factorization
    inputs["dataset_name"] = ""
    inputs["scale_y"] = true
    inputs["scale"] = 10.0 # how are they scaled? - could be better for numerical optimization
    inputs["use_cvs"] = false
    inputs["diag_var"] = false
    inputs["class_alg"] = "ols"
    inputs["H1"] = 1
    inputs["threshold"] = 1e-1
    inputs["full_cvs"] = false

    output_path = "./warmup_garbage"
    file_inds = 1:1

    validate_datasets(inputs, file_inds, mil_path, output_path, verb = false)

    rm("output_path", force=true, recursive=true)
    print("done. \n")
end

"""
    table_summary(res_mat::Array{Float64,2}; verb::Bool = true)

Returns and possibly prints mean values of error rates for a result of the validate_datasets() function.
"""
function table_summary(class_res::Dict{String,Any}; verb::Bool = true)
    inputs = class_res["inputs"]
    res_mat = class_res["res_mat"]

    p_vec =inputs["p_vec"]
    np = size(p_vec)[1]
    ndiag = size(res_mat)[2]
    ncvs = 0

    mean_table = Array{Any,2}(np, ndiag) # matrix of resulting error numbers

    # to be able to work with old results
    use_cvs = get(inputs, "use_cvs", false)

    # first the cv indexes statistics
    if use_cvs
        cvs_mean_table = Array{Any,2}(1, ndiag) # matrix of resulting error numbers
        cv_res_mat = res_mat[res_mat[:,1] .== "cvs", :] # cv result matrix
        cvs_mean_table[1,1] = "cvs"

        cv_res_mat = cv_res_mat[cv_res_mat[:,2] .!= -1.0, :] # throw away lines with computation errors
        if size(cv_res_mat)[1] == 0
            cvs_mean_table[1,2:end] = repmat([-1.0], 1, ndiag-1)
        else
            for i in 2:ndiag
                cvs_mean_table[1,i] = mean(cv_res_mat[!isnan(convert(Array{Float64,1},cv_res_mat[:,i])),i])  # throw away nans
            end
        end
    end

    for n in 1:np
        p = p_vec[n]
        p_mat = res_mat[res_mat[:,1] .== p, :] # extract just rows with current p-val
        p_mat = p_mat[p_mat[:,2] .!= -1.0, :] # throw away lines with computation errors
        if size(p_mat)[1] == 0
            p_mat = repmat([-1.0], 1, ndiag)
            p_mat[1] = p
        end
        for i in 1:ndiag
            mean_table[n,i] = mean(p_mat[!isnan(convert(Array{Float64,1},p_mat[:,i])),i])  # throw away nans
        end
    end

    heteroscedastic = get(inputs, "diag_var", false)
    if heteroscedastic
        noise_model = "heteroscedastic"
    else
        noise_model = "homoscedastic"
    end

    if verb
        dataset_name = inputs["dataset_name"]
        H = inputs["H"]
        H1 = inputs["H1"]
        nclass_iter = inputs["nclass_iter"]
        method = inputs["solver"]
        print("\nMean classsification error, $method solver, dataset $dataset_name, H = $H, H0 = $(H-H1), $nclass_iter samples, $noise_model noise: \n \n")
        print(" perc. of known labels | error rate | EER | false pos. | false neg. | neg. samples | pos. samples \n")
        print("------------------------------------------------------------------------------------------------------\n")
        if use_cvs
            @printf "        cvs                 %0.3f    %0.3f     %0.1f       %0.1f          %0.1f         %0.1f \n" cvs_mean_table[1,2] cvs_mean_table[1,3] cvs_mean_table[1,4] cvs_mean_table[1,5] cvs_mean_table[1,6] cvs_mean_table[1,7]
        end
        for n in 1:np
            @printf "        %0.2f                %0.3f    %0.3f     %0.1f       %0.1f          %0.1f         %0.1f \n" mean_table[n,1] mean_table[n,2] mean_table[n,3] mean_table[n,4] mean_table[n,5] mean_table[n,6] mean_table[n,7]
        end
    end

    if use_cvs
        res = cat(1, cvs_mean_table, mean_table)
    else
        res = mean_table
    end

    return res
end

"""
   plot_statistics(clas_res::Dict{String,Any})

Plots statistics of a validate_datasets() function result.
"""
function plot_statistics(class_res::Dict{String,Any}; verb::Bool = false, save_path::String = "")
    inputs = class_res["inputs"]
    res_mat = class_res["res_mat"]

    p_vec =inputs["p_vec"]
    np = size(p_vec)[1]
    ndiag = size(res_mat)[2]
    stat_names = ["mean error", "equal error rate", "false positives", "false negatives", "negative samples", "positive samples"]

    dataset_name = inputs["dataset_name"]
    H = inputs["H"]
    nclass_iter = inputs["nclass_iter"]
    method = inputs["solver"]
    mean_table = table_summary(class_res, verb = verb)

    # to be able to work with old results
    use_cvs = get(inputs, "use_cvs", false)

    #noise model
    heteroscedastic = get(inputs, "diag_var", false)
    if heteroscedastic
        noise_model = "heteroscedastic"
    else
        noise_model = "homoscedastic"
    end

    # because of cv indexes results are in the same file
    if use_cvs
        cv_mean_table = mean_table[1, :]
        mean_table = mean_table[2:end, :]

        if size(mean_table)[1] == 0
            # nothing to draw
            return
        end
    end
    mean_table = convert(Array{Float64,2}, mean_table)

    # plots
    ioff() # Interactive plotting OFF, necessary for inline plotting in IJulia
    fig = figure("vbmfa classification statistics",figsize=(8,13))
    #suptitle("$dataset_name, $method solver, H = $H, $nclass_iter samples")
    subplots_adjust(hspace=0.5)

    # mean values of error rates
    subplot(411) # Create the 1st axis of a 3x1 array of axes
    #ax = gca()
    #ax[:set_yscale]("log") # Set the y axis to a logarithmic scale
    plot(1:np, mean_table[:,2], label = stat_names[1])
    plot(1:np, mean_table[:,3], label = stat_names[2])
    title("$dataset_name, $method solver, H = $H, $nclass_iter samples, $noise_model noise: \n
        Mean error values")
    xlabel("percentage of known labels")
    ylabel("")
    xticks(1:np, p_vec)
    legend()

    # false negatives and positives
    subplot(412)
    plot(1:np, mean_table[:,4], label = stat_names[3])
    plot(1:np, mean_table[:,5], label = stat_names[4])
    plot(1:np, mean_table[:,6], label = stat_names[5])
    plot(1:np, mean_table[:,7], label = stat_names[6])
    xlabel("percentage of known labels")
    ylabel("")
    title("Identification statistics")
    xticks(1:np, p_vec)
    legend()

    # boxplots
    subplot(413) # Create the 2nd axis of a 3x1 array of axes
    data = []
    for n in 1:np
        p = p_vec[n]
        curr_vec = res_mat[res_mat[:,1] .== p, 2] # extract just rows with current p-val
        curr_vec = curr_vec[curr_vec .!= -1.0] # throw away lines with computation errors
        curr_vec = curr_vec[!isnan(convert(Array{Float64,1},curr_vec))]
        push!(data, curr_vec)
    end
    boxplot(data)
    title("box plot of mean error rate")
    xlabel("percentage of known labels")
    xticks(1:np, p_vec)

    subplot(414) # Create the 2nd axis of a 3x1 arrax of axes
    data = []
    for n in 1:np
        p = p_vec[n]
        curr_vec = res_mat[res_mat[:,1] .== p, 3] # extract just rows with current p-val
        curr_vec = curr_vec[curr_vec .!= -1.0] # throw away lines with computation errors
        curr_vec = curr_vec[!isnan(convert(Array{Float64,1},curr_vec))]
        push!(data, curr_vec)
    end
    boxplot(data)
    title("box plot of equal error rate")
    xlabel("percentage of known labels")
    xticks(1:np, p_vec)

     fig[:canvas][:draw]() # Update the figure
     gcf() # Needed for IJulia to plot inline

     # save the figure
     if save_path != ""
        filename = string(save_path, "/$dataset_name", "_$method", "_$H", "_$nclass_iter.eps")
        savefig(filename, format="eps", dpi=1000);
        println("Saving the figure to $filename.")
     end
     close()
end


text:
for Typ in (:Ldiv, :Rdiv)
    @eval begin
        struct $Typ{StyleA, StyleB, AType, BType}
            A::AType
            B::BType
        end

        $Typ{StyleA, StyleB}(A::AType, B::BType) where {StyleA,StyleB,AType,BType} =
            $Typ{StyleA,StyleB,AType,BType}(A,B)

        $Typ(A::AType, B::BType) where {AType,BType} =
            $Typ{typeof(MemoryLayout(AType)),typeof(MemoryLayout(BType)),AType,BType}(A, B)

        BroadcastStyle(::Type{<:$Typ}) = ApplyBroadcastStyle()
        broadcastable(M::$Typ) = M

        similar(A::$Typ, ::Type{T}, axes) where T = similar(Array{T}, axes)
        similar(A::$Typ, ::Type{T}) where T = similar(A, T, axes(A))
        similar(A::$Typ) = similar(A, eltype(A))

        copy(M::$Typ) = copyto!(similar(M), M)
        materialize(M::$Typ) = copy(instantiate(M))
    end
end

size(L::Ldiv{<:Any,<:Any,<:Any,<:AbstractMatrix}) = (size(L.A, 2),size(L.B,2))
size(L::Ldiv{<:Any,<:Any,<:Any,<:AbstractVector}) = (size(L.A, 2),)
axes(L::Ldiv{<:Any,<:Any,<:Any,<:AbstractMatrix}) = (axes(L.A, 2),axes(L.B,2))
axes(L::Ldiv{<:Any,<:Any,<:Any,<:AbstractVector}) = (axes(L.A, 2),)
length(L::Ldiv{<:Any,<:Any,<:Any,<:AbstractVector}) =size(L.A, 2)

size(L::Rdiv) = (size(L.A, 1),size(L.B,1))
axes(L::Rdiv) = (axes(L.A, 1),axes(L.B,1))

_ldivaxes(::Tuple{}, ::Tuple{}) = ()
_ldivaxes(::Tuple{}, Bax::Tuple) = Bax
_ldivaxes(::Tuple{<:Any}, ::Tuple{<:Any}) = ()
_ldivaxes(::Tuple{<:Any}, Bax::Tuple{<:Any,<:Any}) = (OneTo(1),last(Bax))
_ldivaxes(Aax::Tuple{<:Any,<:Any}, ::Tuple{<:Any}) = (last(Aax),)
_ldivaxes(Aax::Tuple{<:Any,<:Any}, Bax::Tuple{<:Any,<:Any}) = (last(Aax),last(Bax))

@inline ldivaxes(A, B) = _ldivaxes(axes(A), axes(B))

ndims(L::Ldiv) = ndims(last(L.args))
eltype(M::Ldiv) = promote_type(Base.promote_op(inv, eltype(M.A)), eltype(M.B))
eltype(M::Rdiv) = promote_type(eltype(M.A), Base.promote_op(inv, eltype(M.B)))

check_ldiv_axes(A, B) =
    axes(A,1) == axes(B,1) || throw(DimensionMismatch("First axis of A, $(axes(A,1)), and first axis of B, $(axes(B,1)) must match"))

check_rdiv_axes(A, B) =
    axes(A,2) == axes(B,2) || throw(DimensionMismatch("Second axis of A, $(axes(A,2)), and second axis of B, $(axes(B,2)) must match"))

function instantiate(L::Ldiv)
    check_ldiv_axes(L.A, L.B)
    Ldiv(instantiate(L.A), instantiate(L.B))
end

function instantiate(L::Rdiv)
    check_rdiv_axes(L.A, L.B)
    Rdiv(instantiate(L.A), instantiate(L.B))
end

_ldiv!(A, B) = ldiv!(factorize(A), B)
_ldiv!(A::Factorization, B) = ldiv!(A, B)

_ldiv!(dest, A, B) = ldiv!(dest, factorize(A), B)
_ldiv!(dest, A::Factorization, B) = ldiv!(dest, A, B)
_ldiv!(dest, A::Transpose{<:Any,<:Factorization}, B) = ldiv!(dest, A, B)
_ldiv!(dest, A::Adjoint{<:Any,<:Factorization}, B) = ldiv!(dest, A, B)

materialize!(M::Ldiv) = _ldiv!(M.A, M.B)
materialize!(M::Rdiv) = materialize!(Lmul(M.B', M.A'))'
copyto!(dest::AbstractArray, M::Rdiv) = copyto!(dest', Ldiv(M.B', M.A'))'

if VERSION ≥ v"1.1-pre"
    copyto!(dest::AbstractArray, M::Ldiv) = _ldiv!(dest, M.A, M.B)
else
    copyto!(dest::AbstractArray, M::Ldiv) = _ldiv!(dest, M.A, copy(M.B))
end

const MatLdivVec{styleA, styleB, T, V} = Ldiv{styleA, styleB, <:AbstractMatrix{T}, <:AbstractVector{V}}
const MatLdivMat{styleA, styleB, T, V} = Ldiv{styleA, styleB, <:AbstractMatrix{T}, <:AbstractMatrix{V}}
const BlasMatLdivVec{styleA, styleB, T<:BlasFloat} = MatLdivVec{styleA, styleB, T, T}
const BlasMatLdivMat{styleA, styleB, T<:BlasFloat} = MatLdivMat{styleA, styleB, T, T}

const MatRdivMat{styleA, styleB, T, V} = Rdiv{styleA, styleB, <:AbstractMatrix{T}, <:AbstractMatrix{V}}
const BlasMatRdivMat{styleA, styleB, T<:BlasFloat} = MatRdivMat{styleA, styleB, T, T}

# function materialize!(L::BlasMatLdivVec{<:AbstractColumnMajor,<:AbstractColumnMajor})

# end

macro lazyldiv(Typ)
    esc(quote
        LinearAlgebra.ldiv!(A::$Typ, x::AbstractVector) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
        LinearAlgebra.ldiv!(A::$Typ, x::AbstractMatrix) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
        LinearAlgebra.ldiv!(A::$Typ, x::StridedVector) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))
        LinearAlgebra.ldiv!(A::$Typ, x::StridedMatrix) = ArrayLayouts.materialize!(ArrayLayouts.Ldiv(A,x))

        Base.:\(A::$Typ, x::AbstractVector) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(A,x))
        Base.:\(A::$Typ, x::AbstractMatrix) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(A,x))

        Base.:\(x::AbstractMatrix, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(x,A))
        Base.:\(x::Diagonal, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(x,A))

        Base.:\(x::$Typ, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Ldiv(x,A))

        Base.:/(A::$Typ, x::AbstractVector) = ArrayLayouts.materialize(ArrayLayouts.Rdiv(A,x))
        Base.:/(A::$Typ, x::AbstractMatrix) = ArrayLayouts.materialize(ArrayLayouts.Rdiv(A,x))

        Base.:/(x::AbstractMatrix, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Rdiv(x,A))
        Base.:/(x::Diagonal, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Rdiv(x,A))

        Base.:/(x::$Typ, A::$Typ) = ArrayLayouts.materialize(ArrayLayouts.Rdiv(x,A))
    end)
end

ext: jl


text:
export ImplicitSolverType

"""
# Description
    ImplicitSolverType(;
        solver_method = KenCarp4,
    )

This solver type constructs an ODE solver using a _fully implicit_
method.

# Arguments
- `solver_method` (Function): Function defining the implicit
    solver.
    Default: `KenCarp4`
"""
struct ImplicitSolverType <: AbstractSolverType
    # Function for an implicit method
    solver_method::Function

    function ImplicitSolverType(
        alg;
        solver_method = ODESolvers.DiffEqJLConstructor(alg),
    )

        return new(solver_method)
    end
end

"""
    getdtmodel(ode_solver::AbstractSolverType, bl)

A function which returns a model representing the dynamics
with the most restrictive time-stepping requirements.
"""
function getdtmodel(ode_solver::ImplicitSolverType, bl)
    # For implicit methods, the entire model itself
    # contributes to the total stability of the time-integrator
    return bl
end

"""
# Description
    function solversetup(
        ode_solver::ImplicitSolverType,
        dg,
        Q,
        dt,
        t0,
        diffusion_direction,
    )

Creates an implicit ODE solver.
"""
function solversetup(
    ode_solver::ImplicitSolverType,
    dg,
    Q,
    dt,
    t0,
    diffusion_direction,
)

    solver = ode_solver.solver_method(dg, Q; dt = dt, t0 = t0)

    return solver
end

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA


text:
### stat functions

@doc """ return value type must be real
""" ->
function isGenomeWideScreen{T}(row::Array{T,1})
    # Genome-wide screen is in the 16th column
    gws = row[16]
    if gws == "y"
        return 1
    elseif gws == "n"
        return 0
    else
        error("GenomeWideScreen column is neither y or n")
    end
end

function getGeneCDSLength(row::Array{ASCIIString,1})
    # getGeneCDSLength is in the 3th column
    parse(Int64, row[3])
end

### typical read pipeline
function readDiscreteLine(file::IOStream)
    map(x->convert(ASCIIString,x), split(strip(readline(file)),'\t'))
end

function statCosmic(funcs::Array{Function,1}, filepath)
    i = 0
    function ProgressInfo(numCounter)
        i += 1
        if i % numCounter == 0
            info("In the line $i")
        end
    end

    stats = zeros(Int64, length(funcs))

    open(filepath) do file
        header = readDiscreteLine(file)
        while !eof(file)
            row = readDiscreteLine(file)
            stats = stats + map(f->f(row), funcs)
            ProgressInfo(100_000)
        end
    end
    stats
end

function stat()
    statCosmic(Function[isGenomeWideScreen, getGeneCDSLength], "../data/CosmicCompleteExport.tsv")
end


text:
using Distributions
using Distances
using Random
using LinearAlgebra
using Zygote
using Zygote: @adjoint

include("./kernel.jl")
include("./transform.jl")
include("./covariate.jl")

abstract type AbstractParameter{T} end

struct Parameter{T}
    param::Vector{T}
    probs::Vector{Float64}
end

mutable struct ScalarParameter{T<:Real} <: AbstractParameter{T}
    param::T
    dist::Distribution{Univariate,Continuous}
end

struct VectorParameter{T<:Real} <: AbstractParameter{T}
    param::Vector{T}
    dist::Vector{Distribution{Univariate,Continuous}}
end

log_jacobian(g, λ, y) = sum(x -> log∂(g, λ, x), y)

function logprob(k::AbstractKernel, g::AbstractTransform, f::AbstractCovariate, U::AbstractMatrix, X::AbstractMatrix, yu::AbstractVector, yx::AbstractVector)
    Fu::Matrix{Float64} = f(U)
    W::Matrix{Float64} = Fu \ f(X)
    h = let W = W, k = k, U = U, X = X, g = g, yu = yu, yx = yx, Fu = Fu
        (θ, λ) -> logprob(k, θ, g, λ, U, Fu, X, W, yu, yx)
    end
    return h
end

Zygote.@adjoint function mul!(
    C::AbstractMatrix,
    A::AbstractMatrix,
    B::AbstractMatrix,
    α::Real,
    β::Real,
)
    tmp, back = pullback(C, A, B, α, β) do C, A, B, α, β
        (A * B) .* α .+ C .* β
    end
    return C, c::AbstractMatrix -> back(c)
end

function update_Kxx′!(
    Kxx′::AbstractMatrix,
    B::AbstractMatrix,
    Kux::AbstractMatrix,
    W::AbstractMatrix,
    W′::AbstractMatrix
)
    mul!(Kxx′, W', B, -1, 1)
    mul!(Kxx′, Kux', W′, -1, 1)
    return Kxx′
end

function update_Kx′!(
    Kx′::AbstractMatrix,
    B::AbstractMatrix,
    Kux′::AbstractMatrix,
    W′::AbstractMatrix
) # B = Kux′ .- Ku * W′
    mul!(Kx′, (W′)', B, -1, 1)
    mul!(Kx′, (Kux′), W′, -1, 1)
    return Kx′
end

struct Block
    Ku::Matrix{Float64}
    Kux::Matrix{Float64}
    rKx::Matrix{Float64}
end

function logprob(k, θ, g, λ, U, Fu, X, W, yu, yx)
    p, m = size(W)
    Ku = correlation(k, θ, U)
    Kux = correlation(k, θ, U, X)
    Kx = correlation(k, θ, X)
    B = Kux .- Ku * W
    rKx = Kx .- W' * B .- Kux' * W
    zu = (x -> g(λ, x)).(yu)
    zx = (x -> g(λ, x)).(yx)
    z = zx .- W' * zu
    c = rKx \ z
    lnJλ = log_jacobian(g, λ, yu) + log_jacobian(g, λ, yx)
    lnDet = 2 * logabsdet(Fu)[1] + logdet(rKx)
    lnQ = log(dot(z, c))
    return m / (m + p) * lnJλ - m / 2 * lnQ - 0.5 * lnDet
end

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
module AlgebraPDF
#
# using Base: NamedTuple
using Parameters
using QuadGK
#
using PyCall
using Optim
#
using Random
using RecipesBase
using StaticArrays
using Interpolations
using ForwardDiff
using LinearAlgebra
using Measurements

#
const ∅ = NamedTuple()
export ∅
#
const NumberOrTuple = Union{Number,Tuple{Vararg{Number}}}
export NumberOrTuple

import Base: +,-,*,==
import Base: getproperty
import Base: getindex
import Base: abs2, log
import Base: keys
import Base: length
import Optim: minimizer, minimum
import ForwardDiff: hessian
#
#
export ParTypes
include("parameters.jl")
const Ext = FlaggedNamedTuple
export Ext, FlaggedNamedTuple

#
export updatepar, updatepars
export fixpar, fixpars
export releasepar, releasepars
#
export @makefuntype
export AbstractFunctionWithParameters
export FunctionWithParameters
#
export func
export nfreepars
export freepars, fixedpars
export pars
export p2v, v2p
#
include("functionwithparameters.jl")

#
export @makepdftype
export AbstractPDF
export Normalized
export integral, integrals
export fixedshapepdf
export noparsf, noparsnormf
export lims
export normalizationintegral
#
include("pdf.jl")
#

export FSum, FAbs2, FLog
export NegativeLogLikelihood
export Extended
export FSum
include("arithmetics.jl")

const pdf = Normalized
export pdf

export fit
export BFGSApproxHesse
export MigradAndHesse
export isiminuitimported
export minussum
export obj2nt
include("fit.jl")
include("minuit_interface.jl")
include("optim_interface.jl")

export generate
include("generation.jl")

export convGauss
include("convolution.jl")

export scaletobinneddata
include("plottingrecipe.jl")

export FGauss, FBreitWigner, FExp
export FPowExp, FPol
export FDoubleGaussFixedRatio
export FBreitWignerConvGauss
export FTabulated
include("densities.jl")

export nt
export inrange
include("utils.jl")

end # module


text:
function build!(m::EnergyModel)
    # If subnetworks are not set, compute them
    isempty(m.subnetworks) && determine_subnetworks!(m)

    # Clear jump model
    empty!(m.jumpobjects)
    MOI.empty!(backend(m.jumpmodel))
    for vartype in [:fix, :integrality, :lower_bound, :upper_bound, :zero_one]
        empty!(getfield(m.jumpmodel, Symbol(:variable_to_, vartype)))
    end

    # Fill it up again
    addto!(m.jumpmodel, m)
end

function addto!(jm::JuMP.AbstractModel, m::EnergyModel)
    # @info("* Equations for individual devices")
    for d = devices(m)
        # @info("  - $(naming(d)) ($(naming(typeof(d))))")
        addto!(jm, m, d)
    end
    # @info("* Equations for energy balance")
    for d = buses(m)
        addto!(jm, m, d)
    end
    # @info("* Equations for subnetworks")
    for sn = subnetworks(m)
        addto!(jm, m, sn)
    end
    # @info("* Cost minimization objective")
    @objective(jm, Min, sum(cost(d) for d = devices(m)))

    jm
end

max_stars_repo_name: coroa/EnergyModels

max_issues_repo_name: coroa/EnergyModels

max_forks_repo_name: coroa/EnergyModels


text:
θ = rand() * -2π + π
a = U1(θ)
b = U1(rand())
c = U1(rand())

## Properties of the Group ##
@test isapprox(arg(a), θ)
@test isapprox(len(a), 1)

## Properties of the Unary Operators ##

@test a == +a
@test a == -(-a)

## Properties of Multiplication ##

@test a * (U1(π / 2) / a) == U1(π / 2) # Identity
@test a * b == b * a # Commutativity
@test (a * b) * c == a * (b * c) # Associativity


text:
using Base.Iterators, LinearAlgebra, Optim;
import DelimitedFiles.readdlm;

n = 4;
H = [1, 0];
V = [0, 1];
Dia = [1/sqrt(2), 1/sqrt(2)];
A = [1/sqrt(2), -1/sqrt(2)];
R = [1/sqrt(2), -im/sqrt(2)];
L = [1/sqrt(2), im/sqrt(2)];
statevars = [H, V, Dia, A, R, L];

HadamardMatrix_2 = [1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2)];
PauliMatrix_1 = [0 1; 1 0];

Trafo = kron(HadamardMatrix_2 .* Matrix{Float64}(I, 2,2), PauliMatrix_1 * HadamardMatrix_2);
NDimensions = 4;

D2Array = readdlm("D:\\Binto\\ddd\\Julia\\Tomography_Mathematica_to_Julia\\measurement_bases_full_edit_measurement_7.txt", skipstart=2)[1:end,2:5];
Data = collect(flatten(D2Array));
SigmaData = sqrt.(Data);

function BadnessPolynomial!(t)
    pM = [t[1] 0 0 0;
        t[5] + im*t[6]  t[2] 0 0;
        t[11] + im*t[12] t[7] + im*t[8] t[3] 0;
        t[15] + im*t[16] t[13] + im*t[14] t[9] + im*t[10] t[4]];

    pMd = [t[1] t[5] + im*t[6] t[11] + im*t[12] t[15] + im*t[16];
        0 t[2] t[7] + im*t[8] t[13] + im*t[14];
        0 0 t[3] t[9] + im*t[10];
        0 0 0 t[4]];

    global GeneralDM = pMd * pM;
    Prediction(DM, State) = sum(conj(State) .* DM .* State);
    StateMeasured = [kron(i,j) for i in statevars, j in statevars];
    #global PredictionPart = [sum(conj(StateMeasured[i]) .* GeneralDM .* StateMeasured[i]) for i = 1:length(StateMeasured)];
    PredictionPart = [];
    for i = 1:length(StateMeasured)
        push!(PredictionPart, Prediction(GeneralDM, StateMeasured[i]));
    end
    return BD = real(sum((PredictionPart[i] - Data[i])^2 /PredictionPart[i] for i=1:length(StateMeasured)))
end
lower = fill(-5.0, 16);
upper = fill(5.0, 16);
t0 = fill(0.05, 16);
#t0 = randn(16)
#@time res = optimize(BadnessPolynomial!, t0, NelderMead(), Optim.Options(g_tol = 10e-10, iterations = 20000))
#@time res = optimize(BadnessPolynomial!, t0, LBFGS(), Optim.Options(iterations = 20000)) #Requires a function and gradient (will be approximated if omitted)
inner_optimizer = GradientDescent()
@time res = optimize(BadnessPolynomial!, lower, upper, t0, Fminbox(inner_optimizer), Optim.Options(iterations = 50000)) #Box constrained minimization
temp = Optim.minimizer(res);

MaxLikDM = GeneralDM/tr(GeneralDM)

n1 = 0;
n2 = 6;
n3 = 4;
n4 = 4;


text:
import BothOldAndNew
import ScientificTypes
import Test

Test.@testset "BothOldAndNew.jl" begin
    Test.@test BothOldAndNew.f(1) == 2
end