# visualize the final states of vocab charlm model after the roots of words

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 unique roots from data: 8864 of them
if !@isdefined(rootpos)
    posdict = Dict{String,Int}()
    rootpos = Dict{String,Set{String}}()
    for line in eachline("trmor2018.filtered")
        local sur, lex, root, pos
        sur,lex = split(line,'\t')
        root, pos = split(lex,['+','^'])[1:2]
        if !haskey(posdict, pos)
            posdict[pos] = 1+length(posdict)
        end
        push!(get!(rootpos,root,Set{String}()), pos)
    end
    roots = collect(keys(rootpos))
end

# construct a dataset with batchsize = 1
if !@isdefined(rootdata)
    rootdata = map(roots) do root
        chars = split(root,"")
        ints = map(chars) do char
            if char == "þ"; char = "ş"; end # a bug in dataset
            if char == "ý"; char = "ı"; end
            if char == "ð"; char = "ğ"; end
            surface_vocab.w2i[char]
        end
        permutedims(ints)       # batch needs (B,T) shape
    end
end

# 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 rootfhs
    rootfhs = final_hidden_states(model, rootdata) |> permutedims # (8864, 350)
end

if !@isdefined root_tsne
    root_tsne = tsne(rootfhs)
end

using Plots, ColorSchemes
plotly()
labels = map(roots) do r; posdict[first(rootpos[r])]; end
poscolor = get(ColorSchemes.tab20, labels, :extrema)
scatter(root_tsne[:,1], root_tsne[:,2]; hover=roots, color=poscolor)


# 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
