#  C:\dev\GitHub\autodiff\GMM.jl
#  AUTOGENERATED FROM C:\dev\GitHub\autodiff\GMM.ipynb on 2015-08-03T13:16:36
using Base.Test
using IPynbToJl
ipynb_to_jl("GMM.ipynb")
ipynb_to_jl("logsumexp.ipynb")
ipynb_to_jl("LowerTriangular.ipynb")

@printf("julia version = %s, dir %s\n", Base.VERSION_STRING, pwd())
include("MatVec.jl")
include("LowerTriangular.jl")
# Conventional GMM.
# This doesn't even use inverse covariance, because as soon as you
# start down that route, you may as well go for lpGMM below.
type GMM
  n::Int           # number of Gaussians
  d::Int           # dimension of Gaussian
  alphas::Vec      # weights: n, require sum(alphas)==1
  mus::Array{Vec}  # means: n, each dx1
  sigmas::Array{SymMat}  # covariances: n, each dxd symmetric positive definite
end

function log_likelihood(g::GMM, x::Vec)
  total = 0
  for k=1:g.n
    mean = g.mus[k]
    weight =  g.alphas[k]
    Σ = g.sigmas[k]
    mahalanobis = dot(mean - x, inv(Σ) * (mean - x))
    total += weight / sqrt(det(2pi*full(Σ))) * exp(-0.5*mahalanobis)
  end
  log(total)
end
# Test it a bit

n=3
d=2
alphas=rand(n); alphas /= sum(alphas);
mus=[randn(d) for k=1:n]
sigmas=[AAt(randn(d,d)) for k=1:n]
test_gmm = GMM(n,d,alphas,mus,sigmas)
@printf("An example gmm = %s\n", test_gmm)

x = randn(d) # Test point

ll0 = log_likelihood(test_gmm, x)
@printf("Tes log likelihood ll0=%f\n", ll0)

#[markdown]
# ---------

# 

# Log-parameterized GMM

# =====================

# 

# This is a GMM parameterized by log-weights, and the Cholesky factor of the inverse covariance.  This means we can do unconstrained optimization, and in fact makes for a more efficient computation of the normalization constant, as well as more numerical stability all round.    

##########################################################################
# Log-parametrized GMM.
# Weights are strictly positive, covariances are parameterized by their inverse
# square roots (lower triangular).
type lpGMM
  n::Int           # number of Gaussians
  d::Int           # dimension of Gaussian
  alphas::Vec      # log weights: n
  mus::Array{Vec}  # means: n, each dx1
  qs::Array{Vec}  # square-root-inverse-covariances, log(diagonal): n, each d x 1
  Ls::Array{Vec}  # square-root-inverse-covariances, lower triangle: n, each d*(d-1)/2 x 1
end

# Convert simple GMM to lpGMM
function lpGMM(g::GMM)
  Ls = Array{Vec}(g.n)
  qs = Array{Vec}(g.n)
  for k=1:g.n
    L = inv(chol(g.sigmas[k].data, Val{:L}))
    q, L = ltri_pack(L)
    qs[k], Ls[k] = vec(log(q)), vec(L)
  end
  lpGMM(g.n,g.d,log(g.alphas),g.mus,qs,Ls)
end

# Convert log-parameterized-GMM to simple GMM UnivariateGMM
function GMM(l::lpGMM)
  alphas::Vec = exp(l.alphas)/sum(exp(l.alphas))
  mus::Array{Vec} = l.mus
  Qs = [ltri_unpack(exp(l.qs[i]), l.Ls[i]) for i=1:l.n]
  sigmas::Array{SymMat} = map(A->inv(Symmetric(A'*A)), Qs)
  GMM(l.n,l.d,alphas,mus,sigmas)
end

g = lpGMM(test_gmm)
@printf("testgmm=%s\n**\n", test_gmm)
@printf("gmm=%s\n**\n", GMM(g))

#[markdown]
# lpGMM log-likelihood

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

const halflog2π = log(2π)/2

# Compute log-likelihood
# This version is easy to read, but we can do better (see below)
function log_likelihood_reference(g::lpGMM, x::Vec)
  total = 0
  weights = exp(g.alphas)
  weights /= sum(weights) 
  for k=1:g.n
    L_diagonal = exp(g.qs[k])
    L_ltri_entries = g.Ls[k] 
    InvLowerTriangle = ltri_unpack(L_diagonal, L_ltri_entries)
    mean = g.mus[k]
    mahalanobis = sumsq(InvLowerTriangle * (mean - x))
    total += weights[k] * det(InvLowerTriangle) * exp(-0.5*mahalanobis)
  end
  log(total) - halflog2π*g.d
end

ll1 = log_likelihood_reference(g, x)
@printf("ll0=%f, ll1=%f, ratio to true=%f\n", ll0, ll1, ll0/ll1)
@test_approx_eq_eps ll0 ll1 1e-12

include("logsumexp.jl")

# Cleanest log_likelihood implementation
function log_likelihood(g::lpGMM, x::Vec)
    # function to combine log-diagonal an lower triangle
    get_Q(L_log_diagonal, L_ltri_entries) = ltri_unpack(exp(L_log_diagonal), L_ltri_entries)

    # mahalanobis distances squared
    d_mahals = [0.5*sumsq(get_Q(g.qs[i], g.Ls[i])*(g.mus[i] - x)) for i in 1:g.n]
    
    # log determinants
    log_determinants = [sum(g.qs[i]) for i in 1:g.n]
    
    logsumexp(g.alphas + log_determinants - d_mahals) - logsumexp(g.alphas) - halflog2π*g.d
end

ll2 = log_likelihood(g, x)
@printf("ll0=%f, ll2=%f, ratio to true=%f\n", ll0, ll2, ll0/ll2)
@test_approx_eq_eps ll0 ll2 1e-12

using ReverseDiffSource
rdiff(log_likelihood, (g,x))

