using Flux
using CuArrays

export NeuralNet

# This code is originally written by Avik Pal
struct ResidualBlock
    conv_layers
    norm_layers
    shortcut
end

Flux.@treelike ResidualBlock

# ResidualBlock Function allows us to define a Residual Block having any number of Convolution and Batch Normalization Layers
function ResidualBlock(filters, kernels::Array{Tuple{Int,Int}},
                       pads::Array{Tuple{Int,Int}},
                       strides::Array{Tuple{Int,Int}},
                       shortcut = identity)
    local conv_layers = []
    local norm_layers = []
    for i in 2:length(filters)
        @inbounds push!(conv_layers, Conv(kernels[i-1], filters[i-1]=>filters[i], pad = pads[i-1], stride = strides[i-1]) |> gpu)
        push!(norm_layers, BatchNorm(filters[i])|>gpu)
    end
    ResidualBlock(Tuple(conv_layers),Tuple(norm_layers),shortcut)
end

# Function converts the Array of scalar kernel, pad and stride values to tuples
function ResidualBlock(filters, kernels::Array{Int}, pads::Array{Int}, strides::Array{Int}, shortcut = identity)
    ResidualBlock(filters, [(i,i) for i in kernels], [(i,i) for i in pads], [(i,i) for i in strides], shortcut)
end

function (block::ResidualBlock)(input)
    local value = copy.(input)
    for i in 1:length(block.conv_layers)-1
        @inbounds value = relu.((block.norm_layers[i])((block.conv_layers[i])(value)))
    end
    @inbounds relu.(((block.norm_layers[end])((block.conv_layers[end])(value))) + block.shortcut(input))
end


struct NeuralNet
    base_net::Chain
    value::Chain
    policy::Chain
end

Flux.@treelike NeuralNet

function NeuralNet(N::Int = 9, tower_height::Int = 20, features::Int=256)

    res_block() = ResidualBlock([features,features,features], [3,3], [1,1], [1,1])
    # 9 residual blocks
    tower = [res_block() for i = 1:tower_height]

    # feature extraction
    base_net = Chain(Conv((3,3), 9=>features, pad=(1,1)),
                     BatchNorm(features, leakyrelu),
                     tower...) |> gpu

    # value head
    value = Chain(Conv((1,1), features=>1),
                  BatchNorm(1, relu), x->reshape(x, N^2, :),
                  Dense(N^2, features, leakyrelu),
                  Dense(features, 1, tanh)) |> gpu

    policy = Chain(Conv((1,1), features=>2),
                   BatchNorm(2, relu), x->reshape(x, 2N^2, size(x, 4)),
                   Dense(2N^2, 9), softmax) |> gpu

    NeuralNet(base_net, value, policy)
end

function Base.show(io::IO, nn::NeuralNet)
    compact = get(io, :compact, false)
    if !compact
        print(io, "NeuralNet()")
    end
end


function (nn::NeuralNet)(input::AbstractArray{Float32,4})
    features = nn.base_net(input)
    p = nn.policy(features)
    v = nn.value(features)
    return p, v
end
