using Knet, Test, Base.Iterators, Printf, LinearAlgebra, Random,  IterTools, CUDA
Knet.atype() = KnetArray{Float32}
include("/scratch/users/bberkay14/workfolder/language_to_logic/seq2tree/seq2tree/atis/tree.jl")
include("main.jl")
include("/scratch/users/bberkay14/workfolder/language_to_logic/seq2tree/seq2tree/atis/util.jl")
using CUDA
Knet.atype() = KnetArray{Float32}
using JLD2

function convert_to_string(idx_list, form_manager)
    w_list = []
    for i in range(1, stop=length(idx_list))
        push!(w_list, form_manager.get_idx_symbol(Int(idx_list[i])))
    end
    return join(w_list, " ")
end


function do_generate( l::L2L,  enc_w_list, word_manager, form_manager, opt,  checkpoint)
    # initialize the rnn state to all zeros
    
    prev_c  = torch.zeros((1, encoder.hidden_size), requires_grad=False)
    prev_h  = torch.zeros((1, encoder.hidden_size), requires_grad=False)
    prev_c  = param(1, enc_max_len; init=zeros)
    prev_h  = param(1, enc_max_len; init=zeros)

    # reversed order
    push!(enc_w_list, get_symbol_idx(word_manager,"<S>"))
    insert!(enc_w_list, 1, get_symbol_idx(word_manager, "<E>"))
    end_w = length(enc_w_list)
    enc_outputs = torch.zeros((1, end_w, encoder.hidden_size), requires_grad=False)
    enc_outputs = param(1, end_w, checkpoint["opt"]["rnn_size"] ; init=zeros)
    
    for i in range(end_w, 1, step=-1)

        cur_input = Knet.KnetArray(enc_w_list[i])

        prev_c, prev_h = encode(l, cur_input, prev_c, prev_h)
        enc_outputs[:, i, :] = prev_h
    end
    # decode
    queue_decode = []
    #queue_decode.append({"s": (prev_c, prev_h), "parent":0, "child_index":1, "t": Tree()})
    push!(queue_decode, Dict("s"=> (prev_c, prev_h), "parent"=>0, "child_index"=>1, "t"=> Tree()))    
    
    head = 1
    while head <= size(queue_decode,1) and head <=100
        s = queue_decode[head-1]["s"]
        parent_h = s[1]
        t = queue_decode[head-1]["t"]
        if head == 1
            prev_word = KnetArray(get_symbol_idx(form_manager, "<S>"))
            
        else
            prev_word = KnetArray(get_symbol_idx(form_manager, "("))
        end

        i_child = 1
        while True
            curr_c, curr_h = decode(l, prev_word,  parent_h)
            prediction = attention_decode(l, enc_outputs, curr_h)
            
            s = (curr_c, curr_h)
            _prev_word = maximum(prediction, dims=2)
            
            prev_word = _prev_word
            if Int.(prev_word[1]) == get_symbol_idx(form_manager, "<E>") || t.num_children >= checkpoint["opt"]["dec_seq_length"]
                break
            elseif Int64.(prev_word[1]) == get_symbol_idx(form_manager, "<N>")
                #print("we predicted N");exit()
                push!(queue_decode, Dict("s"=> (deepcopy(s[1]), deepcopy(s[2])), "parent"=> head, "child_index"=>i_child, "t"=> Tree()))
                #deepcopy(s[1])
                #push!(queue_decode, Dict("s" => (s[0].clone(), s[1].clone()), "parent": head, "child_index":i_child, "t": Tree()))
                add_child(t, Int.(prev_word[1]))
            else
                add_child(t, Int.(prev_word[1]))
            end
            i_child = i_child + 1
        end
        head = head + 1
    # refine the root tree (TODO, what is this doing?)
    end
    for i in range(size(queue_decode,1), stop=1, step=-1)
        cur = queue_decode[i]
        queue_decode[cur["parent"]]["t"].children[cur["child_index"]] = cur["t"]
    end
    return to_list(queue_decode[1]["t"], form_manager)
end



s = ArgParseSettings()
@add_arg_table m begin
    "--gpuid"
        arg_type = Int
        help = "which gpu to use. -1 = use CPU"
        default = 0
    "--temperature"
        help = "temperature of sampling"
        arg_type = Int
        default = 0
    "--sample"
        help = "0 to use max at each timestep (-beam_size=1), 1 to sample at each timestep, 2 to beam search"
        arg_type = Int
        default = 0
    "--beam_size"
        help = "beam size"
        arg_type = Int
        default = 20
    "--display"
        help = "whether display on console"
        arg_type = Int
        default = 1
    "--data_dir"
        help = "data path"
        arg_type = String
        default = "../data/"
    "-input"
        help = "input data filename"
        arg_type = String
        default = "test.t7"
    "--output"
        help =  "input data filename"
        arg_type = String
        default = "output/seq2seq_output.txt"
    "--model"
        help = "model checkpoint to use for sampling"
        arg_type = String
        default = "checkpoint_dir/model_seq2seq"
    "-seed"
        help = "torch manual random number generator seed"
        arg_type = Int
        default = 123
end


args = parse_args(s)


# load the model checkpoint
checkpoint = Knet.load(args["model"])
l_model = checkpoint["lang2logic_model"]



# initialize the vocabulary manager to display text
#managers = pkl.load( open("{}/map.pkl".format(args.data_dir), "rb" ) )
@load  (string(checkpoint["data_dir"]) * "/map.jld2") managers
word_manager, form_manager = managers
# load data
#data = pkl.load(open("{}/test.pkl".format(args.data_dir), "rb"))
@load  (string(checkpoint["data_dir"]) * "/test.jld2") managers
reference_list = []
candidate_list = []
#with open(args.output, "w") as output:
open(args["output"], "w") do output
    # TODO change when running full -- this is to just reproduce the error
    #for i in range(30,50):
    #for i in range(278,280):
    for i in range(1, stop=length(data))
        #print("example {}\n".format(i))
        x = data[i]
        reference = x[2]
        candidate = do_generate(l_model, x[1], word_manager, form_manager, args, checkpoint)
        candidate = [Int(c) for c in candidate]


        num_left_paren = sum(1 for c in candidate if form_manager.idx2symbol[Int(c)] == "(")
            
        num_right_paren = sum(1 for c in candidate if form_manager.idx2symbol[Int(c)]== ")")
        diff = num_left_paren - num_right_paren
        #print(diff)
        if diff > 0
            for i in range(1, stop= diff)
                push!(candidate, form_manager.symbol2idx[")"])
            end
        elseif diff < 0
            candidate = candidate[1:diff]
        end
        ref_str = convert_to_string(reference, form_manager)
        cand_str = convert_to_string(candidate, form_manager)

        push!(reference_list, reference)
        push!(candidate_list, candidate)
        # print to console
        if args["display"] > 0
            print("results: ")
            print(ref_str)
            print(cand_str)
            print(' ')
        #output.write("{}\n".format(ref_str))
        write(output, string(ref_str) *"\n")                    
        #output.write("{}\n".format(cand_str))
        write(output, string(cand_str) *"\n") 
    end
    val_acc = compute_tree_accuracy(candidate_list, reference_list, form_manager)
    print("ACCURACY = " * string(val_acc) * "\n")
    #output.write("ACCURACY = {}\n".format(val_acc))
    write(output,"ACCURACY = " * string(val_acc) *"\n") 
    
end
