using Random: seed!
endof(a) = lastindex(a)
linspace(start, stop, length) = range(start, stop=stop, length=length)
const srand = seed!

using Distributions
using PyPlot
using Optim

d(a,b) = MixtureModel([Normal(0,1), Normal(b,1)], [a, 1-a])

#loglik(a, b, Y) = sum(logpdf(d(a,b), Y[i]) for i in 1:endof(Y)) # very slow
lpdf(a, b, y) = log(a+(1-a)*exp(b*y-b^2/2)) - y^2/2 - log(2π)/2
loglik(a, b, Y) = sum(lpdf(a, b, Y[i]) for i in 1:endof(Y))

function plot_lik(a₀, b₀, n; seed = 4649, bmin=-1.0, bmax=1.0)
    srand(seed)
    Y = rand(d(a₀, b₀), n)
    L = 201
    a = linspace(0, 1, L)
    b = linspace(bmin, bmax, L)
    f(a, b) = loglik(a, b, Y)
    z = f.(a', b)
    zmax, k = findmax(z)
    #i, j = (k - 1) ÷ L + 1, mod1(k, L) # for Julia v0.6
    j, i = k.I
    z .= exp.(z .- zmax)

    sleep(0.1)
    figure(figsize=(5,5))
    pcolormesh(a, b, z, cmap="CMRmap")
    scatter([a₀], [b₀], label="true", color="cyan")
    scatter([a[i]], [b[j]], label="MLE", color="red")
    legend()
    xlim(0,1)
    xlabel("a")
    ylabel("b")
    title("\$a_0 = $a₀,  b_0 = $b₀,  n = $n\$")
end

@time plot_lik(0.5, 1.0,  2^4, bmin=0, bmax=2)

@time plot_lik(0.5, 1.0,  2^6, bmin=0, bmax=2)

@time plot_lik(0.5, 1.0,  2^8, bmin=0, bmax=2)

@time plot_lik(0.5, 1.0,  2^10, bmin=0, bmax=2)

@time plot_lik(0.5, 0.1, 2^4)

@time plot_lik(0.5, 0.1, 2^6)

@time plot_lik(0.5, 0.1, 2^8)

@time plot_lik(0.5, 0.1, 2^10)

@time plot_lik(0.5, 0.1, 2^12)

@time plot_lik(0.5, 0.1, 2^14)

@time plot_lik(0.5, 0.1, 2^16)

@time plot_lik(0.5, 0.1, 2^18)