using Flux

# ResidualBlock is used with `SkipConnection`
# using Parallel, first extract features, then combine them together
# separate filters
struct ResidualBlock
    block
end
Flux.@functor ResidualBlock
(b::ResidualBlock)(x) = x |> b.block .|> leakyrelu


function basic_block()
    layer = Chain(Conv((3, 3, 3), 16 => 16, leakyrelu; stride = 1, pad = 1),
                  Conv((3, 3, 3), 16 => 16; stride = 1, pad = 1))
    return ResidualBlock(SkipConnection(layer, +))
end

function residual_block()
    layer = Chain(Conv((1, 1, 1), 8 => 16, leakyrelu; stride = 1),
                  basic_block(),
                  Conv((1, 1, 1), 16 => 8; stride = 1))
    return ResidualBlock(SkipConnection(layer, +))
end

function residual_net()

    init_conv = Conv((1, 1, 1), 1 => 8, leakyrelu; stride = 1)
    resblock1 = residual_block()
    resblock2 = residual_block()
    resblock3 = residual_block()
    final_conv = Conv((1, 1, 1), 8 => 1; stride = 1)

    return Chain(init_conv, resblock1, resblock2, resblock3, final_conv)
end

# final_conv = Conv((1, 1, 1), 16 => 2; stride = 1)
# basic_conv = Chain(Conv((3, 3, 3), 16 => 16, leakyrelu; stride = 1, pad = 1),
#               Conv((3, 3, 3), 16 => 16; stride = 1, pad = 1))
#
# function mymodel()
#         return Chain(Parallel((x, y) -> cat(x, y, dims = 4),
#                      residual_net(),
#                      residual_net(),
#                      ),
#                      basic_block(),
#                      basic_block(),
#                      basic_block(),
#                      final_conv
#                      )
# end

# xx = randn(Float32, 16, 16, 16, 2, 1)
# xtrue = ones(Float32, 16, 16, 16, 2, 1)
# model = residual_net()
# model(xx)
# ps = Flux.params(model)
# opt = ADAM(0.001)
# loss(x, y) = Flux.Losses.mse(model(x), y)
# evalcb() = @show(loss(xx, xtrue))
# data = [(xx, xtrue)]
# @epochs 500 Flux.train!(loss, ps, data, opt, cb = evalcb)
