

################################################################################
### GENERAL-PURPOSE OBJECTIVE FUNCTIONS
################################################################################
"""
    likelihood(X, P, normalized=true)

Computes the likelihood of `X` given distribution `P`
"""
function likelihood(X, P::AbstractBinaryVectorDistribution, normalized=true)
    Px = [pdf(P, X[:,k]) for k = 1:size(X,2)]
    if normalized
        Px = Px .^ (1 / size(X,2))
    end
    return prod(Px)
end

"""
    loglikelihood(X, P, normalized=true)

Computes the log likelihood of `X` given distribution `P`. Returns `0.0` if `P`
assigns probability 0 to any column of `X`
"""
function loglikelihood(X, P::AbstractBinaryVectorDistribution, normalized=true)
    Px = [pdf(P, X[:,k]) for k = 1:size(X,2)]
    if any(Px .== 0.0)
        return 0.0
    end
    return sum_kbn(map(log,Px)) / (normalized ? size(X,2) : 1)
end

################################################################################
### ISING OBJECTIVE FUNCTIONS - NLopt.jl
################################################################################
"""
    loglikelihood(X, Jtilde, grad)

Computes the loglikelihood of the data given the Ising distribution generated by
`Jtilde`. If `grad` is a vector of length > 0, it is modified in-place with the
gradient of the function.

"""
function loglikelihood(X, Jtilde::Vector, grad::Vector=[]; mu_X=(X*X'/size(X,2)), kwargs...)
    (N_neurons, N_samples) = size(X)

    Jtilde = reshape(Jtilde, N_neurons,N_neurons)
    P = IsingDistribution(Jtilde)

    if length(grad) > 0
        # performing this computation first, if it's necessary, results in all
        # the pdf values being cached.
        # mu_X = X * X' / N_samples
        mu_P = expectation_matrix(P)
        grad[:] = -(mu_P - mu_X)[:]
        grad[1:(N_neurons+1):end] = diag(mu_P - mu_X)[:]

    end
    # return sum_kbn([log(pdf(P, X[:,k])) for k = 1:N_samples]) / N_samples
    return -log(_get_Z(P)) - (sum_kbn([_E_Ising(P, X[:,i]) for i = 1:N_samples])) / N_samples
end

"""
    MPF_objective(X, J, grad)

Uses function `K` from the MPF paper (or maybe from the MPF sample code for
Matlab). Using this to fit J will typically result in values close to optimal,
but the advantage is that this method does not involve computing the partition
function Z.

"""
function MPF_objective(X, Jtilde::Vector, grad::Vector=[])
    (N_neurons, N_samples) = size(X)
    Jtilde = reshape(Jtilde, N_neurons, N_neurons)
    theta = diag(Jtilde)
    J = Jtilde - Diagonal(Jtilde)
    DeltaX = 2 * X - 1 # this is Δx_l for each codeword x, for each index l
    # Kfull = exp.((-0.5 * DeltaX .* (J * X) + DeltaX .* theta)/2)
    Kfull = exp.((DeltaX .* theta - DeltaX .* (J * X)) / 2)
    K = sum_kbn(Kfull[:]) / N_samples
    if length(grad) > 0
        # M = zeros(N_neurons, N_neurons)
        # M[1:(N_neurons+1):end] = sum(0.5 * Kfull .* DeltaX, 2) / N_samples
        # for p = 1:(N_neurons - 1)
        #     for q = (p+1):N_neurons
        #         M[p,q] = M[q,p] = -0.5 * sum([(Kfull[p,w] * DeltaX[p,w] * X[q,w] + Kfull[q,w] * DeltaX[q,w] * X[p,w]) for w = 1:N_samples]) / N_samples
        #     end
        # end
        DK = Kfull .* DeltaX
        dJ = -0.5 * DK * X'
        dJ = (dJ + dJ') / 2
        dJ[1:(n+1):end] = 0.5 * sum(DK, 2)
        grad[:] = dJ[:] / N_samples
    end
    return K
end



################################################################################
### ISING OBJECTIVE FUNCTIONS - Optim.jl
################################################################################
"""
    negloglikelihood(X, Jtilde)

Computes negative LogLikelihood of the data given the Ising distribution generated by
`Jtilde`. See also `dnegloglikelihood!` for gradient. These functions are intended for use
with `Optim.jl` package

"""
function negloglikelihood(W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    N = size(buf[1], 2)
    return log(get_Z(buf[3])) - sum_kbn([F_Ising(buf[3], buf[1][:,α]) for α = 1:N]) / N
end
function dnegloglikelihood!(G, W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    #TODO compute gradient, modify G in place
    mu_P = expectation_matrix(buf[3])
    G[:] = lt2v(mu_P) - lt2v(buf[2])
end
function Hnegloglikelihood!(H, W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    # if we use lt2v on mu_P and take its outer product, we'll get the matrix of
    # E[x_s]E[x_t] that we subtract from E[x_s x_t] to get the Hessian matrix. So the only
    # thing we need to be careful about is making sure we line up terms correctly.
    N_neurons = n_bits(buf[3])
    p = get_pdf(buf[3])
    for k in sortperm(p)
        x = digits(Bool, k-1, 2, N_neurons)
        xx = lt2v(x * x')
        H[:] = H[:] + (p[k] * (xx * xx'))[:]
    end
    mm = lt2v(buf[4])
    H[:] = H[:] - (mm * mm')[:]
end
function negloglikelihood_common!(W, W_prev, buf)
    # buf[1] = X
    # buf[2] = mu_X
    # buf[3] = P
    # buf[4] = mu_P
    if W != W_prev
        copy!(W_prev, W)
        #TODO set the appropriate buffer entry to be P
        buf[3] = IsingDistribution(W)
        buf[4] = expectation_matrix(buf[3])
    end
end


"""
    K_MPF(J, J_prev, buf)

Computes the function ``K_X`` for data in `buf[1]` and parameters `J`. This version is meant
for use with the `Optim` package, which in particular uses a separate function for the
optimization function and gradient function. See also `dK_MPF!`

"""
function K_MPF(J, J_prev, buf)
    # N_neurons, N_samples = size(X)
    # J = reshape(Jtilde, N_neurons, N_neurons)
    # theta = diag(J)
    # J[1:(N_neurons+1):end] = 0.0
    # ΔX = 2X - 1 # flip bits in X
    # Kfull = exp.(0.5 * (ΔX .* theta - ΔX .* (J * X)))
    # return sum_kbn(Kfull) / N_samples
    K_common!(J, J_prev, buf)
    # N_samples = size(buf[1], 2)
    return sum_kbn(buf[3]) # / N_samples
end
"""
    dK_MPF!(G, J, J_prev, buf)

Computes the gradient of `K_MPF(J, J_prev, buf)` and modifies `G` in-place.
"""
function dK_MPF!(G, J, J_prev, buf)
    # N_neurons, N_samples = size(X)
    # J = reshape(Jtilde, N_neurons, N_neurons)
    # theta = diag(J)
    # J[1:(N_neurons+1):end] = 0.0
    # ΔX = 2X - 1 # flip bits in X
    # Kfull = exp.(0.5 * (ΔX .* theta - ΔX .* (J * X)))
    # # M = zeros(N_neurons, N_neurons)
    # # M[1:(N_neurons+1):N_neurons] = sum(0.5 * ΔX .* Kfull, 2)
    # DK = Kfull .* ΔX
    # dJ = -0.5 * (Kfull .* ΔX) * X'
    # dJ = (dJ + dJ') / 2
    # dJ[1:(N_neurons+1):end] = 0.5 * sum(DK, 2)
    # G[:] = dJ[:] / N_samples
    K_common!(J, J_prev, buf)
    N_neurons, N_samples = size(buf[1])
    DK = buf[3] .* buf[2]
    dJ = -0.5 * DK * buf[1]'
    dJ = (dJ + dJ') / 2
    dJ[1:(N_neurons+1):end] = 0.5 * sum(DK, 2)
    G[:] = dJ[:] # / N_samples
end
function K_common!(J, J_prev, buf)
    # buf[1] = X # fixed ahead of time
    # buf[2] = ΔX # fixed ahead of time
    # buf[3] = Kfull / N_samples
    # buf[4] = theta
    # buf[5] = J (no diagonal)
    if J != J_prev
        copy!(J_prev, J)

        N_neurons, N_samples = size(buf[1])
        J = reshape(J, N_neurons, N_neurons)
        theta = diag(J)
        J[1:(N_neurons+1):N_neurons] = 0.0
        buf[4] = theta
        buf[5] = J
        buf[3] = exp.(0.5 * (buf[2] .* theta - buf[2] .* (J * buf[1]))) / N_samples
    end
end
