# Code Question 1
# Takes a residual matrix R and (element wise) calculates the huber loss
function huberify(R,e)
    HR = huberloss.(R,e)
end

function huberloss(r, e)
    if (abs(r) < e)
        return (1/2)r^2
    else
        return e*(abs(r) - (1/2)*e)
    end
end

function robustPCA(X,k)
    (n,d) = size(X)
    e = 0.01

    # Subtract mean
    mu = mean(X,dims=1)
    X -= repeat(mu,n,1)

    # Initialize W and Z
    W = randn(k,d)
    Z = randn(n,k)

    R = Z*W - X
    f = sum(huberify(R,e))
    funObjZ(z) = pcaObjZHuber(z,X,W,e)
    funObjW(w) = pcaObjWHuber(w,X,Z,e)
    for iter in 1:50
        fOld = f

        # Update Z
        Z[:] = findMin(funObjZ,Z[:],verbose=false,maxIter=10)

        # Update W
        W[:] = findMin(funObjW,W[:],verbose=false,maxIter=10)

        R = Z*W - X
        f = sum(R.^2)
        @printf("Iteration %d, loss = %f\n",iter,f/length(X))

        if (fOld - f)/length(X) < 1e-2
            break
        end
    end


    # We didn't enforce that W was orthogonal so we need to optimize to find Z
    compress(Xhat) = compress_gradientDescent(Xhat,W,mu)
    expand(Z) = expandFunc(Z,W,mu)

    return CompressModel(compress,expand,W)
end

function pcaObjZHuber(z,X,W,e)
    # Rezie vector of parameters into matrix
    n = size(X,1)
    k = size(W,1)
    Z = reshape(z,n,k)

    # Comptue function value
    R = Z*W - X
    HR = huberify(R,e)
    f = sum(HR)

    # Comptue derivative with respect to each residual
    dR = R

    # Multiply by W' to get elements of gradient
    G = HR*W'

    # Return function and gradient vector
    return (f,G[:])
end

function pcaObjWHuber(w,X,Z,e)
    # Rezie vector of parameters into matrix
    d = size(X,2)
    k = size(Z,2)
    W = reshape(w,k,d)

    # Comptue function value
    R = Z*W - X
    HR = huberify(R, e)
    f = sum(HR)

    # Comptue derivative with respect to each residual
    dR = R

    # Multiply by Z' to get elements of gradient
    G = Z'HR

    # Return function and gradient vector
    return (f,G[:])
end