# Optimization Code
## X: n x p design matrix (feature matrix)
## y: n x K outcome matrix
## rho: Sparsity level (integer)
## B: p*K initial solution
## K: number of tasks
## L: Lipschitz constant
## n: sample size
## eig: max eigenvalue of all design matrix
## p: number of features/covariates
## idx: vector of active sets
## lambda1>=0: the ridge coefficient
## lambda2>=0: the coefficient value strength sharing coefficient (Bbar penalty)
## lambda_z>=0: the coefficient support strength sharing coefficient (Zbar penalty)
## maxIter_in: number of inner coordinate descent iterations (within active sets)
## maxIter_out: number of outer coordinate descent iterations

#### Active set
# sparse regression with IHT
function BlockComIHT_inexactAS_opt_old_MT(; X::Array{Float64,2},
                    y::Array{Float64,2},
                    rho::Integer,
                    B::Array{Float64,2},
                    K::Integer,
                    L::Float64,
                    n::Integer,
                    lambda1::Float64,
                    lambda2::Float64,
                    lambda_z::Float64,
                    p::Integer,
                    maxIter_in::Integer = 1000,
                    maxIter_out::Integer = 1000
                    )::Array{Float64,2}

    eig = L # initialize for active set
    L = L + lambda1 + lambda2 # L for complete dataset for outer loop
    idxFlag = true # indicates whether indices have changed and whether we nee to recalculate eigenvalues

    ncol = p + 1

    # initialize

    z = zeros(p, K)
    B_bar = zeros(ncol)
    z_bar = zeros(p)
    Z_bar = zeros(p, K)

    r = zeros(n, K)
    g = zeros(ncol, K);

    obj = 1e20;
    iter_in = 1
    iter_out = 1

    # do not calculate for intercept -- initialize based on beta warm start
    B_summed = sum( B.^2, dims = 2)
    B_summed = B_summed[:]

    # projection step
    idx = partialsortperm(B_summed[2:end], # do not include intercept
                          1:rho,
                          by=abs,
                          rev=true)

    idxInt = cat(1, idx .+ 1, dims = 1) # add one for intercept
    z[idx] = ones(size(idx))
################################################################################

# outer while loop
    while (iter_out <= maxIter_out)

        objPrev = obj
        BPrev = copy(B)

        if idxFlag
            # if active set has changed (or first round) recalculate svd
            L_active = tsvd( X[ :, idxInt] )[2][1]; # 0

            L_active = L_active^2 * sqrt(K) / n + lambda1 + lambda2
            idxFlag = false # default to see if we need to recalculate eigenvalues

        end

        B_bar_active = zeros(length(idxInt))
        z_bar_active = zeros(length(idx))
        r_active = zeros(n, K)
        g_active = zeros(length(idxInt) ,K)
        B_active = B[idxInt,:]
        z_active = z[idx, :]

        # inner while loop
########################################################
        while (iter_in <= maxIter_in)

            objPrev = obj
            B_bar_active = sum(B_active[2:end, :], dims=2) / K
            z_bar_active = sum(z_active, dims=2) / K
            z_bar_active = z_bar_active[:]
            B_bar_active = B_bar_active[:]
            f = 0

            for k = 1:K
                r_active[:,k] = X[ :, idxInt] * B_active[:,k] - y[ :, k ]
                g_active[:, k] = (1 / n) * X[ :, idxInt]' * r_active[:,k] # gradient update for whole vector
                g_active[2:end, k] = g_active[2:end, k] + lambda1 * B_active[2:end, k] + lambda2 * (B_active[2:end, k] - B_bar_active) # gradient update for non-intercept
                f = f + r_active[:, k]' * r_active[:, k] / (2 * n) +
                        lambda1 * B_active[2:end, k]' * B_active[2:end, k] / 2 +
                        lambda2 * (B_active[2:end, k] - B_bar_active)' * (B_active[2:end, k] - B_bar_active) / 2 +
                        lambda_z * (z[idx, k] - z_bar_active)' * (z[idx,k] - z_bar_active) / 2
            end

            obj = f

            if (abs(obj - objPrev)/objPrev < 1e-6)
                break
            end

            iter_in = iter_in + 1
            Z_bar_active = z_bar_active * ones(K)'
            B_temp_active = B_active - 0.5 * (1/L) * g_active
            z_active = zeros(length(idx), K)
            idx_active = findall(x-> x.>1e-9, abs.(B_temp_active[2:end,:]) )
            z_active[idx_active] = ones(size(idx_active))
            cost_active = lambda_z / L_active * ( Z_bar_active.^2  - (z_active-Z_bar_active).^2 ) # part of cost vector, rest in for loop below
            z_active = zeros(length(idx), K)

            for k=1:K
                costAct_k = B_temp_active[2:end, k].^2 + n * cost_active[:,k] # scale by nVec since our objective is altered
                idx1 = partialsortperm(costAct_k, 1:rho, rev=true)
                z_active[idx1, k] = ones(size(idx1))
            end


            B_active[2:end,:] = B_temp_active[2:end,:] .* z_active
            B_active[1,:] = B_temp_active[1,:] # update intercept

            B[idxInt,:] = B_active[:,:]
            z[idx,:] = z_active[:,:]
        end
        # end inner while loop
########################################################

        B_bar = sum(B[2:end, :], dims=2) / K
        z_bar = sum(z, dims=2) / K
        z_bar = z_bar[:]
        B_bar = B_bar[:]
        f = 0

        for k = 1:K

            # residual for kth study
            r[:, k] = X * B[:,k] - y[ :, k ]

            # gradient for kth study
            g[:, k] = (1 / n) * X' * r[:,k] # gradient update for whole vector
            g[2:end, k] = g[2:end, k] + lambda1 * B[2:end, k] + # gradient update for non-intercept terms
                                        lambda2 * (B[2:end, k] - B_bar)

            # objective for kth study
            f = f + r[:,k]' * r[:,k] / (2 * n) +
                        lambda1 / 2 * B[2:end, k]' * B[2:end, k] +
                        lambda2 / 2 * (B[2:end, k] - B_bar)' * (B[2:end, k] - B_bar) +
                        lambda_z / 2 * (z[:,k] - z_bar)' * (z[:, k] - z_bar)
        end

        obj = f

        iter_out = iter_out + 1
        Z_bar = z_bar * ones(K)'
        B_temp = B - 0.5 * (1/L) * g

        z = zeros(p, K)
        idx1 = findall(x-> x.>1e-9, abs.(B_temp[2:end,:]) )
        z[idx1] = ones(size(idx1))
        cost =  lambda_z / L * (Z_bar.^2 - (z-Z_bar).^2) # scale by nVec (below) since our objective is altered
        z = zeros(p, K)
        flag = 0

        for k=1:K
            cost_k = B_temp[2:end, k].^2 + n * cost[:,k] # scale lambda_z by nVec
            idx1 = partialsortperm(cost_k,
                                  1:rho,
                                  rev=true)
            z[idx1, k] = ones(size(idx1))

            for jj = 1:length(idx1)
                if ( sum(idx .== idx1[jj]) == 0 )
                    idx = [ idx; idx1[jj] ]
                    idxInt = [ idxInt; idx1[jj] + 1] # for intercept version
                    flag = 1
                    idxFlag = true # set to recalculate eigenvalues
                end
            end

        end

        B[2:end,:] = B_temp[2:end,:] .* z
        B[1,:] = B_temp[1,:] # update intercept

        if (flag == 0)
           break
        end

    end
    # end outer while loop
######################################################################

    return B

end
