# visualize the final hidden states of the vocabulary charlm model.

using TSne, Statistics

include("charlm.jl")

# load model
if !@isdefined(model)
    include(Knet.dir("src/models21/deepmap.jl"))
    include(Knet.dir("src/models21/fileio.jl"))
    model = charlm_model()
    charlm_loss(model, tst) # necessary to initialize weights, fix this
    setweights!(model, "charlm.jld2")
end
#@show charlm_loss(model,tst)

# load lexical forms for labels
if !@isdefined(lexical_forms)
    lexical_forms = last.(collect(trmor_data))
    lexical_vocab = trmor_data.tgt.vocab
    ltst = lexical_forms[rperm[1:80]]
    lval = lexical_forms[rperm[81:160]]
    ltrn = lexical_forms[rperm[161:end]]
end
#@show lexical_vocab.i2w[ltst[1][1,:]]
#@show surface_vocab.i2w[tst[1][1,:]]

# try clustering the final hidden state
# model(charlm_input(tst[1])) |> "37×128×8 CUDA.CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}"
# model[1:end-1](charlm_input(tst[1])) |> "350×128×8 CUDA.CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}"

hidden_states(model, batch) = model[1:end-1](charlm_input(batch)) |> Array{Float64} # Warning: Float32 makes tsne much slower!
final_hidden_states(model, batch) = hidden_states(model, batch)[:,:,end]
final_hidden_states(model, data::Vector{Matrix{Int64}}) = reduce(hcat, (x->final_hidden_states(model,x)).(data))
#@show summary(final_hidden_states(model, tst))
surface_labels(batch) = reduce(*, surface_vocab.i2w[batch]; dims=2)
surface_labels(data::Vector{Matrix{Int64}}) = reduce(vcat, surface_labels.(data))


if !@isdefined tstfhs
    tstfhs = final_hidden_states(model, tst) |> permutedims # (9904, 350)
    tstlbl = surface_labels(tst)
end

if !@isdefined tst_tsne
    n = 3000
    r = randperm(size(tstfhs,1))[1:n]
    tst_tsne = tsne(tstfhs[r,:])
    tst_tsne_labels = tstlbl[r]
end

using Plots
plotly()
scatter(tst_tsne[:,1], tst_tsne[:,2]; hover=tst_tsne_labels)


# From: https://github.com/nassarhuda/JuliaTutorials/blob/master/TSNE/TSNE.ipynb
# function my_tsne_plot(Y,d,labels; marker=5)
#     colorsvec = get(ColorSchemes.tab10, labels, :extrema)
#     plot(legend=false,axis=false,grid=false)
#     tsneplot = scatter!(Y[:,1], Y[:,2], 
#                         marker=(marker,marker,:auto,stroke(0)),alpha=0.8,
#                         background_color=:black,
#                         color=colorsvec)
#     for i = 0:9
#         mx,my = mean(Y[findall(labels.==i),:],dims=1)
#         s = Shape(mx.+[-d,d,d,-d],my.+[-d,-d,d,d])
#         plot!(s,color=ColorSchemes.tab10[i+1],linewidth=0,alpha=0.9)
#         annotate!(mx,my,text(string(i),color=:white))
#     end
#     return tsneplot
# end

# function tsne_mnist(N=3000; plot=false)
#     # using MLDatasets, ColorSchemes, Plots
#     alldata, labels = MNIST.traindata(Float64)
#     data = reshape(permutedims(alldata[:, :, 1:N], (3, 1, 2)),
#                    N, size(alldata, 1)*size(alldata, 2))
#     labels = labels[1:N]
#     rescale(A; dims=1) = (A .- mean(A, dims=dims)) ./ max.(std(A, dims=dims), eps())
#     # Normalize the data, this should be done if there are large scale differences in the dataset
#     X = rescale(data, dims=1)
#     @show summary(X)
#     Y1 = tsne(X, 2, 50, 1000, 20.0)
#     plot ? my_tsne_plot(Y1, 5, labels) : Y1
# end

# Timing: tsne(tstfhs[1:1000,:], 2, 50, 1000, 20.0);

# mnist tsne:        charlm tsne: (why is it taking so much longer with Float32?)
# N    secs  KL      secs  KL
# 1000 3     0.7470  3     0.3175
# 2000 14    1.0122  13    0.3681
# 3000 31    1.1529  25    0.4125
# 4000 48    1.2079

nothing
