using CuArrays
using CUDAnative
using LinearAlgebra

# ::: STRUCTURE

# GPU
GlobalArrayType = CuArray
# CPU
# GlobalArrayType = Array

mutable struct NetworkNode
    forward::Function
    apply_derivs!::Function
    params::Dict{String, Any}
    derivs::Dict{String, Function}
end

struct Optimizer
    f::Function
    params::Dict{String, Any}
end

# :: HELPER FUNCTIONS
include("helper_functions.jl")


# ::: LAYERS
DenseLayer(input_size::Int, output_size::Int; norm = (f,_)->f, w_init = rand_glorot, b_init = CuArrays.zeros) = begin
    w = w_init(input_size, output_size)
    b = b_init(input_size, output_size)

    p = Dict(
        "w" => w,
        "b" => b
    )

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        d = Dict("dydx" => (Δ) -> p["w"]' * Δ,
            "dydw" => norm((Δ) -> Δ * x', Dict("params" => p["w"])),
            "dydb" => norm((Δ) -> Δ, Dict("params" => p["b"])))

        return (p["w"] * x .+ p["b"], d)
    end

    function apply_derivs!(n, Δ, opt)::Nothing
        n.params["w"] .-= apply!(opt, Δ, n.derivs["dydw"])
        n.params["b"] .-= apply!(opt, Δ, n.derivs["dydb"])
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Δp)
end
ActivationLayer(input_size::Int, f::Function, df::Function) = begin
    p = Dict()
    Δp = Dict("dydx" => (Δ) -> Δ)

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        return (f(x),
            Dict("dydx" => (Δ) -> df(x) * Δ)
        )
    end

    function apply_derivs!(_, _, _)::Nothing
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end
BatchNormLayer(in_out_size, batch_size; α_init = rand_glorot, β_init = CuArrays.zeros) = begin
    p = Dict(
        "α" => α_init(in_out_size, batch_size)
        "β" => β_init(in_out_size, batch_size)
        "ϵ" => 0.00001 # to prevent devision by zero
    )

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        μ = sum(x) / length(x)
        σ = sum((x .- μ) .^ 2) / length(x)

        bn = (x .- μ) ./ sqrt(σ ^ 2 + p["ϵ"])
        y = bn .* p["α"] .+ p["β"]

        return (p["w"] * x .+ p["b"],
                Dict("dydx" => (Δ) -> Δ .* p["α"] .* sqrt(σ + p["ϵ"]),
                    "dydα" => (Δ) -> Δ .* bn,
                    "dydβ" => (Δ) -> Δ)
                )
    end

    function apply_derivs!(n, Δ, opt)::Nothing
        n.params["α"] .-= apply!(opt, Δ, n.derivs["dydα"])
        n.params["β"] .-= apply!(opt, Δ, n.derivs["dydβ"])
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end
ConvolutionalLayer(input_size, out_channels, k_size; padding=(0,0,0), k_stride=(1,1,1), k_dilation=(1,1,1); w_init = rand_glorot, b_init = CuArrays.zeros) = begin


    function forward(x::GlobalArrayType, p::Dict)
        x = pad(x, P)


    end

    function apply_derivs!(Δ)
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end

RecurrentLayer(input_size, hidden_size, max_buffer_length) = begin
    p = Dict(
        "xh_w"
        "xh_b"
        "hh_w"
        "hh_b"
        "hy_w"
        "hy_b"
    )

    function forward(x::GlobalArrayType, p::Dict)

        pop!(a)
        pushfirst!(a, 4)
    end

    function apply_derivs!(Δ)
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end



# in order (height, width, channels, batch)
is = (11, 13, 20, 20)
P = (2, 2, 0, 0)
K = (3, 3, 3, 20)
S = (2,2,2,1)
D = (2,2,2,1)



# conv 1: get estimated input size
is_p = is .+ P

# conv 2: get the valid output size estimate
total_KS = K .+ (K .- 1) .* (D .- 1)
max_ind = is_p .- total_KS .+ 1
available_indecies = ((max_ind .* (max_ind .> 0)) ./ S) .|> ceil .|> Int
all_alpha_i = get_valid_indecies(is_p, available_indecies, S)


# conv 3: get flat kernel indecies from any 1D location of 4D tensor from arbitrary location
all_beta_i = flat_kernel_indecies(is_p, K, D)


# conv 4: get all indecies
all_x_indecies = [i .+ all_beta_i for i in all_alpha_i] |> Iterators.flatten |> collect
final_x_ind = reshape(all_x_indecies, (prod(K[1:3]), K[4], length(all_alpha_i)))


function index_assigne!(out, X, I)
    I_ind = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    i = I[I_ind]
    out[I_ind] = X[i]

    return nothing
end

x = CuArray(rand(is_p...))
out = CuArray(similar(final_x_ind))

@cuda threads=length(all_beta_i) blocks=length(all_alpha_i) index_assigne!(out, x, final_x_ind |> CuArray)


# TODO:
# index_assigne => tmp_out
#   => reshape(prod(K[1:3]), K[end], number_of_kernels)
#   => permutedims([prod(K[1:3], #_kernels, K[end])])
#   => multiply pointwise with weights [#_output_channels, prod(K[1:3]), #_kernels, K[end]]
#   => sum(dims=2)
#   => reshape into [#_output_channels, #_kernels, K[end]]
#   => add biases [#_output_channels, #_kernels, K[end]]
#   => reshape into [#_output_size..., #_output_channels, K[end]]





# ::: NORMS
no_norm(df::Function, _) = df
L2_norm(df::Function, p::Dict; λ = 0.005) = (x) -> df(x) .+ (λ .* p["params"])


# ::: OPTIMIZERS
SGD(η) = Optimizer((d, df, p) -> df(d) .* p["lr"], Dict("lr" => η))


# ::: ACTIVATION FUNCTIONS
relu(A::GlobalArrayType) = map(a->a * (a>0), A)
Δrelu(A::GlobalArrayType) =  GlobalArrayType{Float32}(I, length(A), length(A)) .* map(a->(a>0), A)

sigmoid(A::GlobalArrayType) = 1 ./ (1 .+ exp.(.-A))
Δsigmoid(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (sigmoid(A) .* (1 .- sigmoid(A)))

diminish(A::GlobalArrayType) = (2 .* A) ./ (1 .+ A .^ 2)
Δdiminish(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (2 .* ((1 .- A .^ 2) ./ (1 .+ 2 .* A .^ 2 .+ A .^ 4)))
diminish(A) = (2 .* A) ./ (1 .+ A .^ 2)

softmax(A::GlobalArrayType) = exp.(A) ./ sum(exp.(A))
function d_softmax(A::GlobalArrayType)
    f1 = (.-softmax(A) * softmax(A)') .* (CuArrays.fill(1, length(A), length(A)) .- GlobalArrayType{Float32}(I, length(A), length(A)))
    f2 = (softmax(A) .* (1 .- softmax(A))) .* GlobalArrayType{Float32}(I, length(A), length(A))
    return f1 .+ f2
end

# ::: LOSS FUNCTIONS
crossentropy(p::Array, q::Array; base=ℯ) = -sum(q .* log.(base, p))
d_crossentropy(p::Array, q::Array; base=ℯ) = .-(q .* (1 ./ (p .* log(base))) .* log.(base, p))

mse(x::GlobalArrayType, y::GlobalArrayType) = sum((x .- y).^2) / length(x)
Δmse(x::GlobalArrayType, y::GlobalArrayType) = 2 .* (x .- y) ./ length(x)

diminish(A::GlobalArrayType) = (2 .* A) ./ (1 .+ A .^ 2)
Δdiminish(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (2 .* ((1 .- A .^ 2) ./ (1 .+ 2 .* A .^ 2 .+ A .^ 4)))
