text:
include("../../src/unet/utils.jl")

function BatchNormWrap(out_ch)
    Chain(x->expand_dims(x,2)|>cgpu,
	  BatchNorm(out_ch)|> cgpu,
	  x->squeeze(x))|> cgpu
end

ConvDown(in_chs,out_chs;kernel = (5,5), σ=elu) =
    Chain(Conv(kernel, in_chs=>out_chs, stride=(2,2), pad = 1; init=_random_normal),
	    BatchNorm(out_chs),
	    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu

ConvUp(in_chs,out_chs;kernel = (5,5), σ=elu) =
    Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
        ConvTranspose(kernel, in_chs=>out_chs, stride=(2, 2), pad = 1; init=_random_normal),
        BatchNorm(out_chs))|> cgpu

struct UNetUpBlock
  upsample
end

(u::UNetUpBlock)(input, bridge) = cat(u.upsample(input), bridge, dims = 3)

@functor UNetUpBlock

UNetUpBlock(in_chs::Int, out_chs::Int; kernel = (5, 5), p = 0.5f0, σ=elu) =
    UNetUpBlock(Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    ConvTranspose(kernel, in_chs=>out_chs, stride=(2, 2), pad = 1; init=_random_normal),
                    BatchNorm(out_chs)))|> cgpu

UNetConvBlock(in_chs, out_chs; kernel = (3, 3), pad=1, σ=elu) =
    Chain(Conv(kernel, in_chs=>out_chs, pad=pad; init=_random_normal),
                BatchNorm(out_chs),
                x-> (σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu

# ResNet Types:

# Preservation of input with with a 1-kernel convolution

struct ResidualBlock
  layers
  shortcut
  bn
end

(r::ResidualBlock)(input) = r.bn(r.layers(input) + r.shortcut(input))

@functor ResidualBlock

function ResidualBlock(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>out_chs,pad = pad;init=_random_normal),
                	BatchNorm(out_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, out_chs=>out_chs,pad = pad; init=_random_normal))|> cgpu
    shortcut = Conv((1,1),in_chs=>out_chs)|> cgpu
    bn = BatchNorm(out_chs)|> cgpu
    ResidualBlock(layers, shortcut, bn)|> cgpu
end

# Preservation of input without convolution

struct SResidualBlock
  layers
  activation
end

(r::SResidualBlock)(input) = r.activation(r.layers(input) + input)

@functor SResidualBlock

function SResidualBlock(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal),
                	BatchNorm(in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, in_chs=>out_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(out_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    SResidualBlock(layers, activation)|> cgpu
end

# Doubling the channels and reducing back between the 2 convolution actions
# C1: chs-> 2*chs, C2: 2*chs->chs

struct TSResidualBlock
  layers
  activation
end

(r::TSResidualBlock)(input) = r.activation(r.layers(input) + input)

@functor TSResidualBlock

function TSResidualBlock(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>(2*in_chs), pad = pad; init=_random_normal),
                	BatchNorm(2*in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, (2*in_chs)=>out_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(out_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    TSResidualBlock(layers, activation)|> cgpu
end

# Convolutions with known kernels

struct WResidualBlock
  layers
  activation
end

@functor WResidualBlock

function WResidualBlock(in_chs::Int, out_chs::Int; σ=elu)
    layers(kernel) = Conv(kernel, u_type.(reshape([0.0],1,1,1,1))|> cgpu)
    activation = Chain(BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    WResidualBlock(layers, activation)|> cgpu
end

function (r::WResidualBlock)(input, kernel1, kernel2)
    c1 = Conv(u_type.(kernel1)|> cgpu, u_type.([0.0])|> cgpu, pad=1)(input)
    ac1 = r.activation(c1)
    c2 = Conv(u_type.(kernel2)|> cgpu, u_type.([0.0])|> cgpu, pad=1)(ac1)
    y = r.activation(c2 + input)
    return y
end

# 3 convolution operations in one ResNet layer

struct SResidualBlock3C
  layers
  activation
end

(r::SResidualBlock3C)(input) = r.activation(r.layers(input) + input)

@functor SResidualBlock3C

function SResidualBlock3C(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal),
                	BatchNorm(in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, in_chs=>out_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(Conv(kernel, out_chs=>out_chs, pad = pad; init=_random_normal),
                    BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    SResidualBlock3C(layers, activation)|> cgpu
end

# Input as is without activation and batch normalization

struct SResidualBlockI
  layers
  activation
end

(r::SResidualBlockI)(input) = r.layers(r.activation(input)) + input

@functor SResidualBlockI

function SResidualBlockI(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal),
                	BatchNorm(in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    SResidualBlockI(layers, activation)|> cgpu
end

# Input as is without activation and batch normalization
# and doubling the channels and reducing back between the 2 convolution actions

struct TSResidualBlockI
  layers
  activation
end

(r::TSResidualBlockI)(input) = r.layers(r.activation(input)) + input

@functor TSResidualBlockI

function TSResidualBlockI(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    layers = Chain(Conv(kernel, in_chs=>(2*in_chs), pad = pad; init=_random_normal),
                	BatchNorm(2*in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, (2*in_chs)=>out_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(out_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    TSResidualBlockI(layers, activation)|> cgpu
end

# Multiply ResNet steps to create waves

struct ResidualBlock2
  c1
  c2
  activation
end

@functor ResidualBlock2

function ResidualBlock2(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    c1 = Chain(Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal),
                	BatchNorm(in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal))|> cgpu
    c2 = Chain(Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal),
                    BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, in_chs=>in_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    ResidualBlock2(c1, c2, activation)|> cgpu
end

function (r::ResidualBlock2)(input)
    f_t = input + r.c1(r.activation(input))
    f_t1 = 2 * f_t - input + r.activation(r.c2(f_t))
    return f_t1
end

struct TResidualBlock2
  c1
  c2
  activation
end

@functor TResidualBlock2

function TResidualBlock2(in_chs::Int, out_chs::Int; kernel = (3, 3), pad =1, σ=elu)
    c1 = Chain(Conv(kernel, in_chs=> (2*in_chs), pad = pad; init=_random_normal),
                	BatchNorm(2*in_chs),
                	x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, (2*in_chs)=>in_chs, pad = pad; init=_random_normal))|> cgpu
    c2 = Chain(Conv(kernel, in_chs=>(2*in_chs), pad = pad; init=_random_normal),
                    BatchNorm(2*in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                    Conv(kernel, (2*in_chs)=>in_chs, pad = pad; init=_random_normal))|> cgpu
    activation = Chain(BatchNorm(in_chs),
                    x->(σ == elu ? σ.(x,0.2f0) : σ.(x)))|> cgpu
    TResidualBlock2(c1, c2, activation)|> cgpu
end

function (r::TResidualBlock2)(input)
    f_t = input + r.c1(r.activation(input))
    f_t1 = 2 * f_t - input + r.activation(r.c2(f_t))
    return f_t1
end

# U-Net Types:

# SUnet - not deep, wide, without stride at the entrance
# n X n X 4 X bs -> (n/4) X (n/4) X 128 X bs

struct SUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor SUnet

function SUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=ResidualBlock)
    conv_down_blocks = Chain(ConvDown(16,32;σ=σ),
		      ConvDown(32,64;σ=σ),
		      ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(channels, 8; kernel = kernel, σ=σ),
		 resnet_type(8, 32; kernel = kernel, σ=σ),
		 resnet_type(64, 64; kernel = kernel, σ=σ),
		 resnet_type(128, 128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
		UNetUpBlock(128, 32; σ=σ),
		Chain(x-> (σ == elu ? σ.(x,0.2f0) : σ.(x)),
		Conv((3, 3), pad = 1, 64=>labels;init=_random_normal)))|> cgpu
    SUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::SUnet)(x::AbstractArray)

    # n X n X 4 X bs -> n X n X 32 X bs
    op = u.conv_blocks[1:2](x)
    # n X n X 32 X bs -> (n/2) X (n/2) X 64 X bs
    x1 = u.conv_blocks[3](u.conv_down_blocks[2](op))
    # (n/2) X (n/2) X 64 X bs -> (n/4) X (n/4) X 128 X bs
    x2 = u.conv_blocks[4](u.conv_down_blocks[3](x1))

    # (n/4) X (n/4) X 128 X bs
    up_x2 = u.conv_blocks[4](x2)
    up_x2 = u.conv_blocks[4](up_x2)

    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 128 X bs
    up_x1 = u.up_blocks[1](up_x2, x1)
    # (n/2) X (n/2) X 128 X bs -> n X n X 64 X bs
    up_x3 = u.up_blocks[2](up_x1, op)
    # n X n X 64 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x3)

end

# SSUnet - not deep, wide, with stride at the entrance
# n X n X 4 X bs -> (n/4) X (n/4) X 128 X bs

struct SSUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor SSUnet

function SSUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 64=>32;init=_random_normal)),
                        ConvUp(32, labels; σ=σ))|> cgpu
    SSUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::SSUnet)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 32 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 32 X bs -> (n/4) X (n/4) X 64 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](op))
    # (n/4) X (n/4) X 64 X bs -> (n/8) X (n/8) X 128 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1))

    # (n/8) X (n/8) X 128 X bs
    up_x2 = u.conv_blocks[3](x2)
    up_x2 = u.conv_blocks[3](up_x2)

    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 128 X bs
    up_x1 = u.up_blocks[1](up_x2, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 64 X bs
    up_x3 = u.up_blocks[2](up_x1, op)
    # (n/2) X (n/2) X 64 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x3)
    # (n/2) X (n/2) X 32 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x4)

end

# DTUnet - deep, wide, without stride at the entrance
# n X n X 4 X bs -> (n/8) X (n/8) X 256 X bs

struct DTUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor DTUnet

function DTUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=ResidualBlock)
    conv_down_blocks = Chain(ConvDown(16,32;σ=σ),
		      ConvDown(32,64;σ=σ),
		      ConvDown(64,128;σ=σ),
		      ConvDown(128,256;σ=σ))|> cgpu

    conv_blocks = Chain(ResidualBlock(channels, 8; kernel = kernel, σ=σ),
		 ResidualBlock(8, 32; kernel = kernel, σ=σ),
		 resnet_type(64, 64; kernel = kernel, σ=σ),
		 resnet_type(128, 128; kernel = kernel, σ=σ),
		 resnet_type(256, 256; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
		UNetUpBlock(128, 32; σ=σ),
        UNetUpBlock(256, 128; σ=σ),
        UNetUpBlock(256, 64; σ=σ),
		Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
		Conv((3, 3), pad = 1, 64=>labels;init=_random_normal)))|> cgpu

    DTUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::DTUnet)(x::AbstractArray)

    # n X n X 4 X bs -> n X n X 32 X bs
    op = u.conv_blocks[1:2](x)
    # n X n X 32 X bs -> (n/2) X (n/2) X 64 X bs
    x1 = u.conv_blocks[3](u.conv_down_blocks[2](op))
    # (n/2) X (n/2) X 64 X bs -> (n/4) X (n/4) X 128 X bs
    x2 = u.conv_blocks[4](u.conv_down_blocks[3](x1))
    # (n/4) X (n/4) X 128 X bs -> (n/8) X (n/8) X 256 X bs
    x3 = u.conv_blocks[5](u.conv_down_blocks[4](x2))

    # (n/8) X (n/8) X 256 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[5](up_x3)

    # (n/8) X (n/8) X 256 X bs -> (n/4) X (n/4) X 256 X bs
    up_x1 = u.up_blocks[3](up_x3, x2)
    # (n/4) X (n/4) X 256 X bs -> (n/2) X (n/2) X 128 X bs
    up_x2 = u.up_blocks[4](up_x1, x1)
    # (n/2) X (n/2) X 128 X bs -> n X n X 64 X bs
    up_x4 = u.up_blocks[2](up_x2, op)
    # n X n X 64 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x4)

end

# DNUnet - deep, narrow, without stride at the entrance
# n X n X 4 X bs -> (n/8) X (n/8) X 128 X bs

struct DNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor DNUnet

function DNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=ResidualBlock)
    conv_down_blocks = Chain(ConvDown(16,8;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(ResidualBlock(channels, 8; kernel = kernel, σ=σ),
         ResidualBlock(8, 16; kernel = kernel, σ=σ),
         resnet_type(32, 32; kernel = kernel, σ=σ),
         resnet_type(64, 64; kernel = kernel, σ=σ),
         resnet_type(128, 128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(64, 32; σ=σ),
        UNetUpBlock(64, 16; σ=σ),
        UNetUpBlock(128, 64; σ=σ),
        UNetUpBlock(128, 32; σ=σ),
        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
        Conv((3, 3), pad = 1, 32=>labels;init=_random_normal)))|> cgpu

    DNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::DNUnet)(x::AbstractArray)

    # n X n X 4 X bs -> n X n X 16 X bs
    op = u.conv_blocks[1:2](x)
    # n X n X 16 X bs -> (n/2) X (n/2) X 32 X bs
    x1 = u.conv_blocks[3](u.conv_down_blocks[2](op))
    # (n/2) X (n/2) X 32 X bs -> (n/4) X (n/4) X 64 X bs
    x2 = u.conv_blocks[4](u.conv_down_blocks[3](x1))
    # (n/4) X (n/4) X 64 X bs -> (n/8) X (n/8) X 128 X bs
    x3 = u.conv_blocks[5](u.conv_down_blocks[4](x2))

    # (n/8) X (n/8) X 128 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[5](up_x3)

    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 128 X bs
    up_x1 = u.up_blocks[3](up_x3, x2)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 64 X bs
    up_x2 = u.up_blocks[4](up_x1, x1)
    # (n/2) X (n/2) X 128 X bs -> n X n X 32 X bs
    up_x4 = u.up_blocks[2](up_x2, op)
    # n X n X 32 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x4)

end

# SDNUnet - deep, narrow, with stride at the entrance
# n X n X 4 X bs -> (n/8) X (n/8) X 128 X bs

struct SDNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor SDNUnet

function SDNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    SDNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::SDNUnet)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](op))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](x2))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2, op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](up_x4)
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# Old networks and their ResNet type:
# SDNUnet1 use SResidualBlock
# SDNUnet1D use SResidualBlock
# SDNUnetI use SResidualBlockI
# SDNUnet2 use ResidualBlock2
# SSDNUnet2 use SResidualBlock2
# SDNUnet21 use SResidualBlock3C

# SinUnetI - deep, narrow, with stride at the entrance, with sin operators
# n X n X 4 X bs -> (n/8) X (n/8) X 128 X bs

struct SinUnetI
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor SinUnetI

function SinUnetI(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    SinUnetI(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::SinUnetI)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = sin.(u.conv_blocks[2](u.conv_down_blocks[2](op)))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = sin.(u.conv_blocks[4](u.conv_down_blocks[4](x2)))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = sin.(u.conv_blocks[5](x3))
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2, op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](up_x4)
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# SDNNUnet - deep, narrow, with stride at the entrance, double ResNet steps
# n X n X 4 X bs -> (n/8) X (n/8) X 128 X bs

struct SDNNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor SDNNUnet

function SDNNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,12;σ=σ),
              ConvDown(12,24;σ=σ),
              ConvDown(24,48;σ=σ),
              ConvDown(48,96;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(12,12; kernel = kernel, σ=σ),
         resnet_type(24,24; kernel = kernel, σ=σ),
         resnet_type(48,48; kernel = kernel, σ=σ),
         resnet_type(96,96; kernel = kernel, σ=σ),
         resnet_type(96,96; kernel = kernel, σ=σ),
         resnet_type(96,96; kernel = kernel, σ=σ),
         resnet_type(96,96; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(96, 48; σ=σ),
                        UNetUpBlock(96, 24; σ=σ),
                        UNetUpBlock(48, 12; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 24=>12;init=_random_normal)),
                        ConvUp(12, labels; σ=σ))|> cgpu
    SDNNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::SDNNUnet)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 12 X bs
    op = u.conv_blocks[1](u.conv_blocks[1](u.conv_down_blocks[1](x)))
    # (n/2) X (n/2) X 12 X bs -> (n/4) X (n/4) X 24 X bs
    x1 = u.conv_blocks[2](u.conv_blocks[2](u.conv_down_blocks[2](op)))
    # (n/4) X (n/4) X 24 X bs -> (n/8) X (n/8) X 48 X bs
    x2 = u.conv_blocks[3](u.conv_blocks[3](u.conv_down_blocks[3](x1)))
    # (n/8) X (n/8) X 48 X bs -> (n/16) X (n/16) X 96 X bs
    x3 = u.conv_blocks[4](u.conv_blocks[4](u.conv_down_blocks[4](x2)))

    # (n/16) X (n/16) X 96 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 96 X bs -> (n/8) X (n/8) X 96 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 96 X bs -> (n/4) X (n/4) X 48 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 96 X bs -> (n/2) X (n/2) X 24 X bs
    up_x4 = u.up_blocks[3](up_x2, op)
    # (n/2) X (n/2) X 24 X bs -> (n/2) X (n/2) X 12 X bs
    up_x5 = u.up_blocks[4](up_x4)
    # (n/2) X (n/2) X 12 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# Architecture Types:

# SplitUNet: the encoder prepares 'k_chs' channels for the entry of the solver network

struct SplitUNet
  kappa_subnet
  solve_subnet
  indexes
end

@functor SplitUNet

function SplitUNet(in_chs, k_chs, s_model, k_model; kernel = (3, 3), indexes=3, σ=elu, resnet_type=SResidualBlock)
    kappa_subnet = k_model(indexes-2,k_chs;kernel=kernel,σ=σ,resnet_type=resnet_type)
    solve_subnet = s_model(k_chs+in_chs,2;kernel=kernel,σ=σ,resnet_type=resnet_type)
    SplitUNet(kappa_subnet, solve_subnet, indexes)
end

function (u::SplitUNet)(x::AbstractArray)
    kappa = reshape(x[:,:,3:u.indexes,1], size(x)[1], size(x)[1], u.indexes-2, 1)
    kappa_features = u.kappa_subnet(kappa)
    u.solve_subnet(cat(x, kappa_features, dims=3))
end

# FeaturesUNet: the encoder prepares hierarchical input for the solver network

struct FeaturesUNet
  kappa_subnet
  solve_subnet
  indexes
end

@functor FeaturesUNet

function FeaturesUNet(in_chs, k_chs, s_model, k_model; kernel = (3, 3), indexes=3, σ=elu, resnet_type=SResidualBlock)
    kappa_subnet = k_model(indexes-2,k_chs;kernel=kernel,σ=σ,resnet_type=resnet_type)
    solve_subnet = s_model(in_chs,2;kernel=kernel,σ=σ,resnet_type=resnet_type)
    FeaturesUNet(kappa_subnet, solve_subnet, indexes)
end

function (u::FeaturesUNet)(x::AbstractArray)
    kappa = reshape(x[:,:,3:u.indexes,1], size(x)[1], size(x)[1], u.indexes-2, 1)
    features = u.kappa_subnet(kappa)
    u.solve_subnet(x, features)
end

# FKappa + FSDNUnet: an encoder that prepares hierarchical input for the restriction step

struct FKappa
  conv_down_blocks
  conv_blocks
end

@functor FKappa

function FKappa(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=ResidualBlock)
    conv_down_blocks = Chain(ConvDown(16,32;σ=σ),
		      ConvDown(32,64;σ=σ),
		      ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(channels, 8; kernel = kernel, σ=σ),
		 resnet_type(8, 16; kernel = kernel, σ=σ),
		 resnet_type(32, 32; kernel = kernel, σ=σ),
		 resnet_type(64, 64; kernel = kernel, σ=σ),
		 resnet_type(128, 128; kernel = kernel, σ=σ))|> cgpu

    FKappa(conv_down_blocks, conv_blocks)
end

function (u::FKappa)(x::AbstractArray)

    # n X n X 4 X bs -> n X n X 16 X bs
    op = u.conv_blocks[1:2](x)
    # n X n X 16 X bs -> (n/2) X (n/2) X 32 X bs
    x1 = u.conv_blocks[3](u.conv_blocks[3](u.conv_down_blocks[1](op)))
    # (n/2) X (n/2) X 32 X bs -> (n/4) X (n/4) X 64 X bs
    x2 = u.conv_blocks[4](u.conv_blocks[4](u.conv_down_blocks[2](x1)))
    # (n/4) X (n/4) X 64 X bs -> (n/8) X (n/8) X 128 X bs
    x3 = u.conv_blocks[5](u.conv_blocks[5](u.conv_down_blocks[3](x2)))

    return [op, x1, x2, x3]
end

struct FSDNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor FSDNUnet

function FSDNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlockI)
    conv_down_blocks = Chain(ConvDown(20,16;σ=σ),
              ConvDown(48,32;σ=σ),
              ConvDown(96,64;σ=σ),
              ConvDown(192,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    FSDNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::FSDNUnet)(x::AbstractArray, features)

    # n X n X 3 X bs + n X n X 16 X bs-> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](cat(x, features[1], dims=3)))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](cat(op, features[2], dims=3)))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](cat(x1, features[3], dims=3)))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](cat(x2, features[4], dims=3)))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[4](x3)
    up_x3 = u.conv_blocks[4](up_x3)
    up_x3 = u.conv_blocks[4](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2, op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](up_x4)
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# FFKappa + FFSDNUnet: an encoder that prepares hierarchical input

struct FFKappa
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor FFKappa

function FFKappa(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    FFKappa(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::FFKappa)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](op))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](x2))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2, op)

    return [op, x1, x2, up_x3, up_x1, up_x2, up_x4]
end

struct FFSDNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor FFSDNUnet

function FFSDNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(32,32;σ=σ),
              ConvDown(64,64;σ=σ),
              ConvDown(128,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(256, 64; σ=σ),
                        UNetUpBlock(256, 32; σ=σ),
                        UNetUpBlock(128, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 64=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    FFSDNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::FFSDNUnet)(x::AbstractArray, features)

    # n X n X 4 X bs + n X n X 16 X bs-> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](cat(op, features[1], dims=3)))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](cat(x1, features[2], dims=3)))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](cat(x2, features[3], dims=3)))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](cat(up_x3, features[4], dims=3), x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](cat(up_x1, features[5], dims=3), x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](cat(up_x2, features[6], dims=3), op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](cat(up_x4, features[7], dims=3))
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# TFFKappa: an encoder that prepares hierarchical inputand uses double ResNet steps

struct TFFKappa
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor TFFKappa

function TFFKappa(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    TFFKappa(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::TFFKappa)(x::AbstractArray)

    # n X n X 4 X bs -> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[2](u.conv_blocks[1](u.conv_down_blocks[1](x)))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[4](u.conv_blocks[3](u.conv_down_blocks[2](op)))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[6](u.conv_blocks[5](u.conv_down_blocks[3](x1)))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[8](u.conv_blocks[7](u.conv_down_blocks[4](x2)))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[9](x3)
    up_x3 = u.conv_blocks[10](up_x3)
    up_x3 = u.conv_blocks[11](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.conv_blocks[15](u.up_blocks[1](up_x3, x2))
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.conv_blocks[14](u.up_blocks[2](up_x1, x1))
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.conv_blocks[13](u.up_blocks[3](up_x2, op))

    return [op, x1, x2, up_x3, up_x1, up_x2, up_x4]
end

# MFFSDNUnet: solver network which multiplies the hierarchical vectors with the input of each layer

struct MFFSDNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor MFFSDNUnet

function MFFSDNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu, resnet_type=SResidualBlock)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(resnet_type(16,16; kernel = kernel, σ=σ),
         resnet_type(32,32; kernel = kernel, σ=σ),
         resnet_type(64,64; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ),
         resnet_type(128,128; kernel = kernel, σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    MFFSDNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::MFFSDNUnet)(x::AbstractArray, features)

    # n X n X 4 X bs + n X n X 16 X bs-> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x))
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](op .* features[1]))
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1 .* features[2]))
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](x2 .* features[3]))

    # (n/16) X (n/16) X 128 X bs
    up_x3 = u.conv_blocks[5](x3)
    up_x3 = u.conv_blocks[6](up_x3)
    up_x3 = u.conv_blocks[7](up_x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3 .* features[4], x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1 .* features[5], x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2 .* features[6], op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](up_x4 .* features[7])
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

# Old encoder-solver and their ResNet type:
# FFSDNUnet,FFKappa = SResidualBlockI
# FFSDNUnet1,FFKappa1 = SResidualBlock

# WKappa + WSDNUnet: an encoder that learns weights for the ResNet steps of the solver network

struct WKappa
  conv_down_blocks
  conv_blocks
end

@functor WKappa

function WKappa(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu)
    conv_down_blocks = Chain(ConvDown(16,32;σ=σ),
		      ConvDown(32,64;σ=σ),
		      ConvDown(64,128;σ=σ),
		      ConvDown(128,256;σ=σ),
		      ConvDown(256,480;σ=σ),
		      ConvDown(480,480;σ=σ))|> cgpu

    conv_blocks = Chain(ResidualBlock(channels, 8; kernel = kernel, σ=σ),
		 ResidualBlock(8, 16; kernel = kernel, σ=σ),
		 ResidualBlock(32, 32; kernel = kernel, σ=σ),
		 ResidualBlock(64, 64; kernel = kernel, σ=σ),
		 ResidualBlock(128, 128; kernel = kernel, σ=σ),
		 ResidualBlock(256, 256; kernel = kernel, σ=σ),
		 ResidualBlock(480, 480; kernel = kernel, σ=σ))|> cgpu

    WKappa(conv_down_blocks, conv_blocks)
end

function (u::WKappa)(x::AbstractArray)

    # n X n X 4 X bs -> n X n X 16 X bs
    op = u.conv_blocks[1:2](x)
    # n X n X 16 X bs -> (n/2) X (n/2) X 32 X bs
    x1 = u.conv_blocks[3](u.conv_blocks[3](u.conv_down_blocks[1](op)))
    # (n/2) X (n/2) X 32 X bs -> (n/4) X (n/4) X 64 X bs
    x2 = u.conv_blocks[4](u.conv_blocks[4](u.conv_down_blocks[2](x1)))
    # (n/4) X (n/4) X 64 X bs -> (n/8) X (n/8) X 128 X bs
    x3 = u.conv_blocks[5](u.conv_blocks[5](u.conv_down_blocks[3](x2)))
    #  (n/8) X (n/8) X 128 X bs -> (n/16) X (n/16) X 256 X bs
    x4 = u.conv_blocks[6](u.conv_blocks[6](u.conv_down_blocks[4](x3)))
    # (n/16) X (n/16) X 256 X bs -> (n/32) X (n/32) X 480 X bs
    x5 = u.conv_blocks[7](u.conv_blocks[7](u.conv_down_blocks[5](x4)))

    return [big_block_filter!(3, x5[1:3,1:3,1:16,:], 16), big_block_filter!(3, x5[1:3,1:3,17:32,:], 16),
            big_block_filter!(3, x5[1:3,1:3,33:64,:], 32), big_block_filter!(3, x5[1:3,1:3,65:96,:], 32),
            big_block_filter!(3, x5[1:3,1:3,97:160,:], 64), big_block_filter!(3, x5[1:3,1:3,161:224,:], 64),
            big_block_filter!(3, x5[1:3,1:3,225:352,:], 128), big_block_filter!(3, x5[1:3,1:3,353:480,:], 128)]
end

struct WSDNUnet
  conv_down_blocks
  conv_blocks
  up_blocks
end

@functor WSDNUnet

function WSDNUnet(channels::Int = 2, labels::Int = channels; kernel = (3, 3), σ=elu)
    conv_down_blocks = Chain(ConvDown(channels,16;σ=σ),
              ConvDown(16,32;σ=σ),
              ConvDown(32,64;σ=σ),
              ConvDown(64,128;σ=σ))|> cgpu

    conv_blocks = Chain(WResidualBlock(16; σ=σ),
         WResidualBlock(32; σ=σ),
         WResidualBlock(64; σ=σ),
         WResidualBlock(128; σ=σ))|> cgpu

    up_blocks = Chain(UNetUpBlock(128, 64; σ=σ),
                        UNetUpBlock(128, 32; σ=σ),
                        UNetUpBlock(64, 16; σ=σ),
                        Chain(x->(σ == elu ? σ.(x,0.2f0) : σ.(x)),
                                Conv((3, 3), pad = 1, 32=>16;init=_random_normal)),
                        ConvUp(16, labels; σ=σ))|> cgpu
    WSDNUnet(conv_down_blocks, conv_blocks, up_blocks)
end

function (u::WSDNUnet)(x::AbstractArray, kernels)

    # n X n X 4 X bs -> (n/2) X (n/2) X 16 X bs
    op = u.conv_blocks[1](u.conv_down_blocks[1](x), kernels[1], kernels[2])
    # (n/2) X (n/2) X 16 X bs -> (n/4) X (n/4) X 32 X bs
    x1 = u.conv_blocks[2](u.conv_down_blocks[2](op), kernels[3], kernels[4])
    # (n/4) X (n/4) X 32 X bs -> (n/8) X (n/8) X 64 X bs
    x2 = u.conv_blocks[3](u.conv_down_blocks[3](x1), kernels[5], kernels[6])
    # (n/8) X (n/8) X 64 X bs -> (n/16) X (n/16) X 128 X bs
    x3 = u.conv_blocks[4](u.conv_down_blocks[4](x2), kernels[7], kernels[8])

    # (n/16) X (n/16) X 128 X bs
    up_x3 = x3 # u.conv_blocks[4](x3)

    # (n/16) X (n/16) X 128 X bs -> (n/8) X (n/8) X 128 X bs
    up_x1 = u.up_blocks[1](up_x3, x2)
    # (n/8) X (n/8) X 128 X bs -> (n/4) X (n/4) X 64 X bs
    up_x2 = u.up_blocks[2](up_x1, x1)
    # (n/4) X (n/4) X 128 X bs -> (n/2) X (n/2) X 32 X bs
    up_x4 = u.up_blocks[3](up_x2, op)
    # (n/2) X (n/2) X 32 X bs -> (n/2) X (n/2) X 16 X bs
    up_x5 = u.up_blocks[4](up_x4)
    # (n/2) X (n/2) X 16 X bs -> n X n X 2 X bs
    u.up_blocks[end](up_x5)

end

function create_model!(e_vcycle_input,kappa_input,gamma_input;kernel=(3,3),type=SUnet,k_type=NaN,resnet_type=SResidualBlock,k_chs=-1, indexes=3, σ=elu, arch=0)
    input = 2
    if e_vcycle_input == true
        input = input+2
    end
    if kappa_input == true
        input = input+1
    end
    if gamma_input == true
        input = input+1
    end

    if arch == 0 # A stand-alone U-Net
        return type(input,2;kernel=kernel,σ=σ,resnet_type=resnet_type)
    else
        if arch == 2 # Encoder with a hierarchical context
            @info "$(Dates.format(now(), "HH:MM:SS")) - FeaturesUNet"
            return FeaturesUNet(input,k_chs,type,k_type;kernel=kernel,indexes=indexes,σ=σ,resnet_type=resnet_type)
        else # Encoder with a simple context
            @info "$(Dates.format(now(), "HH:MM:SS")) - SplitUNet $(input) $(k_chs)"
            return SplitUNet(input,k_chs,type,k_type;kernel=kernel,indexes=indexes,σ=σ,resnet_type=resnet_type)
        end
    end
end


text:
function _extract_txt(body)
    sep = '\n'
    lines = split(body, sep)
    index = findfirst(contains("let txt = "), lines)
    txt_line = strip(lines[index])
    txt = txt_line[12:end-1]
    with_newlines = replace(txt, "\\n" => '\n')
    without_extraneous_newlines = strip(with_newlines, '\n')
    return without_extraneous_newlines
end

function _patch_with_terminal(body::String)
    txt = _extract_txt(body)
    return """
        <pre id="plutouiterminal">
        $txt
        </pre>
        """
end
precompile(_patch_with_terminal, (String,))

ext: jl


text:
const datadir = joinpath(dirname(@__DIR__), "data")

files = readdir(datadir)
filter!(x->endswith(x, ".names"), files)

for file in files
    classes = length(readlines(joinpath(datadir, file)))
    basename = splitext(file)[1]
    open(joinpath(datadir, "$basename.data"), "w") do io
        contents = """
        classes=$classes
        train  = $(datadir)/coco_train.txt
        valid  = $(datadir)/coco_test.txt
        names = $(datadir)/coco.names
        backup = $(datadir)/new-weights/
        eval=coco
        """
        write(io, contents)
    end
end
weightsfile = joinpath(datadir,"yolov3-tiny.weights")
!isfile(weightsfile) && download("https://pjreddie.com/media/files/yolov3-tiny.weights", weightsfile)


text:
function _A_ldiv_B_md!(dest, W::Woodbury, src,  R1, R2)
    _A_ldiv_B_md!(dest, W.A, src, R1, R2)
    tmp1 = _A_mul_B_md(W.V, dest, R1, R2)
    tmp2 = _A_mul_B_md(W.Cp, tmp1, R1, R2)
    tmp3 = _A_mul_B_md(W.U, tmp2, R1, R2)
    # TODO?: would be nice to fuse the next two steps
    tmp4 = _A_ldiv_B_md(W.A, tmp3, R1, R2)
    sub!(dest, tmp4)
end

function sub!(A, B)
    for I in eachindex(A, B)
        A[I] -= B[I]
    end
    A
end

check_matrix(W::Woodbury) = check_matrix(W.A)

max_stars_repo_name: timholy/FilterMD

max_issues_repo_name: timholy/FilterMD

max_forks_repo_name: timholy/FilterMD


text:
using Luxor

width = 400

"""
    drawlogo(s)

Draw the UnitfulCurrencies logo with the given scale `s`.
"""
function drawlogo(s)

    origin()
    scale(s)

    @layer begin # Draw background
        rotate(π/4)
        sethue("cornsilk")
        squircle(Point(0,0), 280, 280, :fill, rt=2.2)
        sethue("gray10")
        setline(2)
        squircle(Point(0,0), 280, 280, :stroke, rt=2.2)
        rotate(-π/4)
    end

    @layer begin # Draw UnitfulAssets text
        sethue("gray10")
        fontface("Luminari") # "Luminari" "Trattatello" "Arial Rounded MT Bold" "Impact"
        fontsize(74)
        text("Unitful", O - (40, 100), halign=:center)
        text("Assets", O - (-50, 28), halign=:center)
    end

    @layer begin # Draw coins
        coin_width = 88
        coin_height = 36
        colors = ["royalblue", "mediumorchid3", "forestgreen", "brown3"]
        for h = -1:1
            for v = 0:3-2*h
                position = Point(110*h, 125-25*v)
                for j=1:4
                    setcolor("gray40")
                    setdash("solid")
                    ellipse(position+(0,j), coin_width, coin_height, :stroke)
                    setcolor("gray10")
                    setdash("dotted")
                    ellipse(position+(0,j), coin_width, coin_height, :stroke)
                end
                setcolor("gray10")
                setdash("solid")
                ellipse(position+(0,5), coin_width, coin_height, :stroke)
                setcolor(colors[mod(v+h+1,4)+1])
                ellipse(position, coin_width, coin_height, :fill)
                setcolor("gray10")
                ellipse(position, coin_width, coin_height, :stroke)
            end
        end
    end

    finish()
end

for s = [1; 0.5; 0.18]
    Drawing(Int(s*width), Int(s*width), joinpath(@__DIR__, "logo_"*string(Int(s*width))*"x"*string(Int(s*width))*".png"))
    drawlogo(s)
end


text:
struct DeepGPModel{T,L}
    layers::T
    nLayers::Int
    likelihood::L
end

function DeepGPModel(layers...; likelihood=Normal())
    # Here do some checks with likelihood and last layer
    # Check dimensions between layers
    return DeepGPModel{typeof(layers),typeof(likelihood)}(
        layers, length(layers), likelihood
    )
end

# Base.:∘(l1::AbstractGPLayer,l2::AbstractGPLayer) = DeepGPModel([l1,l2])
Base.getindex(m::DeepGPModel, i::Int) = m.layers[i]
Base.length(m::DeepGPModel) = m.length
Base.iterate(m::DeepGPModel) = iterate(m.layers)

Base.iterate(m::DeepGPModel, state) = iterate(m.layers, state)

params(m::DeepGPModel) = params.(m.layers) #Eventually add parameters from likelihood

function ELBO(m::DeepGPModel, X, y; nsamples=1)
    return expec_log_like(m, X, y; nsamples=nsamples) - KL(m)
end

function KL(m::DeepGPModel)
    return sum(KL, m)
end

propagate(m::DeepGPModel, f) = propagate(Base.tail(m.layers), propagate(first(m), f))
propagate(::Tuple{}, f) = first.(f) # Final output
propagate(layers::Tuple, f) = propagate(Base.tail(layers), propagate(first(layers), f))

"""
    predict(m, X; nsamples)

Create `nsamples` samples with model `m` given samples `X`
"""
function predict(m::DeepGPModel, X::AbstractVector; nsamples=1)
    return reshape(propagate(m, repeat(X, nsamples)), :, nsamples)
end

function expec_log_like(m::DeepGPModel, X::AbstractVector, y; nsamples=1)
    f = predict(m, X; nsamples=nsamples)
    loss = mean(_logpdf.(Ref(m), eachcol(f), Ref(y)))
    return loss
end

function _logpdf(m::DeepGPModel, f, y)
    return sum(_logpdf.(Ref(m.likelihood), f, y))
end

function _logpdf(l::Normal, f::Real, y::Real)
    return -0.5 * abs2(y - f) / l.σ^2
    # Distributions.logpdf(l,f-y)
end


text:
push!(LOAD_PATH,"../../src/")

using NESSie
using NESSie.BEM
using NESSie.Format: readoff, readpqr
using NESSie.TestModel

using LinearAlgebra: norm
using PyPlot: cm_get_cmap, figure, plot_trisurf, show, title, xlabel, ylabel, zlabel, zlim

#=
    xiebem.jl

    Comparison of nonlocal reaction field potentials (or electrostatic potentials) for a
    Poisson test model generated from a given PQR file. The potentials are computed using
    the NESSie.TestModel and NESSie.BEM modules.
=#
if length(ARGS) > 1 || startswith(get(ARGS, 1, ""), "-")
    println("\n\e[1mUsage\e[0m: julia xiebem.jl [PQR_FILE]")
    exit(1)
end

const εΩ = 2.                           # dielectric constant (protein)
const εΣ = 80.                          # dielectric constant (solvent)
const ε∞ = 1.8                          # dielectric bulk response (solvent)
const λ  = 15.                          # correlation length (Å)
const radius = 1.                       # test model sphere radius (Å)
const numiter = 20                      # number of iterations (test model)
const reac = true                       # generate reaction field only
const lval = 2.                         # grid size [-l, l]² in xz-plane
const grid = 40                         # grid resolution (grid x grid)
const fpqr = "../../data/xie/2LZX.pqr"  # PQR file (ignored when given as argument)

function plotpotential(
        x   ::Vector{T},
        y   ::Vector{T},
        z   ::Vector{T},
        name::String = ""
    ) where T
    figure()
    plot_trisurf(x, y, z, cmap=cm_get_cmap("viridis"))
    length(name) > 0 && (title(name))
    xlabel("x component [\$\\AA\$]")
    ylabel("z component [\$\\AA\$]")
    zlabel(reac ? "Reaction field potential [V]" : "Electrostatic potential [V]")
    nothing
end

# generate observation points
Ξ = [ξ for Ξ in obspoints_plane(
        [-lval, zero(lval),  lval],
        [-lval, zero(lval), -lval],
        [ lval, zero(lval), -lval],
        grid, grid
    ) for ξ in Ξ]

# x- and y-axes of the plots
x = [ξ[1] for ξ in Ξ]
y = [ξ[3] for ξ in Ξ]

#=
    Nonlocal Xie model 1
=#
nlxie1 = NonlocalXieModel1(
            TestModel.XieModel(
                radius,
                readpqr(get(ARGS, 1, fpqr)),
                Option(εΩ, εΣ, ε∞, λ),
                compat=true),
            numiter
        )
molpot  = φmol(Ξ, nlxie1.charges)/4π/εΩ/ε0*NESSie.ec # molecular potential

z = [
        norm(ξ) < nlxie1.radius ?
        TestModel.φΩ(ξ, nlxie1) :
        TestModel.φΣ(ξ, nlxie1) for ξ in Ξ
    ]
reac && (z -= molpot)

zmin = (any(x -> x < 0, z) ? 1.2 : .8) * minimum(z)
zmax = (any(x -> x > 0, z) ? 1.2 : .8) * maximum(z)
plotpotential(x, y, z, "Nonlocal Poisson Test Model")

#=
    BEM
=#
surf = readoff("../../data/xie/unitsphere.off")
surf.charges = nlxie1.charges
surf.params  = nlxie1.params
bem = solve(NonlocalES, surf)

zΣ = BEM.φΣ(Ξ, bem)
zΩ = BEM.φΩ(Ξ, bem)
z  = [norm(ξ) < nlxie1.radius ? zΩ[i] : zΣ[i] for (i, ξ) in enumerate(Ξ)]
reac && (z -= molpot)

zmin = min(zmin, (any(x -> x < 0, z) ? 1.2 : .8) * minimum(z))
zmax = max(zmax, (any(x -> x > 0, z) ? 1.2 : .8) * maximum(z))
zlim((zmin, zmax))
plotpotential(x, y, z, "Nonlocal BEM")
zlim((zmin, zmax))

show()


text:
@doc raw"""
    SymmetricPositiveDefinite{N} <: AbstractEmbeddedManifold{ℝ,DefaultEmbeddingType}

The manifold of symmetric positive definite matrices, i.e.

````math
\mathcal P(n) =
\bigl\{
p ∈ ℝ^{n × n}\ \big|\ a^\mathrm{T}pa > 0 \text{ for all } a ∈ ℝ^{n}\backslash\{0\}
\bigr\}
````

The tangent space at ``T_p\mathcal P(n)`` reads

```math
    T_p\mathcal P(n) =
    \bigl\{
        X \in \mathbb R^{n×n} \big|\ X=X^\mathrm{T}
    \bigr\},
```
i.e. the set of symmetric matrices,

# Constructor

    SymmetricPositiveDefinite(n)

generates the manifold $\mathcal P(n) \subset ℝ^{n × n}$
"""
struct SymmetricPositiveDefinite{N} <: AbstractEmbeddedManifold{ℝ,DefaultEmbeddingType} end

SymmetricPositiveDefinite(n::Int) = SymmetricPositiveDefinite{n}()

@doc raw"""
    check_point(M::SymmetricPositiveDefinite, p; kwargs...)

checks, whether `p` is a valid point on the [`SymmetricPositiveDefinite`](@ref) `M`, i.e. is a matrix
of size `(N,N)`, symmetric and positive definite.
The tolerance for the second to last test can be set using the `kwargs...`.
"""
function check_point(M::SymmetricPositiveDefinite{N}, p; kwargs...) where {N}
    mpv = invoke(check_point, Tuple{supertype(typeof(M)),typeof(p)}, M, p; kwargs...)
    mpv === nothing || return mpv
    if !isapprox(norm(p - transpose(p)), 0.0; kwargs...)
        return DomainError(
            norm(p - transpose(p)),
            "The point $(p) does not lie on $(M) since its not a symmetric matrix:",
        )
    end
    if !all(eigvals(p) .> 0)
        return DomainError(
            eigvals(p),
            "The point $p does not lie on $(M) since its not a positive definite matrix.",
        )
    end
    return nothing
end

"""
    check_vector(M::SymmetricPositiveDefinite, p, X; kwargs... )

Check whether `X` is a tangent vector to `p` on the [`SymmetricPositiveDefinite`](@ref) `M`,
i.e. atfer [`check_point`](@ref)`(M,p)`, `X` has to be of same dimension as `p`
and a symmetric matrix, i.e. this stores tangent vetors as elements of the corresponding
Lie group.
The tolerance for the last test can be set using the `kwargs...`.
"""
function check_vector(M::SymmetricPositiveDefinite{N}, p, X; kwargs...) where {N}
    mpv = invoke(
        check_vector,
        Tuple{supertype(typeof(M)),typeof(p),typeof(X)},
        M,
        p,
        X;
        kwargs...,
    )
    mpv === nothing || return mpv
    if !isapprox(norm(X - transpose(X)), 0.0; kwargs...)
        return DomainError(
            X,
            "The vector $(X) is not a tangent to a point on $(M) (represented as an element of the Lie algebra) since its not symmetric.",
        )
    end
    return nothing
end

function decorated_manifold(M::SymmetricPositiveDefinite)
    return Euclidean(representation_size(M)...; field=ℝ)
end

@doc raw"""
    injectivity_radius(M::SymmetricPositiveDefinite[, p])
    injectivity_radius(M::MetricManifold{SymmetricPositiveDefinite,LinearAffineMetric}[, p])
    injectivity_radius(M::MetricManifold{SymmetricPositiveDefinite,LogCholeskyMetric}[, p])

Return the injectivity radius of the [`SymmetricPositiveDefinite`](@ref).
Since `M` is a Hadamard manifold with respect to the [`LinearAffineMetric`](@ref) and the
[`LogCholeskyMetric`](@ref), the injectivity radius is globally $∞$.
"""
injectivity_radius(::SymmetricPositiveDefinite) = Inf
injectivity_radius(::SymmetricPositiveDefinite, ::ExponentialRetraction) = Inf
injectivity_radius(::SymmetricPositiveDefinite, ::Any) = Inf
injectivity_radius(::SymmetricPositiveDefinite, ::Any, ::ExponentialRetraction) = Inf
eval(
    quote
        @invoke_maker 1 AbstractManifold injectivity_radius(
            M::SymmetricPositiveDefinite,
            rm::AbstractRetractionMethod,
        )
    end,
)

@doc raw"""
    manifold_dimension(M::SymmetricPositiveDefinite)

returns the dimension of
[`SymmetricPositiveDefinite`](@ref) `M`$=\mathcal P(n), n ∈ ℕ$, i.e.
````math
\dim \mathcal P(n) = \frac{n(n+1)}{2}.
````
"""
@generated function manifold_dimension(::SymmetricPositiveDefinite{N}) where {N}
    return div(N * (N + 1), 2)
end

"""
    mean(
        M::SymmetricPositiveDefinite,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolation();
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::AbstractManifold, args...)) of `x` using
[`GeodesicInterpolation`](@ref).
"""
mean(::SymmetricPositiveDefinite, ::Any)

function Statistics.mean!(
    M::SymmetricPositiveDefinite,
    p,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
)
    return mean!(M, p, x, w, GeodesicInterpolation(); kwargs...)
end

@doc raw"""
    project(M::SymmetricPositiveDefinite, p, X)

project a matrix from the embedding onto the tangent space $T_p\mathcal P(n)$ of the
[SymmetricPositiveDefinite](@ref) matrices, i.e. the set of symmetric matrices.
"""
project(::SymmetricPositiveDefinite, p, X)

project!(::SymmetricPositiveDefinite, Y, p, X) = (Y .= Symmetric((X + X') / 2))

@doc raw"""
    representation_size(M::SymmetricPositiveDefinite)

Return the size of an array representing an element on the
[`SymmetricPositiveDefinite`](@ref) manifold `M`, i.e. $n × n$, the size of such a
symmetric positive definite matrix on $\mathcal M = \mathcal P(n)$.
"""
@generated representation_size(::SymmetricPositiveDefinite{N}) where {N} = (N, N)

function Base.show(io::IO, ::SymmetricPositiveDefinite{N}) where {N}
    return print(io, "SymmetricPositiveDefinite($(N))")
end

@doc raw"""
    zero_vector(M::SymmetricPositiveDefinite,x)

returns the zero tangent vector in the tangent space of the symmetric positive
definite matrix `x` on the [`SymmetricPositiveDefinite`](@ref) manifold `M`.
"""
zero_vector(::SymmetricPositiveDefinite, ::Any)

zero_vector!(::SymmetricPositiveDefinite{N}, v, ::Any) where {N} = fill!(v, 0)


text:
#!/usr/bin/env julia

Pkg.update()
PKGS = ["Bio", "JSON", "DataStructures", "Combinatorics", "Iterators", "YAML"]
for p in PKGS
    Pkg.add(p)
    eval(Expr(:import, Symbol(p))) # and import it so that it gets precompiled
end


text:
using DifferentialEquations.DiffEqBase: DEIntegrator

export AbstractSpaceAdapterSparse, AbstractSpaceAdapter, init!, adapt!

abstract type AbstractSpaceAdapter end
abstract type AbstractSpaceAdapterSparse <: AbstractSpaceAdapter end

include("rstepadapters.jl")


text:
# 'cstrt_audet_ngon' provides the vertices coordinates (a,b) of a convex
# equilateral small n-gon for n = 2^s and s >= 4
function cstrt_audet_ngon(n)
if mod(log2(n),1)==0 && n>=16
    # initialization
    x(t) = sin(t)-(2*cos(t)-1)*cos((n/4-2)*t)*sin(n*t/4)/cos(2*t);
    F(t) = (2*x(t)+sin((n/2-1)*t))^2 + cos((n/2-1)*t)^2 - 4*sin(t/2)^2;
    z = find_zero(F,(2*asin(sin(pi/n)/2),pi/n));
    # construction
    a = zeros(n-1); b = zeros(n-1);
    a[Int(n/2-1)] = sin(z); a[Int(n/2+1)] = -a[Int(n/2-1)];
    b[Int(n/2-1)] = cos(z); b[Int(n/2+1)] = b[Int(n/2-1)];
    a[n-1] = a[Int(n/2-1)]-sin(2*z); a[1] = -a[n-1];
    b[n-1] = b[Int(n/2-1)]-cos(2*z); b[1] = b[n-1];
    a[n-2] = a[Int(n/2-1)]-sin(3*z); a[2] = -a[n-2];
    b[n-2] = b[Int(n/2-1)]-cos(3*z); b[2] = b[n-2];
    for i = 1:Int(n/8-1)
        a[Int(mod(n/2-1+i*(3*n/2-2),n))] = sin(z)-(2*cos(z)-1)*sin(i*(2*z-pi/2))*sin((i+1)*(2*z-pi/2))/cos(2*z);
        a[Int(n-mod(n/2-1+i*(3*n/2-2),n))] = -a[Int(mod(n/2-1+i*(3*n/2-2),n))];
        b[Int(mod(n/2-1+i*(3*n/2-2),n))] = cos(z)-(2*cos(z)-1)*sin(i*(2*z-pi/2))*cos((i+1)*(2*z-pi/2))/cos(2*z);
        b[Int(n-mod(n/2-1+i*(3*n/2-2),n))] = b[Int(mod(n/2-1+i*(3*n/2-2),n))];
        a[Int(mod(n-1+i*(3*n/2-2),n))] = a[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*sin((4*i+2)*z);
        a[Int(n-mod(n-1+i*(3*n/2-2),n))] = -a[Int(mod(n-1+i*(3*n/2-2),n))];
        b[Int(mod(n-1+i*(3*n/2-2),n))] = b[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*cos((4*i+2)*z);
        b[Int(n-mod(n-1+i*(3*n/2-2),n))] = b[Int(mod(n-1+i*(3*n/2-2),n))];
        a[Int(mod(n-2+i*(3*n/2-2),n))] = a[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*sin((4*i+3)*z);
        a[Int(n-mod(n-2+i*(3*n/2-2),n))] = -a[Int(mod(n-2+i*(3*n/2-2),n))];
        b[Int(mod(n-2+i*(3*n/2-2),n))] = b[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*cos((4*i+3)*z);
        b[Int(n-mod(n-2+i*(3*n/2-2),n))] = b[Int(mod(n-2+i*(3*n/2-2),n))];
        a[Int(mod(n+i*(3*n/2-2),n))] = a[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*sin((4*i+1)*z);
        a[Int(n-mod(n+i*(3*n/2-2),n))] = -a[Int(mod(n+i*(3*n/2-2),n))];
        b[Int(mod(n+i*(3*n/2-2),n))] = b[Int(mod(n/2-1+i*(3*n/2-2),n))]-(-1)^i*cos((4*i+1)*z);
        b[Int(n-mod(n+i*(3*n/2-2),n))] = b[Int(mod(n+i*(3*n/2-2),n))];
    end
    a[Int(n/2)] = 0; b[Int(n/2)] = 1;
    return(a,b)
end
end

max_stars_repo_name: cbingane/optigon

max_issues_repo_name: cbingane/optigon

max_forks_repo_name: cbingane/optigon


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -7.9284408379108005 <= q <= 14.377266175441108)

                     add_NL_constraint(m, :(swish(0.4343485145630286 + -0.5735104656090533*swish(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + -0.40833038450807946*swish(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.09714914815406805*swish(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.19623354791767778*swish(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + -0.3237276355524821*swish(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + swish(-0.11905941659917563 + 0.17224590981590104*swish(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.8406566605626429*swish(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.24719777641241691*swish(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.24308200319021767*swish(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + -0.39557038452022475*swish(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + swish(0.967203829156619 + 0.3785230261408685*swish(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + -0.23118999327456224*swish(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + -0.5197073892915105*swish(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + 0.6005003791404913*swish(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.15463985115916623*swish(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + swish(-0.18340968019959902 + 0.01980109053282808*swish(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.8767310690037822*swish(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.42243485881005904*swish(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.1694649191201152*swish(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.5278959282338329*swish(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + swish(0.7043455967116397 + -0.48930630508951234*swish(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.7375930151357304*swish(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + -0.9188673755139103*swish(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + 0.5801272291202126*swish(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.06743755305487698*swish(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
configs = Dict()
setups = [
    (false, 100, 200),
    (true, 100, 1001),
    (true, 1001, 100),
    (true, 100, 200)
]
for (i, setup) in enumerate(setups)
    configs[i] = EpidemicModel.Config(
        n_nodes=300,
        n0=3,
        k=3,
        agentcount=1000,
        transmission_prob=0.05,
        test_prob=0.8,
        quarantine_duration=140,
        behavior_reduction_factor=0.25,
        behavior_activated=setup[1],
        tick_authority_recommendation=setup[2],
        tick_authority_policy=setup[3],
        n_public_spaces=100,
        initial_infectives=5,
        n_iter=1000
    )
end


text:
#!/usr/bin/env julia

using SimradEK60
using SimradEK60TestData
using EchogramPyPlot
using EchogramColorSchemes
using ColorSchemes
using PyPlot
using LaTeXStrings

filename = EK60_SAMPLE
ps = pings(filename)
ps38 = [p for p in ps if p.frequency == 38000]
Sv38 = Sv(ps38)
r38 = R(ps38)

c = addwhite(ColorSchemes.viridis.colors[1:16:end])
echogram(Sv38, vmin=-95, vmax=-50, cmap=ColorMap("EK500", c.colors, length(c)), range=maximum(r38))
cbar = PyPlot.colorbar()
cbar[:set_label](L" $S_v$ [dB re 1 $m^{-1}$] ")
xlabel("Pings")
ylabel("Range [m]")
savefig("example2.png")


text:
module Telperion

using PropertyUtils: PropertyUtils, @with
using OrderedCollections: OrderedDict

export PropertyUtils, OrderedDict, @xy, @xydata, dummy

#-----------------------------------------------------------------------------# @xy
macro xy(df, ex)
    ex.args[1] === :~ && length(ex.args) == 3 || error("Formula must have the form `lhs ~ rhs`")
    lhs = [ex.args[2]]  # TODO: allow multivariate responses
    rhs = split_terms(ex.args[3])

    yblock = Expr(:call, :OrderedDict)
    for term in lhs
        push!(yblock.args, :($(string(term)) => Telperion.@with($df, $term)))
    end
    xblock = Expr(:call, :OrderedDict)
    for term in rhs
        push!(xblock.args, :($(string(term)) => Telperion.@with($df, $term)))
    end
    esc(quote
        Telperion.process!($df, $xblock), Telperion.process!($df, $yblock)
    end)
end

split_terms(x) = [x]
function split_terms(x::Expr)
    x.head === :call && x.args[1] === :+ || error("Formula must have the form `y ~ 1 + x`")
    x.args[2:end]
end

#-----------------------------------------------------------------------------# process!
# Deal with special things : Numbers and OrderedDicts
function process!(df, dict)
    for (k,v) in dict
        if v isa OrderedDict
            for (k2, v2) in v
                dict["$k [$k2]"] = v2
            end
            delete!(dict, k)
        elseif length(v) == 1
            dict[k] = fill(v, size(df, 1))
        end
    end
    dict
end

#-----------------------------------------------------------------------------# functions
dummy(x) = OrderedDict("$level" => x .== level for level in sort(unique(x))[2:end])

end


text:
using Distributions, Statistics, LinearAlgebra

import Base: size, length, show

import Statistics: mean, var, std

export  PropagationFunction,
        MeasurementFunction,
        FilteringFunction,
        RealMeasurementFunction

struct PropagationFunction end

struct MeasurementFunction end

struct FilteringFunction end

struct RealMeasurementFunction end

max_stars_repo_name: mleprovost/EnKF

max_issues_repo_name: mleprovost/EnKF

max_forks_repo_name: mleprovost/EnKF


text:
similar_from_indices(T::Type, p1::IndexTuple, p2::IndexTuple, A, CA::Symbol) =
    checked_similar_from_indices(nothing, T, p1, p2, A, CA)

function cached_similar_from_indices(sym::Symbol, T::Type, p1::IndexTuple, p2::IndexTuple, A, CA::Symbol)
    if use_cache()
        C = get(cache, sym, nothing)
        C′ = checked_similar_from_indices(C, T, p1, p2, A, CA)
        cache[sym] = C′
        return C′
    else
        return checked_similar_from_indices(nothing, T, p1, p2, A, CA)
    end
end

similar_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple, p1::IndexTuple, p2::IndexTuple,
    A, B, CA::Symbol, CB::Symbol) = checked_similar_from_indices(nothing, T, poA, poB, p1, p2, A, B, CA, CB)

function cached_similar_from_indices(sym::Symbol, T::Type, poA::IndexTuple, poB::IndexTuple,
    p1::IndexTuple, p2::IndexTuple, A, B, CA::Symbol, CB::Symbol)

    if use_cache()
        C = get(cache, sym, nothing)
        C′ = checked_similar_from_indices(C, T, poA, poB, p1, p2, A, B, CA, CB)
        cache[sym] = C′
        return C′
    else
        return checked_similar_from_indices(nothing, T, poA, poB, p1, p2, A, B, CA, CB)
    end
end


text:
# Gamma family
# Acknowledgement:
# This file is adapted from codes in https://github.com/rachtsingh/lgamma/blob/master/src/internals.c,
# which is again ported from https://bitbucket.org/eigen/eigen/overview.

function cuda1gammafamily()
    sprint() do s
        print(s,
"""
#define _USE_MATH_DEFINES
#include <float.h>
#include <math.h>

__device__ __host__ float polynomial_evaluation_32(float x, const float *f, int n) {
  float result = 0.0;
  for (int i = 0; i < n; i++) {
    result *= x;
    result += f[i];
  }
  return result;
}

__device__ __host__ double polynomial_evaluation_64(double x, const double *f, int n) {
  double result = 0.0;
  for (int i = 0; i < n; i++) {
    result *= x;
    result += f[i];
  }
  return result;
}

__device__ __host__ float digamma_impl_maybe_poly_32(const float s) {
  const float A[] = {-4.16666666666666666667E-3f, 3.96825396825396825397E-3f,
                     -8.33333333333333333333E-3f, 8.33333333333333333333E-2f};
  float z;
  if (s < 1.0e8f) {
    z = 1.0f / (s * s);
    return z * polynomial_evaluation_32(z, A, 4);
  } else {
    return 0.0f;
  }
}

__device__ __host__ double digamma_impl_maybe_poly_64(const double s) {
  const double A[] = {8.33333333333333333333E-2, -2.10927960927960927961E-2,
                      7.57575757575757575758E-3, -4.16666666666666666667E-3,
                      3.96825396825396825397E-3, -8.33333333333333333333E-3,
                      8.33333333333333333333E-2};

  double z;
  if (s < 1.0e17) {
    z = 1.0 / (s * s);
    return z * polynomial_evaluation_64(z, A, 7);
  } else
    return 0.0;
}
""")

for (T,F) in [("float","32"),("double","64")]

floor_str = (T == "float") ? "floorf" : "floor"
tan_str = (T == "float") ? "tanf" : "tan"
log_str = (T == "float") ? "logf" : "log"
max_str = (T == "float") ? "FLT_MAX" : "DBL_MAX"
pow_str = (T == "float") ? "powf" : "pow"
fabs_str = (T == "float") ? "fabsf" : "fabs"
one_str = (T == "float") ? "1.0f" : "1.0"
zeta_impl_series_2nd_cond = (T == "float") ? "" : "||(*a <= 9.0)"

print(s,
"""
__device__ __host__ $T digamma_impl_$F(const $T u) {

  $T xi = u;
  $T p, q, nz, s, w, yi;
  bool negative;

  const $T maxnum = $max_str;
  const $T m_pi = M_PI;

  negative = 0;
  nz = 0.0;

  const $T zero = 0.0;
  const $T one = 1.0;
  const $T half = 0.5;

  if (xi <= zero) {
    negative = one;
    q = xi;
    p = $floor_str(q);
    if (p == q) {
      return maxnum;
    }
    /* Remove the zeros of tan(m_pi x)
    * by subtracting the nearest integer from x
    */
    nz = q - p;
    if (nz != half) {
      if (nz > half) {
        p += one;
        nz = q - p;
      }
      nz = m_pi / $tan_str(m_pi * nz);
    } else {
      nz = zero;
    }
    xi = one - xi;
  }

  /* use the recurrence psi(x+1) = psi(x) + 1/x. */
  s = xi;
  w = zero;
  while (s < 10.0) {
    w += one / s;
    s += one;
  }

  yi = digamma_impl_maybe_poly_$F(s);

  yi = $log_str(s) - (half / s) - yi - w;

  return (negative) ? yi - nz : yi;

}

__device__ __host__ int zeta_impl_series_$F($T *a, $T *b, $T *s, const $T x,
                                         const $T machep) {
  int i = 0;
  while ((i < 9) $zeta_impl_series_2nd_cond) {
    i += 1;
    *a += $one_str;
    *b = $pow_str(*a, -x);
    *s += *b;
    if ($fabs_str(*b / *s) < machep) {
      return true;
    }
  }

  // Return whether we are done
  return false;
}

__device__ __host__ $T zeta_impl_$F($T x, $T q) {
  int i;
  $T p, r, a, b, k, s, t, w;

  const $T A[] = {
      12.0,
      -720.0,
      30240.0,
      -1209600.0,
      47900160.0,
      -1.8924375803183791606e9, /*1.307674368e12/691*/
      7.47242496e10,
      -2.950130727918164224e12,  /*1.067062284288e16/3617*/
      1.1646782814350067249e14,  /*5.109094217170944e18/43867*/
      -4.5979787224074726105e15, /*8.028576626982912e20/174611*/
      1.8152105401943546773e17,  /*1.5511210043330985984e23/854513*/
      -7.1661652561756670113e18  /*1.6938241367317436694528e27/236364091*/
  };

  const $T maxnum = $max_str;
  const $T zero = 0.0, half = 0.5, one = 1.0;
  const $T machep = 1e-15;

  if (x == one) return maxnum;

  if (x < one) {
    return zero;
  }

  if (q <= zero) {
    if (q == $floor_str(q)) {
      return maxnum;
    }
    p = x;
    r = $floor_str(p);
    if (p != r) return zero;
  }

  /* Permit negative q but continue sum until n+q > +9 .
   * This case should be handled by a reflection formula.
   * If q<0 and x is an integer, there is a relation to
   * the polygamma function.
   */
  s = $pow_str(q, -x);
  a = q;
  b = zero;

  // Run the summation in a helper function that is specific to the floating
  // precision
  if (zeta_impl_series_$F(&a, &b, &s, x, machep)) {
    return s;
  }

  w = a;
  s += b * w / (x - one);
  s -= half * b;
  a = one;
  k = zero;
  for (i = 0; i < 12; i++) {
    a *= x + k;
    b /= w;
    t = a * b / A[i];
    s = s + t;
    t = $fabs_str(t / s);
    if (t < machep) return s;
    k += one;
    a *= x + k;
    b /= w;
    k += one;
  }
  return s;
};

__device__ __host__ $T gamma_impl_$F($T x) {
  return tgamma(x);
}

__device__ __host__ $T polygamma_impl_$F(int n, $T x) {
  if (n == 0) {
    return digamma_impl_$F(x);
  }

  // dumb code to calculate factorials
  $T factorial = $one_str;
  for (int i = 0; i < n; i++) {
    factorial *= (i + 1);
  }

  $T s = n % 2 == 0 ? -$one_str : $one_str;

  return s * factorial * zeta_impl_$F(n + 1, x);
}

__device__ __host__ $T trigamma_impl_$F($T x) {
  return polygamma_impl_$F(1, x);
}
""")
        end
    end
end


text:
using Parameters

################################################################
### Lookup table for 2d gaussian (sigma = 2.5) where (1,1) is top left and (7,7) is bottom right
gauss25 = [
    0.02546481f0 0.02350698f0 0.01849125f0 0.01239505f0 0.00708017f0 0.00344629f0 0.00142946f0
    0.02350698f0 0.02169968f0 0.01706957f0 0.01144208f0 0.00653582f0 0.00318132f0 0.00131956f0
    0.01849125f0 0.01706957f0 0.01342740f0 0.00900066f0 0.00514126f0 0.00250252f0 0.00103800f0
    0.01239505f0 0.01144208f0 0.00900066f0 0.00603332f0 0.00344629f0 0.00167749f0 0.00069579f0
    0.00708017f0 0.00653582f0 0.00514126f0 0.00344629f0 0.00196855f0 0.00095820f0 0.00039744f0
    0.00344629f0 0.00318132f0 0.00250252f0 0.00167749f0 0.00095820f0 0.00046640f0 0.00019346f0
    0.00142946f0 0.00131956f0 0.00103800f0 0.00069579f0 0.00039744f0 0.00019346f0 0.00008024f0
]

################################################################
### AKAZE Descriptor Type
@enum DESCRIPTOR_TYPE begin
    SURF_UPRIGHT = 0   ###< Upright descriptors, not invariant to rotation
    SURF = 1
    MSURF_UPRIGHT = 2  ###< Upright descriptors, not invariant to rotation
    MSURF = 3
    MLDB_UPRIGHT = 4   ###< Upright descriptors, not invariant to rotation
    MLDB = 5
end

################################################################
### AKAZE Diffusivities
@enum DIFFUSIVITY_TYPE begin
    PM_G1 = 0
    PM_G2 = 1
    WEICKERT = 2
    CHARBONNIER = 3
end

select_diffusivity(diffusivity::DIFFUSIVITY_TYPE) =
    if diffusivity == PM_G1
        pm_g1_diffusivity
    elseif diffusivity == PM_G2
        pm_g2_diffusivity
    elseif diffusivity == WEICKERT
        weickert_diffusivity
    elseif diffusivity == CHARBONNIER
        charbonnier_diffusivity
    end

################################################################
### AKAZE Timing structure
mutable struct AKAZETiming
    kcontrast::Float64       ###< Contrast factor computation time in ms
    scale::Float64           ###< Nonlinear scale space computation time in ms
    derivatives::Float64     ###< Multiscale derivatives computation time in ms
    detector::Float64        ###< Feature detector computation time in ms
    extrema::Float64         ###< Scale space extrema computation time in ms
    subpixel::Float64        ###< Subpixel refinement computation time in ms
    descriptor::Float64      ###< Descriptors computation time in ms
end

################################################################
### AKAZE configuration options structure
@with_kw mutable struct AKAZEOptions
    omin::Int32                           ###< Initial octave level (-1 means that the size of the input image is duplicated)
    omax::Int32 = 4                       ###< Maximum octave evolution of the image 2^sigma (coarsest scale sigma units)
    nsublevels::Int32 = 4                 ###< Default number of sublevels per scale level
    img_width::Int32                      ###< Width of the input image
    img_height::Int32                     ###< Height of the input image
    soffset::Float32 = 1.6f0              ###< Base scale offset (sigma units)
    derivative_factor::Float32 = 1.5f0    ###< Factor for the multiscale derivatives
    sderivatives::Float32 = 1f0           ###< Smoothing factor for the derivatives
    diffusivity::DIFFUSIVITY_TYPE = PM_G2 ###< Diffusivity type

    dthreshold::Float32 = 3f-5            ###< Detector response threshold to accept point
    min_dthreshold::Float32 = 1f-5        ###< Minimum detector threshold to accept a point

    descriptor::DESCRIPTOR_TYPE = MLDB    ###< Type of descriptor
    descriptor_size::Int32 = 0            ###< Size of the descriptor in bits. 0->Full size
    descriptor_channels::Int32 = 3        ###< Number of channels in the descriptor (1, 2, 3)
    descriptor_pattern_size::Int32 = 10   ###< Actual patch size is 2*pattern_size*point.scale

    kcontrast::Float32 = 1f-3             ###< The contrast factor parameter
    kcontrast_percentile::Float32 = 7f-1  ###< Percentile level for the contrast factor
    kcontrast_nbins::UInt32 = 300         ###< Number of bins for the contrast factor histogram

    save_scale_space::Bool = false        ###< Set to true for saving the scale space images
    save_keypoints::Bool = false          ###< Set to true for saving the detected keypoints and descriptors
    show_results::Bool = true             ###< Set to true for displaying results
    verbosity::Bool = false               ###< Set to true for displaying verbosity information
end

################################################################
### AKAZE nonlinear diffusion filtering evolution
@with_kw mutable struct TEvolution{T}
    Lx::T            ###< First order spatial derivatives
    Ly::T            ###< First order spatial derivatives
    Lxx::T  ###< Second order spatial derivatives
    Lxy::T  ###< Second order spatial derivatives
    Lyy::T  ###< Second order spatial derivatives
    Lflow::T                ###< Diffusivity image
    Lt::T                   ###< Evolution image
    Lsmooth::T              ###< Smoothed image
    Lstep::T                ###< Evolution step update
    Ldet::T                 ###< Detector response
    etime::Float32 = 0f0      ###< Evolution time
    esigma::Float32 = 0f0     ###< Evolution sigma. For linear diffusion t = sigma^2 / 2
    octave::UInt32 = 0x0      ###< Image octave
    sublevel::UInt32 = 0x0    ###< Image sublevel in each octave
    sigma_size::UInt32 = 0x0  ###< Integer sigma. For computing the feature detector responses

end

construct_tevolution(; image_width, image_height, esigma, octave, sublevel) =
    TEvolution(
        Lx = zeros(image_height, image_width),
        Ly = zeros(image_height, image_width),
        Lxx = zeros(image_height, image_width),
        Lxy = zeros(image_height, image_width),
        Lyy = zeros(image_height, image_width),
        Lflow = zeros(image_height, image_width),
        Lt = zeros(image_height, image_width),
        Lsmooth = zeros(image_height, image_width),
        Lstep = zeros(image_height, image_width),
        Ldet = zeros(image_height, image_width),
        esigma = esigma,
        sigma_size = round(Int, esigma),
        etime = 0.5f0 * (esigma * esigma),
        octave = octave,
        sublevel = sublevel,
    )

        # evolution = map(Iterators.product(0:octavemax, 0:options.nsublevels-1)) do (i,j)
        #     construct_tevolution(
        #         image_width = options.img_width >> i,
        #         image_height = options.img_height >> i,
        #         esigma = options.soffset * 2.0^(j/options.nsublevels + i),
        #         octave = i,
        #         sublevel = j
        #     )
        # end

# Float32[1.2800001, 1.8101934, 2.5600002, 3.6203868, 5.1200004, 7.2407737, 10.240001, 14.481547, 20.480001, 28.963095, 40.960003, 57.92619, 81.920006, 115.85238, 163.84001, 231.70476]


text:
using MPI, Test

function runmpi(file; ntasks = 1)
    # by default some mpi runtimes will
    # complain if more resources (processes)
    # are requested than available on the node
    if MPI.MPI_LIBRARY == MPI.OpenMPI
        oversubscribe = `--oversubscribe`
    else
        oversubscribe = ``
    end
    @info "Running MPI test..." file ntasks
    @time @test mpiexec() do cmd
        run(`$cmd $oversubscribe -n $ntasks $(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $file`)
        true
    end
end


text:
# --- Utilities

"""
    _unzip(xs::Tuple)

# Examples
```jldoctest; setup = :(using Transducers: _unzip)
julia> _unzip(((1, 2, 3), (4, 5, 6)))
((1, 4), (2, 5), (3, 6))
```
"""
_unzip(xs::Tuple{Vararg{NTuple{N,Any}}}) where {N} = ntuple(i -> map(x -> x[i], xs), N)

if isdefined(Iterators, :Zip1)  # VERSION < v"1.1-"
    arguments(xs::Iterators.Zip1) = (xs.a,)
    arguments(xs::Iterators.Zip2) = (xs.a, xs.b)
    arguments(xs::Iterators.Zip) = (xs.a, arguments(xs.z)...)
else
    arguments(xs::Iterators.Zip) = xs.is
end

if VERSION < v"1.3"
    _Channel(f, ::Type{T}, size; kwargs...) where {T} =
        Channel(f; ctype = T, csize = size, kwargs...)
else
    _Channel(f, ::Type{T}, size; kwargs...) where {T} = Channel{T}(f, size; kwargs...)
end

function _materializer(xs)
    T = Tables.materializer(xs)
    return T isa Type ? T : _materializer(typeof(xs))
end

function _materializer(::Type{T}) where T
    S = ConstructionBase.constructorof(T)
    return S isa Type ? S : T
end

prefixed_type_name(@nospecialize x) =
    sprint(show, typeof(x), context = :module => Base)
# `:module => Base` to enforce that the type is prefixed even when
# `typeof(x)` is imported in `Main`.  See:
# https://github.com/JuliaLang/julia/pull/29466

const DenseSubVector{T} =
    SubArray{T, 1, Vector{T}, Tuple{UnitRange{Int64}}, true}

# https://github.com/JuliaLang/julia/pull/33533
if VERSION < v"1.4"
    const PartitionableArray = Vector
else
    const PartitionableArray = AbstractArray
end

const _non_executable_transducer_msg = """
Output type of the transducer is inferred to be a `Union{}`.  This
probably means one or more of the composed transducers throw.
"""

@inline _poptail(xs) = _poptail_impl(xs...)
@inline _poptail_impl(a) = (), a
@inline function _poptail_impl(a, xs...)
    head, tail = _poptail_impl(xs...)
    return (a, head...), tail
end

Base.@pure nthtype(::Val{1}, ::Type{T}) where {T <: Tuple} =
    Base.tuple_type_head(T)

Base.@pure nthtype(::Val{n}, ::Type{T}) where {n, T <: Tuple} =
    nthtype(Val(n - 1), Base.tuple_type_tail(T))

_cljapiurl(name) =
    "https://clojure.github.io/clojure/clojure.core-api.html#clojure.core/$name"
_cljref(name) =
    "[`$name` in Clojure]($(_cljapiurl(name)))"
_thx_clj(name) =
    "This API is modeled after $(_cljref(name))."

# https://github.com/JuliaLang/julia/pull/30575
const _true_str = sprint(show, true; context=:typeinfo => Bool)
const _false_str = sprint(show, false; context=:typeinfo => Bool)

# Just like `Function` but for defining some common methods.
abstract type _Function <: Function end

# Avoid `Function` fallbacks:
@nospecialize
Base.show(io::IO, ::MIME"text/plain", f::_Function) = show(io, f)
Base.print(io::IO, f::_Function) = show(io, f)
@specialize


text:
@inline Base.:^(v::AbstractSIMD{W,T}, i::Integer) where {W,T} = Base.power_by_squaring(v, i)
@inline Base.:^(v::AbstractSIMD{W,T}, i::Integer) where {W,T<:Union{Float32,Float64}} = Base.power_by_squaring(v, i)
@inline relu(x) = (y = zero(x); IfElse.ifelse(x > y, x, y))


text:
"Computes the lower bound of the posterior mean function of a Gaussian process in an interval."
function compute_μ_lower_bound(gp, x_L, x_U, theta_vec_train_squared; upper_flag=false)
    # Set minmax_factor to -1 if maximizing
    minmax_factor = upper_flag ? -1. : 1.
    # global variables: training_data, theta_vec_train_squared
    x_train = gp.x # Confirmed
    m = gp.nobs # Confirmed
    n = gp.dim # Dimension of input
    α_train = minmax_factor*gp.alpha

    sigma_prior = gp.kernel.σ2 # confirmed
    α_train *= sigma_prior # confirmed

    # Create the bounds on ϕ (z bounds)U
    z_i_vector = [compute_z_intervals(gp, gp.x[:, i], x_L, x_U) for i=1:m]

    # Get the theta vector
    theta_vec = ones(n) * 1 ./ (2*gp.kernel.ℓ2) # confirmed

    a_i_sum = 0.
    b_i_vec = zeros((1,m))
    C = 0.

    # For each training point
    for idx=1:m
        a_i, b_i = linear_lower_bound(α_train[idx], z_i_vector[idx][1], z_i_vector[idx][2]) # Confirmed!
        b_i_vec[idx] = b_i
        a_i_sum += a_i
        C += b_i * theta_vec_train_squared[idx]
    end

    # Hessian object, with respect to each "flex" point
    H = 2*sum(b_i_vec)*theta_vec   # nx1 vector
    f = zeros((n,1))
    for idx=1:n
        #               # 1x1            # 1xm       # mx1
        f[idx] = (-2 * theta_vec[idx] * b_i_vec * x_train[idx, :])[1]
    end

    x_mu_lb, f_val = separate_quadratic_program(H, f, x_L, x_U)
    x_mu_lb = hcat(x_mu_lb)

    lb = minmax_factor*(f_val + C + a_i_sum)
    gp.alpha *= minmax_factor
    # TODO: Check this - does this work as intended?
    ub, _ = predict_f(gp, x_mu_lb)
    ub = ub[1]*minmax_factor
    gp.alpha *= minmax_factor

    if upper_flag
        return x_mu_lb, ub, lb
    else
        return x_mu_lb, lb, ub
    end
end

function split_region(x_L, x_U)

    n = length(x_L)
    # TODO: Does not work in general
    x_avg = (x_L + x_U)/2
    # if n == 1
    #     return [(x_L, x_avg), (x_avg, x_U)]
    # end

    # if n == 2
    #     return [(x_L, x_avg), (x_avg, x_U),
    #             ([x_avg[1], x_L[2]], [x_U[1], x_avg[2]]),
    #             ([x_L[1], x_avg[2]], [x_avg[1], x_U[2]])]
    # end

    pairs = []
    for i=1:n
        push!(pairs, [[x_L[i], x_avg[i]], [x_avg[i], x_U[i]]])
    end

    # if n==1
    #     return pairs
    # end

    lowers = []
    uppers = []
    for i=1:n
        push!(lowers, [x_L[i], x_avg[i]])
        push!(uppers, [x_avg[i], x_U[i]])
    end
    lower_ranges = collect(Base.product(lowers...))
    upper_ranges = collect(Base.product(uppers...))

    return [[[lower...], [upper...]] for (lower, upper) in zip(lower_ranges, upper_ranges)]

   @error "How did you get here!"
end

"Computes the upper bound of the posterior covariance function of a Gaussian process in an interval."
function compute_σ_upper_bound(gp, x_L, x_U, R_inv)
    # TODO: Verify what R inverse is.
#     σ_noise = exp(gp.logNoise.value)
#     R_inv = [σ_noise^2 0, 0 σ_noise^2] # TODO: Generalize to any dimension
    # global variables: training_data, theta_vec_train_squared
    x_train = gp.x # Confirmed
    m = gp.nobs # Confirmed
    n = gp.dim # Dimension of input
    α_train = gp.alpha
    # TODO: Check the others too (k matrix, etc.)

    sigma_prior = gp.kernel.σ2 # confirmed
    α_train *= sigma_prior # confirmed

    # Create the bounds on ϕ (z bounds)
    z_i_vector = [compute_z_intervals(gp, gp.x[:, i], x_L, x_U) for i=1:m]

    # Get the theta vector

    theta_vec = ones(n) * 1 ./ (2*gp.kernel.ℓ2) # confirmed

    a_i_sum = 0.
    b_i_vec = zeros((1,m))

    # For each training point
    # THIS LOOP takes FOREVER: https://discourse.julialang.org/t/improving-performance-of-a-nested-for-loop/29705/4
    for idx=1:(m::Int)
       for subidx=1:(idx::Int)
            z_il_L = z_i_vector[idx][1] + z_i_vector[subidx][1]
            z_il_U = z_i_vector[idx][2] + z_i_vector[subidx][2]

            a_i, b_i = linear_lower_bound(R_inv[idx, subidx], z_il_L, z_il_U) # Confirmed!
            b_i_vec[idx] += b_i
            if subidx < idx
                a_i *= 2
                b_i_vec[subidx] += b_i
            end
            a_i_sum += a_i

        end
    end

    # Hessian object, with respect to each "flex" point
    H = 4*sum(b_i_vec)*theta_vec   # nx1 vector
#     @info H
    f = zeros((n,1))
    for idx=1:n
        #               # 1x1            # 1xm       # mx1
        f[idx] = (-4 * theta_vec[idx] * b_i_vec * x_train[idx, :])[1]
    end

    C = 0.
    theta_vec_train_squared = zeros(m);
    for i = 1:m
        theta_vec_train_squared[i] = transpose(theta_vec) * (gp.x[:, i].^2);
    end

    # TODO: Check this multiplication, doesn't matter for 1D
                               # Nx1         1x?
#     theta_vec_train_squared = theta_vec*transpose(x_train.^2)
    for idx=1:m
       C += 2 * b_i_vec[idx] * theta_vec_train_squared[idx]
    end
    x_σ_ub, f_val = separate_quadratic_program(H, f, x_L, x_U)

#     x_mu_lb = transpose(x_mu_lb)
    x_σ_ub = hcat(x_σ_ub)
#     @info x_mu_lb

    σ2_ub = sigma_prior*(1. - (f_val + C + a_i_sum))
    if σ2_ub < 0
        @warn "σ2_ub less than zero. Using trivial UB."
        σ_ub = sqrt(sigma_prior)
    else
        σ_ub = sqrt(σ2_ub)
    end

    _, σ2_lb = predict_f(gp, x_σ_ub)

    return x_σ_ub, sqrt(σ2_lb[1]), σ_ub
end

function compute_z_intervals(gp, x_i, x_L, x_U)
    # Might not work for N-D output, N>1
    n = length(x_i)
    z_i_L = 0.
    theta_vec = ones(n)
    theta_vec *= 1 ./(2(gp.kernel.ℓ2))

    for idx=1:n
        # If the point is inside of the range, the lower bound for the component is zero
        # @info x_i, x_L, x_U
        if x_i[idx] <= x_U[idx] && x_i[idx] >= x_L[idx]
            z_i_L_j = 0.
        else
            z_i_L_j = theta_vec[idx] * minimum([(x_i[idx] - x_L[idx])^2  , (x_i[idx] - x_U[idx])^2])
        end
        z_i_L += z_i_L_j
    end

    z_i_U = transpose(theta_vec) * max((x_i - x_L).^2, (x_i - x_U).^2)

    return z_i_L, z_i_U
end

function linear_lower_bound(α, z_i_L, z_i_U)
 # Now compute the linear under approximation (inlined for computational reasons)
    if α >= 0
        z_i_avg = (z_i_L + z_i_U)/2
        # Where does this come from?
        a_i = (1 + z_i_avg)*α*exp(-z_i_avg)
        b_i = -α*exp(-z_i_avg)
    else
        # This looks more like a slope equation
        a_i = α*exp(-z_i_L)
        a_i = α*(exp(-z_i_L) - z_i_L*(exp(-z_i_L) - exp(-z_i_U))/(z_i_L - z_i_U) )
        b_i = α*(exp(-z_i_L) - exp(-z_i_U))/(z_i_L - z_i_U)
    end

    return a_i, b_i
end

"A simple quadratic program solver."
function separate_quadratic_program(H, f, x_L, x_U; C=0.)

    # By default, set the optimal points to the lower bounds
    x_star = copy(x_L)   # Same size as number of dimensions
    f_val = 0.    # Value at x*
    n = length(x_L) # Number of dimensions.
    calc_f_part(ddf, df, point) = 0.5*ddf*point.^2 + df*point

    for idx=1:n
        # This is odd notation. shouldn't the lower bound always be less than upper?
#         if x_L[idx] < x_U[idx]
        x_critic = -f[idx]/H[idx]

        if H[idx] >= 0 && (x_critic <= x_U[idx]) && (x_critic >= x_L[idx])
            x_star[idx] = x_critic
            f_val_partial = calc_f_part(H[idx], f[idx], x_critic)
        else
            vec = [calc_f_part(H[idx], f[idx], x_L[idx]),
                   calc_f_part(H[idx], f[idx], x_U[idx])]
            f_val_partial = minimum(vec)
            if f_val_partial == vec[2]
                x_star[idx] = x_U[idx]
            end
        end
#         else
#             f_val_partial = calc_f_part(H[idx], f[idx], x_L[idx])
#         end
        f_val += f_val_partial
    end

    return x_star, f_val + C
end


text:
function update_param!(o, param)
    for s in fieldnames(typeof(o))
        x = getfield(o, s)
        if fieldcount(typeof(x)) > 0
            update_param!(x, param)
        else
            x′ = get(param, string(s), x)
            x !== x′ && setproperty!(o, s, x′)
        end
    end
end

function update_param!(o::PyObject, param)
    for (k, v) in param
        s = Symbol(k)
        PyCall.hasproperty(o, s) &&
        setproperty!(o, s, v)
    end
end

function gridsearch(f, bounds)
    xs = collect(Iterators.product(bounds...))
    ys = pmap(f ∘ collect, xs)
    yᵒ, i = findmin(ys)
    x = xs[i]
    return x, y, xs, ys
end

amploss(x, λ = zero(x)) = ifelse(x < zero(x), λ, one(x)) * x

entropy(p, ϵ = 1f-9) = -sum(p .* log.(p .+ ϵ)) / log(length(p))

mae(x, ϵ = 1f-2) = sum(z -> max(abs(z) - ϵ , zero(x)), x) / length(x)

eglob(pattern::String, prefix = pwd()) = eglob(Regex(pattern), prefix)

eglob(pattern::Regex, prefix = pwd()) = glob([pattern], prefix)

rglob(pattern, prefix = pwd()) = Iterators.flatten(map(rdf -> glob(pattern, rdf[1]), walkdir(prefix)))

function lcp(strs)
    io = IOBuffer()
    if !isempty(strs)
        i = 1
        while all(i ≤ length(s) for s in strs) &&
            all(s == strs[1][i] for s in getindex.(strs, i))
            print(io, strs[1][i])
            i += 1
        end
    end
    return String(take!(io))
end

function findtoroot(file)
    dir = pwd()
    while !isfile(file)
        dir = joinpath(dir, "..")
        file = joinpath(dir, file)
    end
    return realpath(file)
end

macro everydir(dir, ex)
    cwd = gensym()
    dir_quot, cwd_quot = Expr(:$, dir), Expr(:$, cwd)
    quote
        mkpath($dir); $cwd = pwd(); @everywhere cd($dir_quot)
        try $ex finally @everywhere cd($cwd_quot) end
    end |> esc
end

macro indir(dir, ex)
    cwd = gensym()
    quote
        mkpath($dir); $cwd = pwd(); cd($dir)
        try $ex finally cd($cwd) end
    end |> esc
end

macro roll(ex)
    quote
        if isdir("roll")
            date = Dates.format(today(), dateformat"yymmdd")
            mv("roll", "roll-" * date, force = true)
        end
        $(esc(ex))
        combine("roll")
    end
end

macro NT(xs...)
    @static if VERSION >= v"0.7.0"
        xs = [:($x = $x) for x in xs]
        esc(:(($(xs...),)))
    else
        esc(:(@NT($(xs...))($(xs...))))
    end
end

macro redirect(src, ex)
    src = src == :devnull ? "/dev/null" : src
    quote
        io = open($(esc(src)), "a")
        o, e = stdout, stderr
        redirect_stdout(io)
        redirect_stderr(io)
        try
            $(esc(ex)); sleep(0.01)
        finally
            flush(io); close(io)
            redirect_stdout(o)
            redirect_stderr(e)
        end
    end
end

zeroel(x) = zero(eltype(x))

oneel(x) = one(eltype(x))

lvcha(资金曲线) = 10000 * mean(amploss.(diff(资金曲线), 2f0))

function cumsum_reset(x, r)
    z = zero(x)
    s = zero(eltype(x))
    for i in 1:length(x)
        s = ifelse(r[i], zero(s), s + x[i])
        z[i] = s
    end
    return z
end

function drawdown(资金曲线)
    最大值 = 当前回撤 = 最大回撤 = zeroel(资金曲线)
    回撤期 = 最大回撤期 = 0
    for i in eachindex(资金曲线)
        回撤期 = ifelse(资金曲线[i] > 最大值, 0, 回撤期 + 1)
        最大值 = max(最大值, 资金曲线[i])
        当前回撤 = 资金曲线[i] - 最大值
        最大回撤 = min(当前回撤, 最大回撤)
        最大回撤期 = max(最大回撤期, 回撤期)
    end
    return 最大回撤, 最大回撤期
end

function pct_change(x, T)
    ϵ = eps(eltype(x))
    z = zero(x)
    for t in 1:length(x)
        t′ = clamp(t - T, 1, length(x))
        t1, t2 = minmax(t, t′)
        z[t] = (x[t2] - x[t1]) / (x[t1] + ϵ)
    end
    return z
end

@noinline function sharperatio(资金曲线, 一年天数 = 240)
    日收益率 = pct_change(资金曲线, 1)
    年化夏普率 = eltype(资金曲线)(mean(日收益率) ⧶ std(日收益率) * √一年天数)
end

@noinline function sortinoratio(资金曲线, 一年天数 = 240)
    日收益率 = pct_change(资金曲线, 1)
    eltype(r)(mean(日收益率) ⧶ std(min.(日收益率, 0)) * √一年天数)
end

function moving_max(A, k)
    B = similar(A)
    Q = Int[1]
    sizehint!(Q, length(A))
    for t = 1 : k
        if A[t] >= A[Q[1]]
            Q[1] = t
        end
        B[t] = A[Q[1]]
    end
    for t = (k + 1):length(A)
        while !isempty(Q) && A[t] >= A[Q[end]]
            pop!(Q)
        end
        while !isempty(Q) && Q[1] <=  t - k
            popfirst!(Q)
        end
        push!(Q, t)
        B[t] = A[Q[1]]
    end
    return B
end

function moving_min(A, k)
    B = similar(A)
    Q = Int[1]
    sizehint!(Q, length(A))
    for t = 1 : k
        if A[t] <= A[Q[1]]
            Q[1] = t
        end
        B[t] = A[Q[1]]
    end
    for t = (k + 1):length(A)
        while !isempty(Q) && A[t] <= A[Q[end]]
            pop!(Q)
        end
        while !isempty(Q) && Q[1] <=  t - k
            popfirst!(Q)
        end
        push!(Q, t)
        B[t] = A[Q[1]]
    end
    return B
end

parseenv(key, default::String) = get(ENV, string(key), string(default))

function parseenv(key, default::T) where T
    str = get(ENV, string(key), string(default))
    if hasmethod(parse, (Type{T}, String))
        parse(T, str)
    else
        include_string(Main, str)
    end
end

macro gc(exs...)
    Expr(:block, [:($ex = 0) for ex in exs]..., :(@eval GC.gc())) |> esc
end

function lockphase(x, P)
    P == 1 && return x
    N, T = size(x)
    x′ = zeros(Float32, N * P, T)
    for ph in 1:P, t in ph:P:T, n in 1:N
        n′ = n + N * (ph - 1)
        x′[n′, t] = x[n, t]
    end
    return x′
end

function get_index_price(;update = false, delay = 30)
    pkl = joinpath(DEPOT_PATH[1], "index.pkl")
    if isfile(pkl) && !update
        df = pd.read_pickle(pkl)
        if now() - df.index[end] < Day(delay)
            return df
        end
    end
    haskey(ENV, "JQDATA_USERNAME") || return nothing
    jqdata = pyimport("jqdatasdk")
    dfs = DataFrame[]
    for (pool, code) in zip(["SZ50", "HS300", "ZZ500", "ZZ1000"],
        ["000016.XSHG", "000300.XSHG", "000905.XSHG", "000852.XSHG"])
        df = jqdata.get_price(code, start_date = "2010-01-01", end_date = string(Date(now())), fields = "close")
        df.rename(columns = Dict("close" => pool), inplace = true)
        push!(dfs, df)
    end
    df = pdhcat(dfs...)
    df = df.ffill().bfill()
    df.to_pickle(pkl)
    return df
end

function multisort(xs::AbstractArray...; ka...)
    p = sortperm(first(xs); ka...)
    map(x -> x[p], xs)
end

function multisort!(xs::AbstractArray...; ka...)
    p = sortperm(first(xs); ka...)
    for x in xs
        permute!(x, p)
    end
    return xs
end

unsqueeze(xs, dim) = reshape(xs, (size(xs)[1:dim-1]..., 1, size(xs)[dim:end]...))

function read_csv(csv, a...; ka...)
    tmp = randstring() * ".csv"
    cp(csv, tmp, force = true)
    df = pd.read_csv(tmp, a...; ka...)
    rm(tmp, force = true)
    return df
end

function to_csv(df, csv, a...; ka...)
    tmp = randstring() * ".csv"
    df.to_csv(tmp, a...; ka...)
    mv(tmp, csv, force = true)
    return csv
end

function read_hdf(h5, a...; ka...)
    tmp = randstring() * ".h5"
    cp(h5, tmp, force = true)
    df = pd.read_hdf(tmp, a...; ka...)
    rm(tmp, force = true)
    return df
end

function to_hdf(df, h5, a...; ka...)
    tmp = randstring() * ".h5"
    df.to_hdf(tmp, a...; ka...)
    mv(tmp, h5, force = true)
    return h5
end

function to_trade_hour(timestamp)
    hour = unix2hour(timestamp)
    ifelse(hour > 16, hour - 24, hour)
end