#=
This file contains ParamCtx, Param and KnetModule abstractions.
=#

"""Data structure used to store parameters. Currently, it is a 1d Any array"""
const ParamCtx = Array{Any, 1}

# The context abstraction

let
    global active_ctx, switch_ctx!, default_ctx, reset_ctx!
    
    #=default = ParamCtx()
    active = default=#
    active = ParamCtx()
    
    """`active_ctx()::ParamCtx` returns the context in use"""
    active_ctx()::ParamCtx = active

    """`switch_ctx!(new)::ParamCtx` switches to a new context and returns it"""
    switch_ctx!(new::ParamCtx)::ParamCtx = (active=new)

    #="""`default_ctx()::ParamCtx` returns the default context provided by the system"""
    default_ctx()::ParamCtx = default=#

    """`reset_ctx!()::ParamCtx` removes everything from the context"""
    reset_ctx!()::ParamCtx = (active=ParamCtx())
end

"""
Shorthand form of `active_ctx()`
"""
actx() = active_ctx()

"""`Param` stores an address to the parameters value.
Trainable arrays should be stored as `Param` inside modules.
User is responsible to keep track of `ParamCtx` object the `Param`
belongs to.

# Constructors
    `Param(w; ctx=active_ctx())` registers a parameter to the context ctx.
"""
type Param
    index::Integer
    take_grad::Bool
    Param(w; ctx::ParamCtx=active_ctx(), take_grad=true) =
        new(length(push!(ctx, w)), take_grad)
end


"""
`val(ctx, p::Union{Param, Void})` returns the value of parameter in
ctx. 
This function should be used used to access parameters for AutoGrad to work.
"""
val(ctx, p::Union{Param, Void}) =
    p.take_grad ? ctx[p.index] : getval(ctx[p.index])


"""
`aval(p::Param)` returns value of p in the current `active_ctx()`. 
It will stop differentiation when used inside forward funtions.
"""
aval(p::Param) = val(active_ctx(), p)


"""
`setval!(p::Param, w; ctx=active_ctx())` sets the value of `p` to `w` in `ctx`.
"""
setval!(p::Param, w; ctx=active_ctx()) = (ctx[p.index] = w)


"""`KnetModule` is the abstract type inherited by other modules."""
abstract type KnetModule end

import AutoGrad.grad

"""
`grad(m::KnetModule, loss)` returns a grad function of the form `(args..., y)->g`.
Recorded parameters are lived in `active_ctx`.

 # Arguments

    `m`: A KnetModule instance

    `loss`: A function of the form `(ypred, ygold)->scalar_loss`.
"""
function grad(m::KnetModule, loss::Function)
    _predict(w, x...; o...) = m(w, x...; o...)
    _loss(w, args...; kwargs...) = loss(_predict(w, args[1:end-1]...; kwargs...),
                                        args[end])
    lossgrad = grad(_loss)
    return (args...; kwargs...)->lossgrad(active_ctx(), args...; kwargs...)
end

import AutoGrad.gradloss
"""
`gradloss(m::Module, loss::Function)` 
Same as `grad`, with only difference of returning the loss
"""
function gradloss(m::KnetModule, loss::Function)
    _predict(w, x...; o...) = m(w, x...; o...)
    _loss(w, args...; kwargs...) = loss(_predict(w, args[1:end-1]...; kwargs...),
                                        args[end])
    lossgrad = gradloss(_loss)
    return (args...; kwargs...)->lossgrad(active_ctx(), args...; kwargs...)
end

"""
`getgrad(p::Param, grads)` returns the corresponding gradient of `p` in
the grad array returned by any gradfn (generated by `grad` functions of Knet
or KnetModules).
"""
getgrad(p::Param, grads) = grads[p.index]


import Knet.optimizers

"""`optimizers(m::KnetModule, otype;o...)` 
Creates a group of optimizers that correspond to parameters in module m. 
otype specifies a Knet optimizer and `o...` is its options.
"""
optimizers(m::KnetModule, otype; o...) =
    Array{Any, 1}(map(_->otype(;o...), params(m)))


import Knet.update!
"""
`update!(m::KnetModule, grads, optims)`

Applies optimizers to parameters of `m`, using gradients stored in `grads`.
`optims` should have returned by `KnetModules.optimizers` for consistent behaviour.
"""
function update!(m::KnetModule, grads, optims)
    for (p, o) in zip(params(m), optims)
        update!(aval(p), getgrad(p, grads), o)
    end
end


function _populate_recursive(m, list, match_type)
    if isa(m, match_type)
        push!(list, m)
    end
    if isa(m, KnetModule)
        for fn in fieldnames(m)
            _populate_recursive(getfield(m, fn), list, match_type)
        end
    elseif isa(m, Array) || isa(m, Tuple)
        etype = eltype(m)
        # TODO: expand this filter
        if ~(etype <: Number || etype<:AbstractString)
            for l in m
                _populate_recursive(l, list, match_type)
            end
        end
    elseif isa(m, Associative)
        for k in keys(m)
            _populate_recursive(m[k], list, match_type)
        end
    else
        return
    end
end


"""
`params(m::KnetModule; kwargs)` returns parameters of the KnetModule m.
These include any value of type `Param` stored in fields, sub-modules, 
dictionaries, arrays and tuples. Parameters are sorted based on their
locations in their respective context, so they are always returned in
the same order.
"""
function params(m::KnetModule)
    res = []
    _populate_recursive(m, res, Param)
    # For determinism (not solely for duplicate removal)
    sort!(res; lt=(r1, r2)->r1.index < r2.index)
    # Remove duplicates from the list
    #return res
    isempty(res) && return res
    res2 = Any[res[1]]
    for i = 2:length(res)
        (res[i-1].index != res[i].index) && push!(res2, res[i])
    end
    return res2
end


"""
`submodules(m::KnetModule)` returns list of modules contained in s.
These include any value of type `Param` stored in fields, sub-modules, 
dictionaries, arrays and tuples. 
"""
function modules(m::KnetModule)
    res = Set{KnetModule}()
    _populate_recursive(m, res, KnetModule)
    return Any[r for r in res]
end

"""
`convert_buffers!(m::KnetModule, atype)` 
Creates a method of this function to convert non-parameter 
state to `atype`. 
You may create a method of this function to change cpu! and
gpu! behaviour of a particular module type.
"""
function convert_buffers!(m::KnetModule, atype)
    nothing
end

"""
`convert_params!(m::KnetModule, atype)` 
Creates a method of this function to convert parameters 
to `atype`.
You may create a method of this function to change `cpu!` and
`gpu!` behaviour of a particular module type.
"""
function convert_params!(m::KnetModule, atype)
    for p in params(m)
        setval!(p, atype(aval(p)))
    end
end


"""
`gpu!(m::KnetModule)` transfers all parameters in `m` to supported
gpu devices, which is identical to converting their valus to `KnetArray`.
If no gpu is avaliable, a warning is raised and transfer is ignored without
any error.
"""
function gpu!(m::KnetModule)
    if gpu() < 0
        warn("Gpu transfer is ignored, no available gpu")
        return m
    end
    for sm in modules(m)
        convert_params!(sm, ka)
        convert_buffers!(sm, ka)
    end
    return m
end


"""`cpu!(m::KnetModule)` transfers all parameters in `m` cpu, 
which is identical to converting their values to `Array`"""
function cpu!(m::KnetModule)
    for sm in modules(m)
        convert_params!(sm, Array)
        convert_buffers!(sm, Array)
    end
    return m
end


"""
`training!(m::KnetModule)` makes the fields called `train` true if exists, 
which may effect the execution mode in some modules. It also make `take_grad`
field of parameters true, which make them included in the training process.
"""
function training!(m::KnetModule)
    for sm in modules(m)
        if :train in fieldnames(sm)
            sm.train = true
        end
        for p in params(sm)
            p.take_grad = true
        end
    end
    return m
end


"""
`testing!(m::KnetModule)` makes the fields called `train` `false` if exists, 
and freezes the parameters so that gradient is not taken w.r.t them. This
is done by unrecording the parameters during training, so that particular modules`
weights can be treated as constants by the AutoGrad. 
"""
function testing!(m::KnetModule)
    for sm in modules(m)
        if :train in fieldnames(sm)
            sm.train = false
        end
        for p in params(sm)
            p.take_grad = false
        end
    end
    return m
end

# Macros for simple module operations
# TODO: support kwargs
#="""
`@mc expr` Converts the expression `m(a...)` to `forward(ctx, m, a...)`. 
`ctx` should come from the upper scope.

# Example
    l = Linear(...)
    @mc l(x)
"""
macro mc(expr)
    if typeof(expr) == Expr && expr.head == :call
        return esc(:(forward(ctx, $(expr.args[1]),
                             $(expr.args[2:end]...))))
    end
    return expr
end

"""
`@run expr` Converts the expression m(a...) to forward(active_ctx(), m, a...).

# Example
    l = Linear(...)
    @run l(x)
"""
macro run(expr)
    if typeof(expr) == Expr && expr.head == :call
        return esc(:(forward(active_ctx(),
                             $(expr.args[1]),
                             $(expr.args[2:end]...))))
    end
    return expr
end

# Legacy
forward(ctx, m::KnetModule, args...) =
    m(ctx, args...)=#
    #=error("This forward signature is not implemented",
          "for abstract types and/or type ", typeof(m))=#


#=macro callable(Module, fn=forward)
    esc(:((m::$(Module))(ctx, args...; kwargs...) =
          $(fn)(ctx, m, args...; kwargs...)))
end=#


"""
`clean_ctx!(m::KnetModule)` 
Switches to a new ctx where only parameters of `m` live.
"""
function clean_ctx!(m::KnetModule)
    is_ctx_clean(m) && return m
    new = ParamCtx()
    for p in params(m)
        p.index = length(push!(new, aval(p)))
    end
    switch_ctx!(new)
    return m
end


function is_ctx_clean(m::KnetModule)
    ps = Set{Int}(map(x->x.index, params(m)))
    as = Set{Int}(1:length(active_ctx()))
    return length(setdiff(as, ps)) == 0
end

function ctx_dict(m::KnetModule)
    cd = Dict()
    for p in params(m)
        cd[p.index] = aval(p)
    end
    return cd
end

# This thing is used with a loaded (not initialized) Knet
# module.
function from_ctx_dict!(m::KnetModule, cd::Dict; ctx=nothing)::ParamCtx
    if ctx == nothing; ctx = ParamCtx(); end
    for p in params(m)
        # FIXME: try not to break the abstraction
        p.index = length(push!(ctx, cd[p.index]))
    end
    return ctx
end


if Pkg.installed("JLD") !== nothing
    import JLD
else
    JLD = nothing
end

@inline assert_jld(fnname="this") =
    @assert (JLD !== nothing) "Install JLD for using $fname function"

function save_module(filename::String, m::KnetModule,
                     other_state::Dict=Dict())
    assert_jld("save_module")
    args = ["model", m, "ctx_dict", ctx_dict(m)]
    for k in keys(other_state)
        push!(args, k, other_state[k])
    end
    JLD.save(filename, args...)
end

function load_module(filename::String; ctx_switch=false)
    assert_jld("load_module")
    loaded = JLD.load(filename)
    model = loaded["model"]
    cdict = loaded["ctx_dict"]
    if ctx_switch
        switch_ctx!(from_ctx_dict!(model, cdict))
    else
        from_ctx_dict!(model, cdict; ctx=active_ctx())
    end
    if length(loaded) > 2
        other_state = loaded
        delete!(other_state, "model")
        delete!(other_state, "ctx_dict")
        return model, other_state
    end
    return model
end
