# charlm of lexicon. cluster all hidden states.

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 all hidden states
# 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!
# hidden_states(model, batch{Int64}(B,T)) => (H,B,T+1)

function all_hidden_states(model, data::Vector{Matrix{Int64}})
    hs = map(data) do batch
        hbt = hidden_states(model, batch)
        htb = permutedims(hbt, (1,3,2)) # we need (H,T+1,B) before hcat to get hiddens for a word contiguous
        return reshape(htb, size(htb,1), :)
    end
    reduce(hcat, hs)
end

#@show summary(final_hidden_states(model, tst))

function insert_markers(word::String)
    # produce labels for all positions in the word
    chars = collect(word)
    map(0:length(chars)) do i
        String(vcat(chars[1:i], ['.'], chars[i+1:end]))
    end
end

function all_surface_labels(batch) 
    words = reduce(*, surface_vocab.i2w[batch]; dims=2)
    reduce(vcat, insert_markers.(words))
end

all_surface_labels(data::Vector{Matrix{Int64}}) = reduce(vcat, all_surface_labels.(data))


if !@isdefined tstahs
    tstahs = all_hidden_states(model, tst) |> permutedims # (104452, 350)
    tstalb = all_surface_labels(tst)
end

if !@isdefined tst_ahs_tsne
    n = 3000
    r = randperm(size(tstahs,1))[1:n]
    tst_ahs_tsne = tsne(tstahs[r,:])
    tst_ahs_tsne_labels = tstalb[r]
end

using Plots
plotly()
scatter(tst_ahs_tsne[:,1], tst_ahs_tsne[:,2]; hover=tst_ahs_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
