text:
using Knet, AutoGrad, Logging

function initweights(atype, hidden, vocab, embed, window, onehotworld, numfilters; worldsize=[5, 20], args=nothing, premb=nothing, winit=1.0)
    weights = Dict()
    input = premb == nothing ? embed : size(premb, 2)

    #first layer
    weights["enc_w_f"] = xavier(Float32, input+hidden, 4*hidden)
    weights["enc_b_f"] = zeros(Float32, 1, 4*hidden)
    weights["enc_b_f"][1:hidden] = 1 # forget gate bias

    weights["enc_w_b"] = xavier(Float32, input+hidden, 4*hidden)
    weights["enc_b_b"] = zeros(Float32, 1, 4*hidden)
    weights["enc_b_b"][1:hidden] = 1 # forget gate bias

    #vocab 300 for embeddings, 512 for standard
    weights["emb_word"] = premb == nothing ? xavier(Float32, vocab, embed) : premb

    #decoder
    if args["percp"] && args["encoding"] == "grid"
        if args["worldatt"] != 0 && length(numfilters) == 1
            win1, win2 = window[1]
            worldfeats = (worldsize[1] - win1 + 1) * (worldsize[2] - win2 + 1)
        else
            sw1 = sum(map(x->x[1], window))
            sw2 = sum(map(x->x[2], window))
            worldfeats = (worldsize[1] - sw1 + length(window)) * (worldsize[2] - sw2 + length(window)) * numfilters[end]
        end
        weights["emb_world"] = xavier(Float32, worldfeats, embed)

        fs = Any[]
        bs = Any[]

        for i=1:length(numfilters)
            inpch = i == 1 ? onehotworld : numfilters[i-1]
            if args["worldatt"] != 0 && i == 2
                inpch = 1
            end

            win1, win2 = window[i]
            push!(fs, xavier(Float32, win1, win2, inpch, numfilters[i]) * winit)
            push!(bs, zeros(Float32, 1, 1, numfilters[i], 1))
        end

        if args["worldatt"] != 0
            wattin = args["attinwatt"] != 0 ? 4*hidden : 2*hidden
            weights["wa1"] = xavier(Float32, wattin, args["worldatt"])
            weights["wa2"] = xavier(Float32, args["worldatt"], numfilters[1])
        end
        weights["filters_w"] = fs
        weights["filters_b"] = bs
    elseif args["percp"]
        weights["emb_world"] = xavier(Float32, onehotworld, embed)
    end

    if !args["percp"] && args["preva"] && !args["att"]
        weights["emb_world"] = xavier(Float32, 4, embed)
        weights["dec_w"] = xavier(Float32, embed + hidden*2, 4*hidden*2)
    elseif !args["percp"] && args["att"] && args["preva"]
        weights["emb_world"] = xavier(Float32, 4, embed)
        weights["dec_w"] = xavier(Float32, embed + hidden*2 + hidden*2, 4*hidden*2)
    elseif args["percp"] && args["att"] && args["preva"]
        weights["dec_w"] = xavier(Float32, embed + 4 + hidden*2 + hidden*2, 4*hidden*2)
    elseif args["percp"] && !args["att"] && args["preva"]
        weights["dec_w"] = xavier(Float32, embed + 4 + hidden*2, 4*hidden*2)
    elseif args["percp"] && args["att"] && !args["preva"]
        weights["dec_w"] = xavier(Float32, embed + hidden*2 + hidden*2, 4*hidden*2)
    elseif args["percp"] && !args["att"] && !args["preva"]
        weights["dec_w"] = xavier(Float32, embed + hidden*2, 4*hidden*2)
    end

    weights["dec_b"] = zeros(Float32, 1, 4*hidden*2)
    weights["dec_b"][1:hidden*2] = 1 # forget gate bias

    #attention
    if args["att"] || args["attinwatt"] != 0
        atthidden = args["attinwatt"] != 0 ? args["attinwatt"] : hidden
        #fenc, benc, dechid
        weights["attention_w"] = xavier(Float32, hidden*2+hidden*2, atthidden)
        weights["attention_v"] = xavier(Float32, atthidden, 1)
    end

    #output
    weights["soft_h"] = xavier(Float32, 2*hidden, 4)

    if args["inpout"]
        weights["soft_inp"] = xavier(Float32, embed, 4)
    end

    if args["att"] && args["attout"]
        weights["soft_att"] = xavier(Float32, hidden*2, 4)
    end

    if args["preva"] && args["prevaout"]
        weights["soft_preva"] = xavier(Float32, 4, 4)
    end

    weights["soft_b"] = zeros(Float32, 1,4)

    for k in keys(weights)
        if startswith(k, "filter")
            ws = map(t->convert(atype, t), weights[k])
            weights[k] = ws
        else
            weights[k] = convert(atype, weights[k])
        end
    end

    return weights
end

function cnn(filters, bias, x)
    inp = x
    for i=1:length(filters)-1
        inp = relu.(conv4(filters[i], inp; padding=0) .+ bias[i])
    end
    inp = sigm.(conv4(filters[end], inp; padding=0) .+ bias[end])
    return transpose(mat(inp))
end

function cnn(filters, bias, worldatt, x)
    inp = x
    for i=1:length(filters)-1
        inp = relu.(conv4(filters[i], inp; padding=0) .+ bias[i])
        if i == 1
            #inp = conv4(worldatt, inp; padding=0)
            inp = relu.(conv4(worldatt, inp; padding=0))
        end
    end
    inp = sigm.(conv4(filters[end], inp; padding=0) .+ bias[end])
    if length(filters) == 1
        #inp = conv4(worldatt, inp; padding=0)
        inp = sigm.(conv4(worldatt, inp; padding=0))
    end
    return transpose(mat(inp))
end

function spatial(emb, x) #multihot
    return x * emb
end

function spatial(filters, bias, emb, x) #grid
    h = cnn(filters, bias, x)
    return h * emb
end

function spatial(filters, bias, emb, worldatt, x) # world att
    h = cnn(filters, bias, worldatt, x)
    return h * emb
end

function lstm(weight,bias,hidden,cell,input)
    gates   = hcat(input,hidden) * weight .+ bias
    hsize   = size(hidden,2)
    forget  = sigm.(gates[:,1:hsize])
    ingate  = sigm.(gates[:,1+hsize:2hsize])
    outgate = sigm.(gates[:,1+2hsize:3hsize])
    change  = tanh.(gates[:,1+3hsize:end])
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh.(cell)
    return (hidden,cell)
end

function attention(states, attention_w, attention_v)
    h = hcat(states[1][2], states[3][end])
    hu = hcat(states[5], h)
    for i=2:(length(states[1])-1)
        hp = hcat(states[1][i+1], states[3][end-i+1])
        h = vcat(h, hp)
        hu = vcat(hu, hcat(states[5], hp))
    end

    raw_att = tanh.(hu * attention_w) * attention_v
    raw_att = raw_att .- maximum(raw_att)

    att_s = exp.(raw_att)
    att_s = att_s ./ sum(att_s)

    att = att_s .* h

    return sum(att, 1), att_s
end

function worldattention(prevh, wa1, wa2)
    h = tanh.(prevh * wa1) * wa2
    h = h .- maximum(h)
    att_p = exp.(h)
    att_p = att_p ./ sum(att_p)
    return reshape(att_p, 1, 1, size(wa2, 2), 1)
end

#dropout before the loop might accelerate the execution
function encode(weight1_f, bias1_f, weight1_b, bias1_b, emb, state, words; dropout=false, pdrops=[0.5, 0.5])
    for i=1:length(words)
        x = words[i] * emb

        if dropout && pdrops[1] > 0.0
            x = x .* (rand!(similar(AutoGrad.getval(x))) .> pdrops[1]) * (1/(1-pdrops[1]))
        end

        state[1][i+1], state[2][i+1] = lstm(weight1_f, bias1_f, state[1][i], state[2][i], x)

        x = words[end-i+1] * emb

        if dropout && pdrops[1] > 0.0
            x = x .* (rand!(similar(AutoGrad.getval(x))) .> pdrops[1]) * (1/(1-pdrops[1]))
        end

        state[3][i+1], state[4][i+1] = lstm(weight1_b, bias1_b, state[3][i], state[4][i], x)
    end
end

#x might be view or preva
function decode(weight, bias, soft_h, soft_b, state, x; soft_inp=nothing, soft_att=nothing,
    soft_preva=nothing, preva=nothing, att=nothing, dropout=false, pdrops=[0.5, 0.5], prevainp=false)

    inp = x

    if preva != nothing && prevainp
        inp = hcat(inp, preva)
    end

    if att != nothing
        inp = hcat(inp, att)
    end

    if dropout && pdrops[1] > 0.0
        inp = inp .* (rand!(similar(AutoGrad.getval(inp))) .> pdrops[1]) * (1/(1-pdrops[1]))
    end

    state[5], state[6] = lstm(weight, bias, state[5], state[6], inp)

    inph = state[5]
    if dropout && pdrops[2] > 0.0
        inph = inph .* (rand!(similar(AutoGrad.getval(inph))) .> pdrops[2]) * (1/(1-pdrops[2]))
    end

    q = (inph * soft_h) .+ soft_b

    if soft_inp != nothing
        q = q + inp[:, 1:size(x,2)] * soft_inp
    end

    if soft_preva != nothing
        q = q + preva * soft_preva
    end

    if soft_att != nothing
        q = q + att * soft_att
    end

    return q
end

function probs(linear)
    linear = linear .- maximum(linear, 2)
    ps = exp.(linear) ./ sum(exp.(linear), 2)
    return ps
end

function sample(ps)
    c_probs = cumsum(ps, 2)
    return indmax(c_probs .> rand())
end

function discount(rewards; γ=0.9)
    discounted = zeros(Float32, length(rewards), 1)
    discounted[end] = rewards[end]

    for i=(length(rewards)-1):-1:1
        discounted[i] = rewards[i] + γ * discounted[i+1]
    end
    return discounted
end

function loss(w, state, words, ys; lss=nothing, views=nothing, as=nothing, dropout=false, encpdrops=[0.5, 0.5], decpdrops=[0.5, 0.5], args=nothing)
    total = 0.0; count = 0

    #encode
    encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"],
        state, words; dropout=dropout, pdrops=encpdrops)

    state[5] = hcat(state[1][end], state[3][end])
    state[6] = hcat(state[2][end], state[4][end])

    #decode
    for i=1:length(ys)
        att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

        if !args["percp"]
            x = spatial(w["emb_world"], as[i])
        elseif args["encoding"] == "grid" && args["worldatt"] != 0
            wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
            worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
            x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, views[i]) # world att
        elseif args["encoding"] == "grid"
            x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], views[i])
        else
            x = spatial(w["emb_world"], views[i])
        end

        soft_inp = args["inpout"] ? w["soft_inp"] : nothing
        soft_att = args["attout"] ? w["soft_att"] : nothing
        soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
        preva = !args["preva"] ? nothing : as[i]
        prevainp = args["preva"] && args["percp"]

        att = args["att"] ? att : nothing

        ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
            soft_preva=soft_preva, preva=preva, att=att, dropout=dropout, pdrops=decpdrops, prevainp=prevainp)

        ynorm = logp(ypred,2)
        total += sum(ys[i] .* ynorm)

        count += 1
    end

    nll = -total/count
    lss[1] = AutoGrad.getval(nll)
    lss[2] = AutoGrad.getval(count)
    return nll
end

function loss_global(w, words, goldactions, instruction, maps, state; dropout=false, encpdrops=[0.5, 0.5], decpdrops=[0.5, 0.5], args=nothing, lss=nothing)
    nll = 0.0
    beamsize = args["beamsize"]

    encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"], state, words; dropout=dropout, pdrops=encpdrops)

    state[5] = hcat(state[1][end], state[3][end])
    state[6] = hcat(state[2][end], state[4][end])

    cstate5 = copy(state[5])
    cstate6 = copy(state[6])

    current = instruction.path[1]
    #1: score
    #2: prev hidden
    #3: prev cell
    #4: current location
    #5: depth
    #6: stopped
    #7: isgold
    #8: prev actions
    cands = Array{Tuple{Float64, Any, Any, Any, Any, Any, Any, Any}, 1}()
    push!(cands, (0.0, cstate5, cstate6, current, 0, false, true, Any[4]))

    nactions = 0
    stop = false
    stopsearch = false
    araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
    goldindex = 0

    while !stopsearch
        newcands = Array{Tuple{Any, Any, Any, Any, Any, Any, Any, Any}, 1}()
        newcand = false
        for cand in cands
            current = cand[4]
            state[5] = cand[2]
            state[6] = cand[3]
            depth = cand[5]
            prevActions = cand[end]
            lastAction = prevActions[end]
            stopped = cand[6]
            isgold = cand[7]

            if araw != nothing
                araw[:] = 0.0
                araw[1, lastAction] = 1.0
            end

            if stopped
                push!(newcands, cand)
                continue
            end

            view = !args["percp"] ? nothing : args["encoding"] == "grid" ?
            state_agent_centric(maps[instruction.map], current) : state_agent_centric_multihot(maps[instruction.map], current)
            view = args["percp"] ? convert(KnetArray{Float32}, view) : nothing
            preva = araw != nothing ? convert(KnetArray{Float32}, araw)  : araw

            if !args["percp"]
                x = spatial(w["emb_world"], preva)
            elseif args["encoding"] == "grid" && args["worldatt"] != 0
                worldatt = worldattention(state[5], w["wa1"], w["wa2"])
                x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
            elseif args["encoding"] == "grid"
                x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
            else
                x = spatial(w["emb_world"], view)
            end

            soft_inp = args["inpout"] ? w["soft_inp"] : nothing
            soft_att = args["attout"] ? w["soft_att"] : nothing
            soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
            preva = !args["preva"] ? nothing : preva
            prevainp = args["preva"] && args["percp"]

            att,_ = args["att"] ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)
            ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
            soft_preva=soft_preva, preva=preva, att=att, dropout=dropout, pdrops=decpdrops, prevainp=prevainp)

            for i=1:4
                myisgold = isgold
                mystop = stopped
                nactions = depth + 1
                actcopy = copy(prevActions)

                if isgold
                    myisgold = i == goldactions[nactions]
                end

                if nactions >= length(goldactions) && i < 4
                    mystop = true
                end
                push!(actcopy, i)
                cur = identity(current)

                if i < 4
                    cur = getlocation(maps[instruction.map], cur, i)
                end

                if (i == 1 && !haskey(maps[instruction.map].edges[(current[1], current[2])], (cur[1], cur[2]))) || i==4
                    mystop = true
                end

                newscore = cand[1] + ypred[1, i]

                push!(newcands, (newscore, state[5], state[6], cur, nactions, mystop, myisgold, actcopy))
                newcand = true
            end
        end
        stopsearch = !newcand
        if newcand
            sort!(newcands; by=x->x[1], rev=true)
            l = length(newcands) < beamsize ? length(newcands) : beamsize
            goldindex = findfirst(x->x[7], newcands)
            cands = newcands[1:l]
            debug("l: $l")
            for index=1:l
                debug("Beam #$index : $(AutoGrad.getval(newcands[index][1])) - $(newcands[index][end])")
            end

            if goldindex > beamsize || newcands[goldindex][5] == length(goldactions)
                stopsearch = true
                glob = vcat(first(cands[1]))
                for ind=2:l
                    glob = vcat(glob, first(cands[ind]))
                end

                if goldindex > beamsize
                    debug("gind: $goldindex")
                    glob = vcat(glob, first(newcands[goldindex]))
                    ynorm = logp(glob)
                    debug("ynorm: $(Array(AutoGrad.getval(ynorm)))")
                    nll = -ynorm[end]
                else
                    debug("gind: $goldindex")
                    ynorm = logp(glob)
                    debug("ynorm: $(Array(AutoGrad.getval(ynorm)))")
                    nll = -ynorm[goldindex]
                end

                debug("nll: $(AutoGrad.getval(nll))")
                debug("gold: $(AutoGrad.getval(newcands[goldindex][1]))")
                debug("goldactions: $(goldactions)")
                #=
                debug("glob: $(AutoGrad.getval(glob))")
                z = log(sum(exp(glob)))
                debug("z: $(AutoGrad.getval(z))")
                nll = -newcands[goldindex][1] + z
                debug("gold: $(AutoGrad.getval(newcands[goldindex][1]))")
                debug("nll: $(AutoGrad.getval(nll))")
                debug("goldactions: $(goldactions)")
                debug("beam-gold: $(newcands[goldindex][end])")
                for index=1:l
                    debug("Beam #$index : $(newcands[index][end])")
                end
                =#
            end
        end
    end
    debug("Exiting nll: $(AutoGrad.getval(nll))")
    lss[1] = AutoGrad.getval(nll)
    return nll
end

lossgradient = grad(loss)
lossglobalgradient = grad(loss_global)

function clip_grad(g; args=args)
    gclip = args["gclip"]
    if gclip > 0
        gnorm = 0
        for k in keys(g)
            if startswith(k, "filter")
                for el in g[k]
                    gnorm += sum(abs2, el)
                end
            else
                gnorm += sum(abs2, g[k])
            end
        end
        gnorm = sqrt(gnorm)

        #debug("Gnorm: $gnorm")

        if gnorm > gclip
            for k in keys(g)
                if startswith(k, "filter")
                    gs = Any[]
                    for el in g[k]
                        push!(gs, el * gclip / gnorm)
                    end
                    g[k] = gs
                else
                    g[k] = g[k] * gclip / gnorm
                end
            end
        end
    end
end

function updatew(w,g,prms)
    #update weights
    for k in keys(g)
        if startswith(k, "filter")
            for ind=1:length(g[k])
                Knet.update!(w[k][ind], g[k][ind], prms[k][ind])
            end
        else
            Knet.update!(w[k], g[k], prms[k])
        end
    end
end

function train(w, prms, data; args=nothing, updatelimit=0)
    lss = 0.0
    cnt = 0.0
    nll = Float32[0, 0]
    for (words, views, ys, _) in data
        bs = size(words[1], 1)
        state = initstate(KnetArray{Float32}, args["hidden"], bs, length(words))

        #load data to gpu
        words = map(t->convert(KnetArray{Float32}, t), words)
        ys = map(t->convert(KnetArray{Float32}, t), ys)

        views = args["percp"] ? map(v->convert(KnetArray{Float32}, v), views) : nothing
        acts = nothing

        if args["preva"]
            as = Any[]
            push!(as, reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4))
            append!(as, ys[1:(end-1)])
            acts = args["preva"] ? map(t->convert(KnetArray{Float32}, t), as) : nothing
        end

        g = lossgradient(w, state, words, ys; lss=nll, views=views, as=acts, dropout=true, encpdrops=args["encdrops"], decpdrops=args["decdrops"], args=args)
        clip_grad(g;args=args)
        updatew(w,g,prms)

        lss += nll[1] * nll[2]
        cnt += nll[2]

        if updatelimit != 0 && cnt >= updatelimit
            break
        end
    end
    return lss / cnt
end

function train_global(w, prms, data, maps; args=nothing)
    lss = 0.0
    cnt = 0.0
    nll = Float64[0, 0]

    for (instruction, words) in data
        words = map(t->convert(KnetArray{Float32}, t), words)
        goldactions = getactions(instruction.path)
        bs = size(words[1], 1)
        state = initstate(KnetArray{Float32}, args["hidden"], bs, length(words))

        g = lossglobalgradient(w, words, goldactions, instruction, maps, state; dropout=true, encpdrops=args["encdrops"], decpdrops=args["decdrops"], args=args, lss=nll)
        clip_grad(g;args=args)
        updatew(w,g,prms)

        lss += nll[1]
        cnt += 1.0
    end

    return lss / cnt
end

function train_loss(w, data; args=nothing)
    lss = 0.0
    cnt = 0.0
    nll = Float32[0, 0]
    for (words, views, ys, _) in data
        bs = size(words[1], 1)
        state = initstate(KnetArray{Float32}, args["hidden"], bs, length(words))

        #load data to gpu
        words = map(t->convert(KnetArray{Float32}, t), words)
        ys = map(t->convert(KnetArray{Float32}, t), ys)

        views = args["percp"] ? map(v->convert(KnetArray{Float32}, v), views) : nothing
        acts = nothing

        if args["preva"]
            as = Any[]
            push!(as, reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4))
            append!(as, ys[1:(end-1)])
            acts = args["preva"] ? map(t->convert(KnetArray{Float32}, t), as) : nothing
        end

        loss(w, state, words, ys; lss=nll, views=views, as=acts, dropout=false, args=args)

        lss += nll[1] * nll[2]
        cnt += nll[2]
    end
    return lss / cnt
end

function test(models, data, maps; args=nothing)
    scss = 0.0

    for (instruction, words) in data
        words = map(v->convert(KnetArray{Float32},v), words)
        states = map(weights->initstate(KnetArray{Float32}, args["hidden"], 1, length(words)), models)

        for ind=1:length(models)
            w = models[ind]
            state = states[ind]
            encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"],
                state, words; dropout=false)

            state[5] = hcat(state[1][end], state[3][end])
            state[6] = hcat(state[2][end], state[4][end])
        end

        current = instruction.path[1]
        nactions = 0
        stop = false

        actions = Any[]
        araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
        while !stop
            preva = araw != nothing ? convert(KnetArray{Float32}, araw)  : araw

            view = !args["percp"] ? nothing : args["encoding"] == "grid" ?
                state_agent_centric(maps[instruction.map], current) : state_agent_centric_multihot(maps[instruction.map], current)
            view = args["percp"] ? convert(KnetArray{Float32}, view) : nothing

            cum_ps = zeros(Float32, 1, 4)
            for ind=1:length(models)
                w = models[ind]
                state = states[ind]

                att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

                if !args["percp"]
                    x = spatial(w["emb_world"], preva)
                elseif args["encoding"] == "grid" && args["worldatt"] != 0
                    wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
                    worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
                    x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
                elseif args["encoding"] == "grid"
                    x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
                else
                    x = spatial(w["emb_world"], view)
                end

                soft_inp = args["inpout"] ? w["soft_inp"] : nothing
                soft_att = args["attout"] ? w["soft_att"] : nothing
                soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
                preva = !args["preva"] ? nothing : preva
                prevainp = args["preva"] && args["percp"]

                att = args["att"] ? att : nothing
                ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
                soft_preva=soft_preva, preva=preva, att=att, dropout=false, prevainp=prevainp)

                cum_ps += probs(Array(ypred))
            end

            cum_ps = cum_ps ./ length(models)
            debug("Probs: $(cum_ps)")
            action = 0
            if args["greedy"]
                action = indmax(cum_ps)
            else
                action = sample(cum_ps)
            end

            if args["preva"]
                araw[:] = 0.0
                araw[1, action] = 1.0
            end

            push!(actions, action)
            prev = current
            current = getlocation(maps[instruction.map], current, action)
            nactions += 1

            nowall = false
            if action == 1
                nowall = !haskey(maps[instruction.map].edges[(prev[1], prev[2])], (current[1], current[2]))
            end

            stop = nactions > args["limactions"] || action == 4 || nowall

        end

        debug("$(instruction.text)")
        debug("Path: $(instruction.path)")
        debug("Filename: $(instruction.fname)")

        debug("Actions: $(reshape(collect(actions), 1, length(actions)))")
        debug("Current: $(current)")

        if current == instruction.path[end]
            scss += 1
            debug("SUCCESS\n")
        else
            debug("FAILURE\n")
        end
    end

    return scss / length(data)
end

function test_paragraph(models, groups, maps; args=nothing)
    scss = 0.0
    for data in groups
        debug("\nNew paragraph")
        current = data[1][1].path[1]

        for i=1:length(data)
            instruction, words = data[i]
            words = map(v->convert(KnetArray{Float32},v), words)
            states = map(weights->initstate(KnetArray{Float32}, args["hidden"], 1, length(words)), models)

            for ind=1:length(models)
                w = models[ind]
                state = states[ind]
                encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"],
                state, words; dropout=false)

                state[5] = hcat(state[1][end], state[3][end])
                state[6] = hcat(state[2][end], state[4][end])
            end

            nactions = 0
            stop = false
            action = 0
            actions = Any[]
            araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
            while !stop
                preva = araw != nothing ? convert(KnetArray{Float32}, araw)  : araw

                view = !args["percp"] ? nothing : args["encoding"] == "grid" ?
                state_agent_centric(maps[instruction.map], current) : state_agent_centric_multihot(maps[instruction.map], current)
                view = args["percp"] ? convert(KnetArray{Float32}, view) : nothing

                cum_ps = zeros(Float32, 1, 4)
                for ind=1:length(models)
                    w = models[ind]
                    state = states[ind]

                    att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

                    if !args["percp"]
                        x = spatial(w["emb_world"], preva)
                    elseif args["encoding"] == "grid" && args["worldatt"] != 0
                        wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
                        worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
                        x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
                    elseif args["encoding"] == "grid"
                        x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
                    else
                        x = spatial(w["emb_world"], view)
                    end

                    soft_inp = args["inpout"] ? w["soft_inp"] : nothing
                    soft_att = args["attout"] ? w["soft_att"] : nothing
                    soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
                    preva = !args["preva"] ? nothing : preva
                    prevainp = args["preva"] && args["percp"]

                    att = args["att"] ? att : nothing
                    ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
                    soft_preva=soft_preva, preva=preva, att=att, dropout=false, prevainp=prevainp)

                    cum_ps += probs(Array(ypred))
                end

                cum_ps = cum_ps ./ length(models)
                debug("Probs: $(cum_ps)")
                if args["greedy"]
                    action = indmax(cum_ps)
                else
                    action = sample(cum_ps)
                end

                if args["preva"]
                    araw[:] = 0.0
                    araw[1, action] = 1.0
                end

                push!(actions, action)
                prev = current
                current = getlocation(maps[instruction.map], current, action)
                nactions += 1

                nowall = false
                if action == 1
                    nowall = !haskey(maps[instruction.map].edges[(prev[1], prev[2])], (current[1], current[2]))
                end

                stop = nactions > args["limactions"] || action == 4 || nowall
            end

            debug("$(instruction.text)")
            debug("Path: $(instruction.path)")
            debug("Filename: $(instruction.fname)")

            debug("Actions: $(reshape(collect(actions), 1, length(actions)))")
            debug("Current: $(current)")

            if action != 4
                debug("FAILURE")
                break
            end

            if i == length(data)
                if current[1] == instruction.path[end][1] && current[2] == instruction.path[end][2]
                    scss += 1
                    debug("SUCCESS\n")
                else
                    debug("FAILURE\n")
                end
            end
        end
    end
    return scss / length(groups)
end

function initparams(ws; args=nothing)
    prms = Dict()

    for k in keys(ws)
        if startswith(k, "filter")
            prms[k] = map(w->Adam(lr=args["lr"]), ws[k])
            #prms[k] = map(w->Sgd(;lr=args["lr"]), ws[k])
        else
            prms[k] = Adam(lr=args["lr"])
            #prms[k] = Sgd(;lr=args["lr"])
        end
    end

    return prms
end

# state[2k-1,2k]: hidden and cell for the k'th lstm layer
function initstate(atype, hidden, batchsize, length)
    state = Array{Any}(6)
    #forward
    state[1] = Array{Any}(length+1)
    for i=1:(length+1); state[1][i] = convert(atype, zeros(batchsize, hidden)); end

    state[2] = Array{Any}(length+1)
    for i=1:(length+1); state[2][i] = convert(atype, zeros(batchsize, hidden)); end

    #backward
    state[3] = Array{Any}(length+1)
    for i=1:(length+1); state[3][i] = convert(atype, zeros(batchsize, hidden)); end

    state[4] = Array{Any}(length+1)
    for i=1:(length+1); state[4][i] = convert(atype, zeros(batchsize, hidden)); end

    state[5] = convert(atype, zeros(batchsize, hidden*2))
    state[6] = convert(atype, zeros(batchsize, hidden*2))

    return state
end

function test_beam(models, data, maps; args=nothing)
    beamsize = args["beamsize"]

    scss = 0.0

    for (instruction, words) in data
        words = map(v->convert(KnetArray{Float32},v), words)
        states = map(weights->initstate(KnetArray{Float32}, args["hidden"], 1, length(words)), models)

        for i=1:length(models)
            w = models[i]
            state = states[i]
            encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"], state, words; dropout=false)

            state[5] = hcat(state[1][end], state[3][end])
            state[6] = hcat(state[2][end], state[4][end])
        end

        cstate5 = Any[]
        cstate6 = Any[]

        for m=1:length(models)
            push!(cstate5, copy(states[m][5]))
            push!(cstate6, copy(states[m][6]))
        end

        current = instruction.path[1]
        cands = Any[(1.0, cstate5, cstate6, current, 0, false, Any[4])]

        nactions = 0
        stop = false
        stopsearch = false
        araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
        while !stopsearch
            newcands = Any[]
            newcand = false
            for cand in cands
                current = cand[4]
                for m=1:length(models)
                    states[m][5] = copy(cand[2][m])
                    states[m][6] = copy(cand[3][m])
                end
                depth = cand[5]
                prevActions = cand[end]
                lastAction = prevActions[end]
                stopped = cand[6]

                if araw != nothing
                    araw[:] = 0.0
                    araw[1, lastAction] = 1.0
                end

                if stopped
                    push!(newcands, cand)
                    continue
                end

                view = !args["percp"] ? nothing : args["encoding"] == "grid" ?
                    state_agent_centric(maps[instruction.map], current) : state_agent_centric_multihot(maps[instruction.map], current)
                view = args["percp"] ? convert(KnetArray{Float32}, view) : nothing
                preva = araw != nothing ? convert(KnetArray{Float32}, araw)  : araw

                cum_ps = zeros(Float32, 1, 4)
                #score = KnetArray(zeros(Float32, 1, 4))

                for ind=1:length(models)
                    w = models[ind]
                    state = states[ind]

                    att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

                    if !args["percp"]
                        x = spatial(w["emb_world"], preva)
                    elseif args["encoding"] == "grid" && args["worldatt"] != 0
                        wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
                        worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
                        x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
                    elseif args["encoding"] == "grid"
                        x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
                    else
                        x = spatial(w["emb_world"], view)
                    end

                    soft_inp = args["inpout"] ? w["soft_inp"] : nothing
                    soft_att = args["attout"] ? w["soft_att"] : nothing
                    soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
                    preva = !args["preva"] ? nothing : preva
                    prevainp = args["preva"] && args["percp"]

                    att = args["att"] ? att : nothing
                    ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
                    soft_preva=soft_preva, preva=preva, att=att, dropout=false, prevainp=prevainp)
                    cum_ps += probs(Array(ypred))
                    #score += ypred
                end

                cum_ps = cum_ps ./ length(models)
                #score = score ./ length(models)
                debug("Probs: $(cum_ps)")
                #debug("Score: $(Array(score))")

                for i=1:4
                    mystop = stopped
                    nactions = depth + 1
                    actcopy = copy(prevActions)
                    if nactions > args["limactions"] && i < 4
                        mystop = true
                    end
                    push!(actcopy, i)
                    cur = identity(current)
                    if i < 4
                        cur = getlocation(maps[instruction.map], cur, i)
                    end

                    if (i == 1 && !haskey(maps[instruction.map].edges[(current[1], current[2])], (cur[1], cur[2]))) || i==4
                        mystop = true
                    end

                    cstate5 = Any[]
                    cstate6 = Any[]
                    for m=1:length(models)
                        push!(cstate5, copy(states[m][5]))
                        push!(cstate6, copy(states[m][6]))
                    end

                    push!(newcands, (cand[1] * cum_ps[1, i], cstate5, cstate6, cur, nactions, mystop, actcopy))
                    #push!(newcands, (cand[1] + score[1, i], cstate5, cstate6, cur, nactions, mystop, actcopy))
                    newcand = true
                end
            end
            stopsearch = !newcand
            if newcand
                sort!(newcands; by=x->x[1], rev=true)
                l = length(newcands) < beamsize ? length(newcands) : beamsize
                cands = newcands[1:l]
            end
        end
        current = cands[1][4]
        actions = cands[1][end]
        debug("$(instruction.text)")
        debug("Path: $(instruction.path)")
        debug("Filename: $(instruction.fname)")

        debug("Actions: $(reshape(collect(actions), 1, length(actions)))")
        debug("Current: $(current)")
        debug("Prob: $(cands[1][1])")
        #debug("Score: $(cands[1][1])")

        if current == instruction.path[end]
            scss += 1
            debug("SUCCESS\n")
        else
            debug("FAILURE\n")
        end
    end

    return scss / length(data)
end

function test_paragraph_beam(models, groups, maps; args=nothing)
    beamsize = args["beamsize"]

    scss = 0.0

    for data in groups
        debug("\nNew paragraph")
        current = data[1][1].path[1]
        cands = Any[]
        for indx=1:length(data)
            instruction, words = data[indx]
            words = map(v->convert(KnetArray{Float32},v), words)
            states = map(weights->initstate(KnetArray{Float32}, args["hidden"], 1, length(words)), models)

            for ind=1:length(models)
                w = models[ind]
                state = states[ind]
                encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"], state, words; dropout=false)

                state[5] = hcat(state[1][end], state[3][end])
                state[6] = hcat(state[2][end], state[4][end])
            end
            cstate5 = Any[]
            cstate6 = Any[]
            for m=1:length(models)
                push!(cstate5, copy(states[m][5]))
                push!(cstate6, copy(states[m][6]))
            end

            if length(cands) == 0
                cands = Any[(1.0, cstate5, cstate6, current, 0, false, true, Any[4])]
            else
                cands = map(x->(x[1], cstate5, cstate6, x[4], 0, false, x[7], x[end]), cands)
            end

            nactions = 0
            stop = false
            stopsearch = false

            araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
            while !stopsearch
                newcands = Any[]
                newcand = false
                for cand in cands
                    current = cand[4]
                    for m=1:length(models)
                        states[m][5] = copy(cand[2][m])
                        states[m][6] = copy(cand[3][m])
                    end
                    depth = cand[5]
                    prevActions = cand[end]
                    lastAction = prevActions[end]
                    stopped = cand[6]
                    if araw != nothing
                        araw[:] = 0.0
                        araw[1, lastAction] = 1.0
                    end

                    if stopped || !cand[7]
                        c = (cand[1], cand[2], cand[3], cand[4], cand[5], true, cand[7], cand[end])
                        push!(newcands, c)
                        continue
                    end

                    preva = araw != nothing ? convert(KnetArray{Float32}, araw)  : araw

                    view = !args["percp"] ? nothing : args["encoding"] == "grid" ?
                    state_agent_centric(maps[instruction.map], current) : state_agent_centric_multihot(maps[instruction.map], current)
                    view = args["percp"] ? convert(KnetArray{Float32}, view) : nothing

                    cum_ps = zeros(Float32, 1, 4)

                    for ind=1:length(models)
                        w = models[ind]
                        state = states[ind]

                        att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

                        if !args["percp"]
                            x = spatial(w["emb_world"], preva)
                        elseif args["encoding"] == "grid" && args["worldatt"] != 0
                            wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
                            worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
                            x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
                        elseif args["encoding"] == "grid"
                            x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
                        else
                            x = spatial(w["emb_world"], view)
                        end

                        soft_inp = args["inpout"] ? w["soft_inp"] : nothing
                        soft_att = args["attout"] ? w["soft_att"] : nothing
                        soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
                        preva = !args["preva"] ? nothing : preva
                        prevainp = args["preva"] && args["percp"]

                        att = args["att"] ? att : nothing
                        ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
                        soft_preva=soft_preva, preva=preva, att=att, dropout=false, prevainp=prevainp)
                        cum_ps += probs(Array(ypred))
                    end

                    cum_ps = cum_ps ./ length(models)
                    debug("Probs: $(cum_ps)")
                    for i=1:4
                        mystop = stopped
                        legitimate = true
                        nactions = depth + 1
                        actcopy = copy(prevActions)
                        if nactions > args["limactions"] && i < 4
                            mystop = true
                            legitimate = false
                        end
                        push!(actcopy, i)
                        cur = identity(current)
                        if i < 4
                            cur = getlocation(maps[instruction.map], cur, i)
                        end

                        if (i == 1 && !haskey(maps[instruction.map].edges[(current[1], current[2])], (cur[1], cur[2]))) || i==4
                            mystop = true
                            if i == 1
                                legitimate = false
                            end
                        end

                        cstate5 = Any[]
                        cstate6 = Any[]
                        for m=1:length(models)
                            push!(cstate5, copy(states[m][5]))
                            push!(cstate6, copy(states[m][6]))
                        end

                        push!(newcands, (cand[1] * cum_ps[1, i], cstate5, cstate6, cur, nactions, mystop, legitimate, actcopy))
                        newcand = true
                    end
                end
                stopsearch = !newcand
                if !stopsearch
                    sort!(newcands; by=x->x[1], rev=true)
                    l = length(newcands) < beamsize ? length(newcands) : beamsize
                    cands = newcands[1:l]
                end
            end

            debug("$(instruction.text)")
            debug("Path: $(instruction.path)")
            debug("Filename: $(instruction.fname)")

            if indx == length(data)
                current = cands[1][4]
                actions = cands[1][end]
                debug("Actions: $(reshape(collect(actions), 1, length(actions)))")
                debug("Current: $(current)")
                debug("Prob: $(cands[1][1])")

                if actions[end] != 4
                    debug("FAILURE")
                    break
                end

                if current[1] == instruction.path[end][1] && current[2] == instruction.path[end][2]
                    scss += 1
                    debug("SUCCESS\n")
                else
                    debug("FAILURE\n")
                end
            end
        end
    end

    return scss / length(groups)
end

function predict_beam(models, words, navimap, initial; args=nothing)
    beamsize = args["beamsize"]
    words = map(v->convert(args["atype"],v), words)
    states = map(weights->initstate(args["atype"], args["hidden"], 1, length(words)), models)
    for i=1:length(models)
        w = models[i]
        state = states[i]
        encode(w["enc_w_f"], w["enc_b_f"], w["enc_w_b"], w["enc_b_b"], w["emb_word"], state, words; dropout=false)
        state[5] = hcat(state[1][end], state[3][end])
        state[6] = hcat(state[2][end], state[4][end])
    end

    cstate5 = Any[]
    cstate6 = Any[]

    for m=1:length(models)
        push!(cstate5, copy(states[m][5]))
        push!(cstate6, copy(states[m][6]))
    end

    current = initial
    cands = Any[(1.0, cstate5, cstate6, current, 0, false, Any[4])]

    nactions = 0
    stop = false
    stopsearch = false
    araw = args["preva"] ? reshape(Float32[0.0, 0.0, 0.0, 1.0], 1, 4) : nothing
    while !stopsearch
        newcands = Any[]
        newcand = false
        for cand in cands
            current = cand[4]
            for m=1:length(models)
                states[m][5] = copy(cand[2][m])
                states[m][6] = copy(cand[3][m])
            end
            depth = cand[5]
            prevActions = cand[end]
            lastAction = prevActions[end]
            stopped = cand[6]

            if araw != nothing
                araw[:] = 0.0
                araw[1, lastAction] = 1.0
            end

            if stopped
                push!(newcands, cand)
                continue
            end

            view = !args["percp"] ? nothing : args["encoding"] == "grid" ? state_agent_centric(navimap, current) : state_agent_centric_multihot(navimap, current)
            view = args["percp"] ? convert(args["atype"], view) : nothing
            preva = araw != nothing ? convert(args["atype"], araw)  : araw

            cum_ps = zeros(Float32, 1, 4)

            for ind=1:length(models)
                w = models[ind]
                state = states[ind]

                att,_ = args["att"] || args["attinwatt"] != 0 ? attention(state, w["attention_w"], w["attention_v"]) : (nothing, nothing)

                if !args["percp"]
                    x = spatial(w["emb_world"], preva)
                elseif args["encoding"] == "grid" && args["worldatt"] != 0
                    wattinp = args["attinwatt"] == 0 ? state[5] : hcat(state[5], att)
                    worldatt = worldattention(wattinp, w["wa1"], w["wa2"])
                    x =  spatial(w["filters_w"], w["filters_b"], w["emb_world"], worldatt, view) # world att
                elseif args["encoding"] == "grid"
                    x = spatial(w["filters_w"], w["filters_b"], w["emb_world"], view)
                else
                    x = spatial(w["emb_world"], view)
                end

                soft_inp = args["inpout"] ? w["soft_inp"] : nothing
                soft_att = args["attout"] ? w["soft_att"] : nothing
                soft_preva = args["prevaout"] ? w["soft_preva"] : nothing
                preva = !args["preva"] ? nothing : preva
                prevainp = args["preva"] && args["percp"]

                att = args["att"] ? att : nothing
                ypred = decode(w["dec_w"], w["dec_b"], w["soft_h"], w["soft_b"], state, x; soft_inp=soft_inp, soft_att=soft_att,
                               soft_preva=soft_preva, preva=preva, att=att, dropout=false, prevainp=prevainp)
                cum_ps += probs(Array(ypred))
            end

            cum_ps = cum_ps ./ length(models)

            for i=1:4
                mystop = stopped
                nactions = depth + 1
                actcopy = copy(prevActions)
                if nactions > args["limactions"] && i < 4
                    mystop = true
                end
                push!(actcopy, i)
                cur = identity(current)
                if i < 4
                    cur = getlocation(navimap, cur, i)
                end

                if (i == 1 && !haskey(navimap.edges[(current[1], current[2])], (cur[1], cur[2]))) || i==4
                    mystop = true
                end

                cstate5 = Any[]
                cstate6 = Any[]
                for m=1:length(models)
                    push!(cstate5, copy(states[m][5]))
                    push!(cstate6, copy(states[m][6]))
                end

                push!(newcands, (cand[1] * cum_ps[1, i], cstate5, cstate6, cur, nactions, mystop, actcopy))
                newcand = true
            end
        end

        stopsearch = !newcand

        if newcand
            sort!(newcands; by=x->x[1], rev=true)
            l = length(newcands) < beamsize ? length(newcands) : beamsize
            cands = newcands[1:l]
        end
    end
    current = cands[1][4]
    actions = cands[1][end]
    return actions[2:end]
end

max_stars_repo_name: ozanarkancan/navi-app

max_issues_repo_name: ozanarkancan/navi-app

max_forks_repo_name: ozanarkancan/navi-app


text:
# Linear mean function

@doc """
# Description
Constructor for the Linear mean function

m(x) = xᵀβ
# Arguments:
* `β::Vector{Float64}`: One coefficient for each dimension
""" ->
type MeanLin <: Mean
    β::Vector{Float64}
    dim::Int
    MeanLin(β::Vector{Float64}) = new(β, length(β))
end

mean(mLin::MeanLin,x::Matrix{Float64}) =  x'mLin.β

get_params(mLin::MeanLin) = mLin.β
get_param_names(::MeanLin) = [:β]
num_params(mLin::MeanLin) = mLin.dim

function set_params!(mLin::MeanLin, hyp::Vector{Float64})
    length(hyp) == mLin.dim || throw(ArgumentError("Linear mean function only has $(mLin.dim) parameters"))
    mLin.β = hyp
end

function grad_mean(mLin::MeanLin, x::Vector{Float64})
    dM_theta = x
    return dM_theta
end

ext: jl


text:
module FourierGPE

using Reexport
# using SpecialFunctions
# using LoopVectorization
# using Tullio
using LazyArrays
using FillArrays
using OrdinaryDiffEq
using RecursiveArrayTools
using PaddedViews
using FFTW
import DiffEqCallbacks.FunctionCallingCallback
import OrdinaryDiffEq: ODEProblem, solve

# using LaTeXStrings
# using ColorSchemes
# using Plots
# const c1 = cgrad(ColorSchemes.inferno.colors)
# const c2 = cgrad(ColorSchemes.RdBu_11.colors)
# const c3 =:mediumseagreen

@reexport using Parameters
@reexport using JLD2
@reexport using FileIO

include("types.jl")
include("arrays.jl")
include("transforms.jl")
include("evolution.jl")

# simulation
export Simulation, TransformLibrary, UserParams
export xvec, kvec, xvecs, kvecs, dfft, dfftall
export crandn_array, crandnpartition
export maketransarrays, makearrays, xspace, xspace!, kspace, kspace!
export nlin, nlin!, Lgp, Lgp!, V, Params
export initsim!, runsim, internalnorm
export Transforms, @pack_Transforms!, @unpack_Transforms
export Sim, @pack_Sim!, @unpack_Sim, testsim
export Params, @pack_Params!, @unpack_Params
export k2, @pack!, @unpack, makeT, definetransforms #makeTMixed,
export Field, XField, KField
# export showpsi, c1, c2, c3

end


text:
module TestModScope
using Debug
import Debug.BPNode, Debug.Node, Debug.Scope, Debug.debug_eval
export trap

trap(node::BPNode, scope::Scope) = (global firstline = node.loc.line)
function trap(node::Node, scope::Scope)
    global firstline
    line = node.loc.line - firstline + 1

    if (line == 3); debug_eval(scope, :(x = 5)) end
    if (line == 4)
        debug_eval(scope, :(let; global y = 7; end))
        debug_eval(scope, :(z = 9))
    end
end

module Mod
using Debug, TestModScope
import Debug.@instrument
@instrument trap begin
    @bp             # 1
    x, y, z = 1, 2, 3
    let             # 3
        local z = 1 # 4
    end
    @assert (x, y, z) == (5, 7, 3)
end
end # module Mod

end


text:
using DataFrames, CSV
using Plots
using LaTeXStrings
using BSON: @load
using Dates
pyplot()

# Fonts:

bbfont = Plots.font("Helvetica", 15)
bfont = Plots.font("Helvetica", 13)
mfont = Plots.font("Helvetica", 12)
sfont = Plots.font("Helvetica", 10)
fonts = Dict(:guidefont=>bfont, :xtickfont=>sfont, :ytickfont=>sfont, :legendfont=>sfont)

# Loss graphs:

function loss_graph!(title, path)

    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>bfont)
    df = DataFrame(CSV.File(path))
    iter = range(1, length=length(df.Train))
    p = plot(iter, df.Train, label="train", color="blue",size=(500,400),w=1.5; fonts...)
    plot!(iter, df.Test, label="validation",color="blue",w=1.5,linestyle=:dot)

    ylabel!("loss")
    xlabel!("epochs")
    savefig("../paper/$(title)")
    savefig("../paper/$(title).eps")
end

path = ""
title = ""
loss_graph!(title, path)

function kappa_threshold_loss_graph!(title, l25, l50)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>bfont)
    iterations = 125
    iter = range(1, length=iterations)
    p = plot(iter, min.(l25[1:iterations],l25[1]), label=L"\quad \kappa^2 \in [0.25,1]",color="blue",w=1.5,size=(500,400); fonts...)
    plot!(iter, min.(l50[1:iterations],l50[1]), label=L"\quad \kappa^2 \in [0.5,1]",color="red",w=1.5)

    ylabel!("loss")
    xlabel!("epochs")
    savefig("../paper/$(title)")
    savefig("../paper/$(title).eps")
end

path = ""
title = ""
df25 = DataFrame(CSV.File(path))
df50 = DataFrame(CSV.File(path))
kappa_threshold_loss_graph!(title, df25.Train, df50.Train)

function loss_error_residual_to_graph!(title, path)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>bfont)
    df = DataFrame(CSV.File(path))
    iter = range(1, length=length(df.Train))
    p = plot(iter, df.Train, label="total loss",size=(500,400); fonts...)
    plot!(iter, df.Residual, label="residual loss")
    plot!(iter, df.Error, label="error loss")

    ylabel!("loss")
    xlabel!("epochs")
    savefig("../paper/$(title) residual graph")
end

path = ""
title = ""
loss_error_residual_to_graph!(title, path)

# Heatmaps:

function csv_to_heatmap!(title, path, n, m)
    df = DataFrame(CSV.File(path))
    heatmap(reshape(df.K, n-1, m-1), color=:jet,size=(200*m/n,200))
    savefig("../graphs/$(title)")
end

path = ""
csv_to_heatmap!("kappa_7", path, 256, 512)
csv_to_heatmap!("kappa_6", path, 176, 400)
csv_to_heatmap!("kappa_5", path, 128, 256)

# Preconditioner tset graphs:

function unet_vs_vcycle_graph!(title, vu, v, ju; after_vcycle=false, e_vcycle_input=false)
    iterations = 15 #length(v)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>bfont)
    factor = convergence_factor_106!(v[1:iterations])
    factor_text = "v=$(factor)"
    p = plot(iter,v[1:iterations],label=L"\rho_V = " * "$(factor)",legend=:bottomleft,color="green",size=(500,400); fonts...)
    # Jacobi Unet
    factor,_ = convergence_factor_106!(ju[1:iterations])
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(iter,ju[1:iterations],label=L"\rho_{JU} = " * "$(factor)",color="blue",linestyle=:dot)
    # Vcycle Unet
    factor,_ = convergence_factor_106!(vu[1:iterations])
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu[1:iterations],label=L"\rho_{VU} = " * "$(factor)",color="blue")

    yaxis!(L"\Vert b - Hx \Vert_2", :log10)
    xlabel!("iterations")

    savefig("../paper/$(title)")
    savefig("../paper/$(title).eps")
end

path = ""
title = "08_39_27 SDNUnet 0 128"
df = DataFrame(CSV.File(path))
unet_vs_vcycle_graph!("$(title) preconditioner test", df.VU, df.V, df.JU)

function convergence_factor_106!(vector)
    val0 = vector[1]
    index = length(vector)
    for i=2:length(vector)
        if vector[i] < (val0 ./ (10^6))
            index = i
            break
        end
    end
    # length = argmin(vector)[1]
    if index > 200
        return round(((vector[index] / vector[index-30])^(1.0 / 30)), digits=3), index
    else
        return round(((vector[index] / vector[1])^(1.0 / index)), digits=3), index
    end
end

function kappa_threshold_unet_vs_vcycle_graph!(title, vu25, v25, ju25, vu50, v50, ju50; after_vcycle=false, e_vcycle_input=false)
    iterations = length(vc_unet_res)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor = convergence_factor_106!(v25)
    factor_text = "v25=$(factor)"
    p = plot(iter,v25[1:iterations],label=L"\rho_V = " * "$(factor)",legend=:bottomright,color="green", w = 2,size=(500,400); fonts...)
    factor,_ = convergence_factor_106!(ju25)
    factor_text = "$(factor_text) ju25=$(factor)"
    plot!(iter,ju25[1:iterations],label=L"\rho_{JU} = " * "$(factor); " *L"\kappa^2 \in [0.25,1]", w = 2,color="blue",linestyle=:dot)
    factor,_= convergence_factor_106!(vu25)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu25[1:iterations],label=L"\rho_{VU} = " * "$(factor); " * L"\kappa^2 \in [0.25,1]", w = 2,color="blue")
    factor,_ = convergence_factor_106!(ju50)
    factor_text = "$(factor_text) ju25=$(factor)"
    plot!(iter,ju50[1:iterations],label=L"\rho_{JU} = " * "$(factor); " * L"\kappa^2 \in [0.5,1]", w = 2,color="red",linestyle=:dot)
    factor,_ = convergence_factor_106!(vu50)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu50[1:iterations],label=L"\rho_{VU} = " * "$(factor); " * L"\kappa^2 \in [0.5,1]", w = 2,color="red")

    yaxis!("relative residual norm", :log10)
    ylims!((10^-10,10^-1))
    xlabel!("iterations")

    savefig("../paper/$(title)")
    savefig("../paper/$(title).eps")
end

path50 = ""
n = 128
title = "09_40_06_13_48_25_4_$(n)_3"
df25 = DataFrame(CSV.File(path))
df50 = DataFrame(CSV.File(path))
kappa_threshold_unet_vs_vcycle_graph!("$(title)", df25.VU, df25.V, df25.JU, df50.VU, df50.V, df50.JU)

function retrain_unet_vs_vcycle_graph!(title, vu, v, ju, vu_r, v_r, ju_r; after_vcycle=false, e_vcycle_input=false, len=60, v_factor=0.92)
    iterations = len # length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor,_ = convergence_factor_106!(v_r)
    factor_text = "v=$(factor)"
    p = plot(iter,v_r[1:iterations],label=L"\rho_V = " * "$(factor)",legend=:topright, w = 1.8,color="green",size=(500,400); fonts...)
    factor,_ = convergence_factor_106!(ju)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(iter,ju[1:iterations],label=L"\rho_{JU} = " * "$(factor)", w = 1.8,color="blue")
    factor,_ = convergence_factor_106!(ju_r)
    factor_text = "$(factor_text) ju_r=$(factor)"
    plot!(iter,ju_r[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * " re-train", w = 1.8,color="blue",linestyle=:dot)
    factor,_ = convergence_factor_106!(vu)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu[1:iterations],label=L"\rho_{VU} = " * "$(factor)", w = 1.8,color="red")
    factor,_  = convergence_factor_106!(vu_r)
    factor_text = "$(factor_text) vu_r=$(factor)"
    plot!(iter,vu_r[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * " re-train", w = 1.8,color="red",linestyle=:dot)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-10,10^-1))
    xlabel!("iterations")
    savefig("../paper/$(title)")
    savefig("../paper/$(title).eps")
end

n = 128
path = ""
original = DataFrame(CSV.File(path))
retrain = DataFrame(CSV.File(path))
title = "23_48_23_32_bit_cpu_m10r10_128"
retrain_unet_vs_vcycle_graph!("$(title)",original.VU, original.V,original.JU, original.VU_T, original.V_T,original.JU_T;len=100)

function retrain_unet_vs_vcycle_graph_times!(title, vu, v, ju, vu_r, v_r, ju_r, vu_t, v_t, ju_t, vu_r_t, v_r_t, ju_r_t; xlim=20, sec=1, blocks=10)
    iterations = length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor,T = convergence_factor_106!(v_r)
    if T == 1001
        T = 1018
    end
    factor_text = "v=$(factor)"
    p = plot(v_r_t[1:iterations]./(blocks*sec),v_r[1:iterations],label=L"\rho_V = " * "$(factor);" * L" T = " * "$(T)",legend=:topright, w = 1.8,color="green",size=(500,400); fonts...)
    factor,T = convergence_factor_106!(ju)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(ju_t[1:iterations]./(blocks*sec),ju[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * L" T = " * "$(T)", w = 1.8,color="blue")
    factor,T = convergence_factor_106!(ju_r)
    factor_text = "$(factor_text) ju_r=$(factor)"
    plot!(ju_r_t[1:iterations]./(blocks*sec),ju_r[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * L" T = " * "$(T);"* " re-train", w = 1.8,color="blue",linestyle=:dot)
    factor,T= convergence_factor_106!(vu)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(vu_t[1:iterations]./(blocks*sec),vu[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * L" T = " * "$(T)", w = 1.8,color="red")
    factor,T  = convergence_factor_106!(vu_r)
    factor_text = "$(factor_text) vu_r=$(factor)"
    plot!(vu_r_t[1:iterations]./(blocks*sec),vu_r[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * L" T = " * "$(T);" * " re-train", w = 1.8,color="red",linestyle=:dot)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-7,10^-1))
    xlims!((-(xlim ./ 90),xlim))
    if (sec == 1000)
        xlabel!("seconds per RHS")
    else
        xlabel!("milliseconds per RHS")
    end
    savefig("../graphs/times/$(title)")
    savefig("../graphs/times/$(title).eps")
end

function retrain_unet_vs_vcycle_graph_times_details!(title, vu, v, ju, vu_r, v_r, ju_r, vu_t, v_t, ju_t, vu_r_t, v_r_t, ju_r_t; xlim=20, sec=1, blocks=10)
    iterations = length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor,T = convergence_factor_106!(v_r)
    t = T
    if T == 1001
        T = 1018
        t = 1000
    end
    factor_text = "v=$(factor)"
    p = plot(v_r_t[1:iterations]./(blocks*sec),v_r[1:iterations],label=L"\rho_V = " * "$(factor);" * L" T = " * "$(T); $(v_r_t[t]./(blocks*sec))",legend=:topright, w = 1.8,color="green",size=(500,400); fonts...)
    factor,T = convergence_factor_106!(ju)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(ju_t[1:iterations]./(blocks*sec),ju[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * L" T = " * "$(T); $(ju_t[T]./(blocks*sec))", w = 1.8,color="blue")
    factor,T = convergence_factor_106!(ju_r)
    factor_text = "$(factor_text) ju_r=$(factor)"
    plot!(ju_r_t[1:iterations]./(blocks*sec),ju_r[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * L" T = " * "$(T); $(ju_r_t[T]./(blocks*sec));"* " re-train", w = 1.8,color="blue",linestyle=:dot)
    factor,T= convergence_factor_106!(vu)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(vu_t[1:iterations]./(blocks*sec),vu[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * L" T = " * "$(T); $(vu_t[T]./(blocks*sec))", w = 1.8,color="red")
    factor,T  = convergence_factor_106!(vu_r)
    factor_text = "$(factor_text) vu_r=$(factor)"
    plot!(vu_r_t[1:iterations]./(blocks*sec),vu_r[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * L" T = " * "$(T); $(vu_r_t[T]./(blocks*sec));" * " re-train", w = 1.8,color="red",linestyle=:dot)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-7,10^-1))
    xlims!((-(xlim ./ 90),xlim))
    if (sec == 1000)
        xlabel!("seconds per RHS")
    else
        xlabel!("milliseconds per RHS")
    end
    savefig("../graphs/$(title)")
    savefig("../graphs/$(title).eps")
end

function unet_vs_vcycle_graph_iterations!(title, vu, v, ju; after_vcycle=false, e_vcycle_input=false, len=60, v_factor=0.92)
    iterations = len # length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor,_ = convergence_factor_106!(v)
    factor_text = "v=$(factor)"
    p = plot(iter,v[1:iterations],label=L"\rho_V = " * "$(factor)",legend=:topright, w = 1.8,color="green",size=(500,400); fonts...)
    factor,_ = convergence_factor_106!(ju)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(iter,ju[1:iterations],label=L"\rho_{JU} = " * "$(factor)", w = 1.8,color="blue")
    factor,_ = convergence_factor_106!(vu)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu[1:iterations],label=L"\rho_{VU} = " * "$(factor)", w = 1.8,color="red")

    yaxis!("relative residual norm", :log10)
    ylims!((10^-7,10^-1))
    xlabel!("iterations")
    savefig("../graphs/$(title)")
end

function retrain_unet_vs_vcycle_graph_iterations!(title, vu, v, ju, vu_r, v_r, ju_r; xlim=100)
    iterations = length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    factor,_ = convergence_factor_106!(v_r)
    factor_text = "v=$(factor)"
    p = plot(iter,v_r[1:iterations],label=L"\rho_V = " * "$(factor)",legend=:topright, w = 1.8,color="green",size=(500,400); fonts...)
    factor,_ = convergence_factor_106!(ju)
    factor_text = "$(factor_text) ju=$(factor)"
    plot!(iter,ju[1:iterations],label=L"\rho_{JU} = " * "$(factor)", w = 1.8,color="blue")
    factor,_ = convergence_factor_106!(ju_r)
    factor_text = "$(factor_text) ju_r=$(factor)"
    plot!(iter,ju_r[1:iterations],label=L"\rho_{JU} = " * "$(factor);" * " re-train", w = 1.8,color="blue",linestyle=:dot)
    factor,_ = convergence_factor_106!(vu)
    factor_text = "$(factor_text) vu=$(factor)"
    plot!(iter,vu[1:iterations],label=L"\rho_{VU} = " * "$(factor)", w = 1.8,color="red")
    factor,_  = convergence_factor_106!(vu_r)
    factor_text = "$(factor_text) vu_r=$(factor)"
    plot!(iter,vu_r[1:iterations],label=L"\rho_{VU} = " * "$(factor);" * " re-train", w = 1.8,color="red",linestyle=:dot)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-10,10^-1))
    xlims!((-(xlim ./ 90),xlim))
    xlabel!("iterations")
    savefig("../graphs/$(title)")
    savefig("../graphs/$(title).eps")
end

n = 512
blocks = 10
path = ""
df = DataFrame(CSV.File(path))
df_r= DataFrame(CSV.File(path))

title = "23_48_23_64_cpu_block10_model4_$(n)_3090"
retrain_unet_vs_vcycle_graph_iterations!(title, df.VU, df.V, df.JU, df_r.VU, df_r.V, df_r.JU; xlim=400)

# df = (df1.+df2.+df3.+df4.+df5.+df6.+df7.+df8.+df9.+df10)./10
# df_r = (df_r1.+df_r2.+df_r3.+df_r4.+df_r5.+df_r6.+df_r7.+df_r8.+df_r9.+df_r10)./10
title = "23_48_23_32_gpu_model4_$(n)_b$(blocks)_retrain_times_mean10_details"
retrain_unet_vs_vcycle_graph_times!(title, df.VU, df.V, df.JU, df_r.VU, df_r.V, df_r.JU, df.VU_T, df.V_T, df.JU_T, df_r.VU_T, df_r.V_T, df_r.JU_T; xlim=1000, sec=1, blocks=blocks)
retrain_unet_vs_vcycle_graph_times_details!(title, df.VU, df.V, df.JU, df_r.VU, df_r.V, df_r.JU, df.VU_T, df.V_T, df.JU_T, df_r.VU_T, df_r.V_T, df_r.JU_T; xlim=8, sec=1000, blocks=blocks)

blues = palette(:Blues_5)
greens = palette(:Greens_5)
reds = palette(:Reds_5)
function gpu_cpu_32_64_bit_unet_vs_vcycle_graph_times!(title, vu32c, v32c, ju32c, vu32g, v32g, ju32g, vu64c, v64c, ju64c,
    vu32c_t, v32c_t, ju32c_t, vu32g_t, v32g_t, ju32g_t, vu64c_t, v64c_t, ju64c_t; after_vcycle=false, e_vcycle_input=false, len=60, v_factor=0.92)
    iterations = len # length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    p = plot(v64c_t[1:iterations]./1000,v64c[1:iterations],label=L"\rho_V"*" 64 BIT CPU",legend=:topright, w = 1,color=greens[2],size=(800,400), marker=(:circle,5); fonts...)
    plot!(v32c_t[1:iterations]./1000,v32c[1:iterations],label=L"\rho_V"*" 32 BIT CPU",legend=:topright, w = 1.8,color=greens[3])
    plot!(v32g_t[1:iterations]./1000,v32g[1:iterations],label=L"\rho_V"*" 32 BIT GPU",legend=:topright, w = 1.8,color=greens[4])
    # plot!(v64g_t[1:iterations]./1000,v64g[1:iterations],label=L"\rho_V"*" 64 BIT GPU",legend=:topright, w = 1.8,color=greens[5])
    plot!(ju64c_t[1:iterations]./1000,ju64c[1:iterations],label=L"\rho_{JU}"*" 64 BIT CPU",legend=:topright, w = 1.8,color=blues[2])
    plot!(ju32c_t[1:iterations]./1000,ju32c[1:iterations],label=L"\rho_{JU}"*" 32 BIT CPU",legend=:topright, w = 1.8,color=blues[3])
    plot!(ju32g_t[1:iterations]./1000,ju32g[1:iterations],label=L"\rho_{JU}"*" 32 BIT GPU",legend=:topright, w = 1.8,color=blues[4])
    # plot!(ju64g_t[1:iterations]./1000,ju64g[1:iterations],label=L"\rho_{JU}"*" 64 BIT GPU",legend=:topright, w = 1.8,color=blues[5])
    plot!(vu64c_t[1:iterations]./1000,vu64c[1:iterations],label=L"\rho_{VU}"*" 64 BIT CPU",legend=:topright, w = 1.8,color=reds[2])
    plot!(vu32c_t[1:iterations]./1000,vu32c[1:iterations],label=L"\rho_{VU}"*" 32 BIT CPU",legend=:topright, w = 1.8,color=reds[3])
    plot!(vu32g_t[1:iterations]./1000,vu32g[1:iterations],label=L"\rho_{VU}"*" 32 BIT GPU",legend=:topright, w = 1.8,color=reds[4])
    # plot!(vu64g_t[1:iterations]./1000,vu64g[1:iterations],label=L"\rho_{VU}"*" 64 BIT GPU",legend=:topright, w = 1.8,color=reds[5])
    yaxis!("relative residual norm", :log10)
    ylims!((10^-7,10^-1))
    xlabel!("seconds")
    xlabel!("iterations")
    savefig("../graphs/$(title)")
end

function gpu_cpu_32_64_bit_unet_vs_vcycle_graph_iterations!(title, vu32c, v32c, ju32c, vu32g, v32g, ju32g, vu64c, v64c, ju64c, vu64g, v64g, ju64g; after_vcycle=false, e_vcycle_input=false, len=60, v_factor=0.92)
    iterations = len # length(vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    p = plot(iter,v32c[1:iterations],label=L"\rho_V"*" 32 BIT CPU",legend=:topright, w = 1.8,color=greens[2],size=(700,700); fonts...)
    plot!(iter,v32g[1:iterations],label=L"\rho_V"*" 32 BIT GPU",legend=:topright, w = 1.8,color=greens[3])
    plot!(iter,v64c[1:iterations],label=L"\rho_V"*" 64 BIT CPU",legend=:topright, w = 1.8,color=greens[4])
    # plot!(iter,v64g[1:iterations],label=L"\rho_V"*" 64 BIT GPU",legend=:topright, w = 1.8,color=greens[5])
    plot!(iter,ju32c[1:iterations],label=L"\rho_{JU}"*" 32 BIT CPU",legend=:topright, w = 1.8,color=blues[2])
    plot!(iter,ju32g[1:iterations],label=L"\rho_{JU}"*" 32 BIT GPU",legend=:topright, w = 1.8,color=blues[3])
    plot!(iter,ju64c[1:iterations],label=L"\rho_{JU}"*" 64 BIT CPU",legend=:topright, w = 1.8,color=blues[4])
    # plot!(iter,ju64g[1:iterations],label=L"\rho_{JU}"*" 64 BIT GPU",legend=:topright, w = 1.8,color=blues[5])
    plot!(iter,vu32c[1:iterations],label=L"\rho_{VU}"*" 32 BIT CPU",legend=:topright, w = 1.8,color=reds[2])
    plot!(iter,vu32g[1:iterations],label=L"\rho_{VU}"*" 32 BIT GPU",legend=:topright, w = 1.8,color=reds[3])
    plot!(iter,vu64c[1:iterations],label=L"\rho_{VU}"*" 64 BIT CPU",legend=:topright, w = 1.8,color=reds[4])
    # plot!(iter,vu64g[1:iterations],label=L"\rho_{VU}"*" 64 BIT GPU",legend=:topright, w = 1.8,color=reds[5])
    yaxis!("relative residual norm", :log10)
    xlabel!("iterations")
    savefig("../graphs/$(title)")
end

title = "23_48_23_32_64_cpu_gpu_iterations"

path = ""
c32_df = DataFrame(CSV.File(path))
g32_df = DataFrame(CSV.File(path))
c64_df = DataFrame(CSV.File(path))
g64_df = DataFrame(CSV.File(path))

gpu_cpu_32_64_bit_unet_vs_vcycle_graph_iterations!("$(title)",c32_df.VU, c32_df.V,c32_df.JU,g32_df.VU, g32_df.V,g32_df.JU,c64_df.VU, c64_df.V,c64_df.JU,g64_df.VU, g64_df.V,g64_df.JU;len=400)
title = "23_48_23_32_64_cpu_gpu_times"
iter = 800

c32_df = DataFrame(CSV.File(path))
g32_df = DataFrame(CSV.File(path))
c64_df = DataFrame(CSV.File(path))

gpu_cpu_32_64_bit_unet_vs_vcycle_graph_times!("$(title)",c32_df.VU, c32_df.V,c32_df.JU,g32_df.VU, g32_df.V,g32_df.JU,c64_df.VU, c64_df.V,c64_df.JU,
c32_df.VU_T, c32_df.V_T,c32_df.JU_T,g32_df.VU_T, g32_df.V_T,g32_df.JU_T,c64_df.VU_T, c64_df.V_T,c64_df.JU_T;len=iter)

function gmres_unet_vs_vcycle_graph!(title, gmres5vu, gmres5ju, gmres5v, gmres10vu, gmres10ju, gmres10v,gmres20vu, gmres20ju,gmres20v)
    iterations = length(gmres5v)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    p = plot(iter,gmres5v,label=L"M_{V}"*"; GMRES(5)",legend=:outerright, w = 1.8,size=(600,400),color="green"; fonts...)
    plot!(iter,gmres10v,label=L"M_{V}"*"; GMRES(10)",color="green",linestyle=:dash, w = 1.8)
    plot!(iter,gmres20v,label=L"M_{V}"*"; GMRES(20)",color="green",linestyle=:dot, w = 1.8)
    plot!(iter,gmres5ju,label=L"M_{JU}"*"; GMRES(5)",color="blue", w = 1.8)
    plot!(iter,gmres10ju,label=L"M_{JU}"*"; GMRES(10)",color="blue",linestyle=:dash, w = 1.8)
    plot!(iter,gmres20ju,label=L"M_{JU}"*"; GMRES(20)",color="blue",linestyle=:dot, w = 1.8)
    plot!(iter,gmres5vu,label=L"M_{VU}"*"; GMRES(5)",color="red", w = 1.8)
    plot!(iter,gmres10vu,label=L"M_{VU}"*"; GMRES(10)",color="red",linestyle=:dash, w = 1.8)
    plot!(iter,gmres20vu,label=L"M_{VU}"*"; GMRES(20)",color="red",linestyle=:dot, w = 1.8)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-10,10^-1))
    xlabel!("iterations")
    savefig("../graphs/$(title)")
    savefig("../graphs/$(title).eps")
end

n = 256
title = "23_48_23_10blocks_GMRES_$(n)"

path = ""
gmres5 = DataFrame(CSV.File(path))
gmres10 = DataFrame(CSV.File(path))
gmres20 = DataFrame(CSV.File(path))

gmres_unet_vs_vcycle_graph!("$(title)", gmres5.VU, gmres5.JU, gmres5.V, gmres10.VU, gmres10.JU, gmres10.V, gmres20.VU, gmres20.JU, gmres20.V)

function blocks_unet_vs_vcycle_graph!(title, blocks2vu, blocks2ju, blocks2v, blocks5vu, blocks5ju, blocks5v, blocks10vu, blocks10ju, blocks10v, blocks20vu, blocks20ju, blocks20v)
    iterations = length(blocks2vu)
    iter = range(1, length=iterations)
    fonts = Dict(:guidefont=>bfont, :xtickfont=>bfont, :ytickfont=>bfont, :legendfont=>sfont)

    p = plot(iter,blocks2v,label=L"M_{V}"*"; 2 RHS's",legend=:outerright, w = 1.8,color="green",size=(600,400); fonts...)
    plot!(iter,blocks5v,label=L"M_{V}"*"; 5 RHS's",color="green",linestyle=:dash, w = 1.8)
    plot!(iter,blocks10v,label=L"M_{V}"*"; 10 RHS's",color="green",linestyle=:dashdot, w = 1.8)
    plot!(iter,blocks20v,label=L"M_{V}"*"; 20 RHS's",color="green",linestyle=:dot, w = 1.8)
    plot!(iter,blocks2ju,label=L"M_{JU}"*"; 2 RHS's",color="blue", w = 1.8)
    plot!(iter,blocks5ju,label=L"M_{JU}"*"; 5 RHS's",color="blue",linestyle=:dash, w = 1.8)
    plot!(iter,blocks10ju,label=L"M_{JU}"*"; 10 RHS's",color="blue",linestyle=:dashdot, w = 1.8)
    plot!(iter,blocks20ju,label=L"M_{JU}"*"; 20 RHS's",color="blue",linestyle=:dot, w = 1.8)
    plot!(iter,blocks2vu,label=L"M_{VU}"*"; 2 RHS's",color="red", w = 1.8)
    plot!(iter,blocks5vu,label=L"M_{VU}"*"; 5 RHS's",color="red",linestyle=:dash, w = 1.8)
    plot!(iter,blocks10vu,label=L"M_{VU}"*"; 10 RHS's",color="red",linestyle=:dashdot, w = 1.8)
    plot!(iter,blocks20vu,label=L"M_{VU}"*"; 20 RHS's",color="red",linestyle=:dot, w = 1.8)

    yaxis!("relative residual norm", :log10)
    ylims!((10^-10,10^-1))
    xlabel!("iterations")
    savefig("../graphs/$(title)")
    savefig("../graphs/$(title).eps")
end

n = 256
title = "23_48_23_blocks_$(n)"

path = ""
blocks2 = DataFrame(CSV.File(path))
blocks5 = DataFrame(CSV.File(path))
blocks10 = DataFrame(CSV.File(path))
blocks20 = DataFrame(CSV.File(path))

blocks_unet_vs_vcycle_graph!("$(title)", blocks2.VU, blocks2.JU, blocks2.V, blocks5.VU, blocks5.JU, blocks5.V, blocks10.VU, blocks10.JU, blocks10.V, blocks20.VU, blocks20.JU, blocks20.V)


text:
#=
General notation
p   : large prime number
q   : another large prime number
m   : p x q
k   : (p-1) x (q-1)
e   : a large number relativley prime to k
W   : converted message, relativley prime to m
      and less than m

Base equation
eq  : ex - ky = 1

Public shared
e, m
=#

using RSATools

println(m)
println(e)
println(encrypt("tjobalooobaaa"))

#primepair = prime_pair()

#println(primepair[1]*primepair[2])


text:
const N = 10
const GOAL = 10^6

function oneofn{T<:Integer}(n::T)
    0 < n || error("n = ", n, ", but it should be positive.")
    oon = 1
    for i in 2:n
        rand(1:i) == 1 || continue
        oon = i
    end
    return oon
end

nhist = zeros(Int, N)
for i in 1:GOAL
    nhist[oneofn(N)] += 1
end

println("Simulating oneofn(", N, ") ", GOAL, " times:")
for i in 1:N
    println(@sprintf "   %2d => %6d" i nhist[i])
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
const SMOOTH_NONE   = 0
const SMOOTH_LINE_3	= 1
const SMOOTH_LINE_5	= 2
const SMOOTH_QUAD_5	= 3

const MGL_HIST_IN	= 0
const MGL_HIST_SUM	= 1
const MGL_HIST_UP	= 2
const MGL_HIST_DOWN	= 3

const MGL_DRAW_WIRE = 	0	# fastest, no faces
const MGL_DRAW_FAST = 	1	# fast, no color interpolation
const MGL_DRAW_NORM = 	2	# high quality, slower
const MGL_DRAW_LMEM = 	4	# low memory usage (direct to pixel)
const MGL_DRAW_DOTS = 	8	# draw dots instead of primitives
const MGL_DRAW_NONE = 	9	# no ouput (for testing only)

const MGL_DEF_PAL	= "bgrcmyhlnqeupH"	# default palette
const MGL_DEF_SCH	= "BbcyrR"	        # default palette
const MGL_COLORS	= "kwrgbcymhWRGBCYMHlenpquLENPQU"

const MGL_TRANSP_NORM	= 0x000000
const MGL_TRANSP_GLASS  = 0x000001
const MGL_TRANSP_LAMP	= 0x000002
const MGL_ENABLE_CUT	= 0x000004    # Flag which determines how points outside bounding box are drown.
const MGL_ENABLE_RTEXT 	= 0x000008 	# Use text rotation along axis
const MGL_AUTO_FACTOR	= 0x000010 	# Enable autochange PlotFactor
const MGL_ENABLE_ALPHA 	= 0x000020 	# Flag that Alpha is used
const MGL_ENABLE_LIGHT 	= 0x000040 	# Flag of using lightning
const MGL_TICKS_ROTATE 	= 0x000080 	# Allow ticks rotation
const MGL_TICKS_SKIP	= 0x000100 	# Allow ticks rotation

const MGL_DISABLE_SCALE	= 0x000200 	# Temporary flag for disable scaling (used for axis)
const MGL_FINISHED 		= 0x000400 	# Flag that final picture (i.e. mglCanvas::G) is ready
const MGL_USE_GMTIME	= 0x000800 	# Use gmtime instead of localtime
const MGL_SHOW_POS		= 0x001000 	# Switch to show or not mouse click position
const MGL_CLF_ON_UPD	= 0x002000 	# Clear plot before Update()
const MGL_NOSUBTICKS	= 0x004000 	# Disable subticks drawing (for bounding box)
const MGL_DIFFUSIVE	    = 0x008000 	# Use diffusive light instead of specular
const MGL_VECT_FRAME	= 0x010000 	# Use DrwDat to remember all data of frames
const MGL_REDUCEACC		= 0x020000 	# Reduce accuracy of points (to reduc size of output files)
const MGL_PREFERVC 		= 0x040000 	# Prefer vertex color instead of texture if output format supports
const MGL_ONESIDED 		= 0x080000 	# Render only front side of surfaces if output format supports (for debugging)
const MGL_NO_ORIGIN 	= 0x100000 	# Don't draw tick labels at axis origin

const MGL_FONT_BOLD		= 0x01000000	# This value is used binary
const MGL_FONT_ITAL		= 0x02000000	# This value is used binary
const MGL_FONT_BOLD_ITAL= 0x03000000
const MGL_FONT_WIRE	 	= 0x04000000
const MGL_FONT_OLINE	= 0x08000000	# This value is used binary
const MGL_FONT_ULINE	= 0x10000000
const MGL_FONT_ZEROW	= 0x20000000	# internal codes
const MGL_FONT_UPPER	= 0x40000000
const MGL_FONT_LOWER	= 0x80000000
const MGL_FONT_ROMAN	= 0xfcffffff
const MGL_FONT_MASK		= 0x00ffffff
const MGL_COLOR_MASK	= 0xffffff00
const MGL_FONT_STYLE	= 0x3f000000


text:
distance(x, y; dist=Euclidean()) = dist(x, y)
distances(x; dist=Euclidean()) = pairwise(dist, x, dims=2)
distances(x, y; dist=Euclidean()) = pairwise(dist, x, y, dims=2)

function K(rs, x; dist=Euclidean())
    n = size(x, 2)
    d = distances(x; dist)
    N = n * (n - 1)
    cnt = [count(s -> s < r, d) - n for r in rs]
    cnt / N
end

function D(rs, x, maxk; dist=Euclidean())
    n = size(x, 2)
    d = distances(x; dist)

    nr = length(rs)
    cnt = zeros(Int, nr, maxk)

    for (i, r) in enumerate(rs)
        cnt_r = count(s -> s < r, d, dims=1) .- 1
        cnt[i, :] = [count(s -> s ≥ k, cnt_r) for k in 1:maxk]
    end
    cnt / n
end

function T(rs, x; dist=Euclidean())
    n = size(x, 2)
    d = distances(x; dist)

    nr = length(rs)
    cnt = zeros(Int, nr)

    for (i, r) in enumerate(rs)
        q = d .< r
        q[diagind(q)] .= false
        cnt[i] = sum(q^2 .* q)
    end

    N = n * (n - 1) * (n - 2)
    cnt / N
end

min_distance(x, y; dist=Euclidean()) = minimum(pairwise(dist, x, reshape(y, length(y), 1), dims=2))

function H(rs, x::AbstractMatrix, step::Real=rs.step.hi; dist=Euclidean())
    xs = 0.:step:1.
    H(rs, x, xs; dist)
end

function H(rs, x::AbstractMatrix, xs::AbstractRange=rs; dist=Euclidean())
    dim, n = size(x)
    x0 = Iterators.product(fill(xs, dim)...)
    H(rs, x, x0; dist)
end

function H(rs, x::AbstractMatrix, x0; dist=Euclidean())
    n = length(x0)
    mapreduce(y -> rs .>= min_distance(x, collect(y); dist), +, x0) / n
end


text:
struct Mesh <: FEMContainer
    coords::Array{Float64}
    # blocks::Array{Block}
    blocks::Blocks
    node_sets::Array{NodeSet}
    # add side sets and other relevant stuff
    #
end


text:
using Mimi

@defcomp grosseconomy begin
    YGROSS  = Variable(index=[time])    #Gross output
    K       = Variable(index=[time])    #Capital
    l       = Parameter(index=[time])   #Labor
    tfp     = Parameter(index=[time])   #Total factor productivity
    s       = Parameter(index=[time])   #Savings rate
    depk    = Parameter()               #Depreciation rate on capital - Note that it has no time index
    k0      = Parameter()               #Initial level of capital
    share   = Parameter()               #Capital share
end

function timestep(state::grosseconomy, t::Int)
    v = state.Variables
    p = state.Parameters

    #Define an equation for K
    if t == 1
        v.K[t]  = p.k0  #Note the use of v. and p. to distinguish between variables and parameters
    else
        v.K[t]  = (1 - p.depk)^5 * v.K[t-1] + v.YGROSS[t-1] * p.s[t-1] * 5
    end

    #Define an equation for YGROSS
    v.YGROSS[t] = p.tfp[t] * v.K[t]^p.share * p.l[t]^(1-p.share)
end

@defcomp emissions begin
    E       = Variable(index=[time])    #Total greenhouse gas emissions
    sigma   = Parameter(index=[time])   #Emissions output ratio
    YGROSS  = Parameter(index=[time])   #Gross output - Note that YGROSS is now a parameter
end

function timestep(state::emissions, t::Int)
    v = state.Variables
    p = state.Parameters

    #Define an eqation for E
    v.E[t] = p.YGROSS[t] * p.sigma[t]   #Note the p. in front of YGROSS
end

my_model = Model()

setindex(my_model, :time, [2015:5:2110])

addcomponent(my_model, grosseconomy)  #Order matters here. If the emissions component were defined first, the model would not run.
addcomponent(my_model, emissions)

#Set parameters for the grosseconomy component
setparameter(my_model, :grosseconomy, :l, [(1. + 0.015)^t *6404 for t in 1:20])
setparameter(my_model, :grosseconomy, :tfp, [(1 + 0.065)^t * 3.57 for t in 1:20])
setparameter(my_model, :grosseconomy, :s, ones(20).* 0.22)
setparameter(my_model, :grosseconomy, :depk, 0.1)
setparameter(my_model, :grosseconomy, :k0, 130.)
setparameter(my_model, :grosseconomy, :share, 0.3)

#Set parameters for the emissions component
setparameter(my_model, :emissions, :sigma, [(1. - 0.05)^t *0.58 for t in 1:20])
connectparameter(my_model, :emissions, :YGROSS, :grosseconomy, :YGROSS)  #Note that connectparameter was used here.

run(my_model)

#Check model results
my_model[:emissions, :E]


text:
using FredData
using Documenter

DocMeta.setdocmeta!(FredData, :DocTestSetup, :(using FredData); recursive=true)

makedocs(;
    modules=[FredData],
    authors="Micah Smith <micahjsmith@gmail.com>",
    repo="https://github.com/micahjsmith/FredData.jl/blob/{commit}{path}#{line}",
    sitename="FredData.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://micahjsmith.github.io/FredData.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "API Reference" => "api_reference.md",
        "Contributing" => "contributing.md",
        "History" => "history.md"
    ],
)

deploydocs(;
    repo="github.com/micahjsmith/FredData.jl",
    devbranch="master",
)