
using NNlib: DenseConvDims
import NNlib: conv!, ∇conv_filter!, ∇conv_data!, conv_bias_act!

using cuDNN: scalingParameter, CUDNN_CONVOLUTION, convdims,
             cudnnConvolutionBwdDataAlgoPerf,
             cudnnConvolutionForward!, cudnnConvolutionBwdFilterAlgoPerf,
             cudnnConvolutionBackwardData, cudnnConvolutionBackwardFilter,
             cudnnConvolutionBackwardBias
import cuDNN: cudnnConvolutionDescriptor

const CUDNNFloat = Union{Float16,BFloat16,Float32,Float64}
const CUDNNComplexFloat = Union{ComplexF16,ComplexF32,ComplexF64}

function cudnnConvolutionDescriptorAndPaddedInput(cdims::DenseConvDims, x::DenseCuArray{T}) where T
    # The main purpose of this function is to catch asymmetric padding which cudnn does not support
    # If we find asymmetric padding we'll make a copy of x which is manually padded so that we can
    # call cudnn with symmetric padding.
    pad = NNlib.padding(cdims)
    sdims = NNlib.spatial_dims(cdims)
    all(i -> pad[i] .== pad[i+1], 1:2:2sdims) && return (cudnnConvolutionDescriptor(cdims, x), x, identity)

    # Naive implementation, is there a faster way?
    # How much we need to pad x manually: The absolute difference between pad_left and pad_right, pad_top
    # and pad_bottom etc. respectively. We keep the sign here though because we use it below to figure out
    # which side of x to pad. Oh, and we use a CartesianIndex as we will mainly use this to index in x
    pad_manual = CartesianIndex(ntuple(i -> i > sdims ? 0 : pad[2(i-1)+1] - pad[2(i-1)+2], ndims(x)))

    # How much we can let cudnn pad: The smallest padding amount between pad_left and pad_right, pad_top
    # and pad_bottom etc. respectively
    pad_cudnn = ntuple(i -> min(pad[2(i-1)+1], pad[2(i-1)+2]), sdims)

    x_padded_size = ntuple(i -> i <= sdims ? size(x, i) + abs(pad_manual[i]) : size(x ,i), ndims(x))
    x_padded = similar(x, x_padded_size)
    fill!(x_padded, 0)
    # This is a bit yucky, but we are basically figuring out where in x_padded we shall insert x
    # Haven't benchmarked if this has any advantages over a more readable solution, e.g. writing dim
    # by dim to an array in a loop
    xIs = CartesianIndices(x)
    xI_first = first(xIs)
    xI_last = last(xIs)
    xIs_pad = max(xI_first, xI_first + pad_manual) : max(xI_last, xI_last + pad_manual)
    x_padded[xIs_pad] = x

    return cudnnConvolutionDescriptor(cdims, x_padded, pad_cudnn), x_padded, _x -> _x[xIs_pad]
end

# Compute (accumulation) type for the convolution descriptor. cuDNN's
# TRUE_HALF_CONFIG (compute in Float16) is only supported for a narrow set of 2D
# shapes; in particular it is unsupported for 3D convolutions and for many
# backward-filter shapes, surfacing as CUDNN_STATUS_NOT_SUPPORTED/BAD_PARAM.
# PSEUDO_HALF_CONFIG (Float16 data, Float32 compute) is broadly supported and is
# the standard choice (matching e.g. PyTorch). So we always compute Float16
# convolutions in Float32.
# Fixes https://github.com/FluxML/NNlib.jl/issues/505 and #515.
# BFloat16 has no native cuDNN accumulation config at all, so it likewise computes
# in Float32 (cuDNN requires a Float32 accumulator for CUDNN_DATA_BFLOAT16 data).
conv_compute_type(::Type{Float16}) = Float32
conv_compute_type(::Type{BFloat16}) = Float32
conv_compute_type(::Type{T}) where T = T

function cudnnConvolutionDescriptor(cdims::DenseConvDims, x::DenseCuArray{T}, pad = nnlibPadding(cdims)) where T
    mode=(NNlib.flipkernel(cdims) ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION)
    cudnnConvolutionDescriptor(convdims(pad, size(x),0),
                               convdims(NNlib.stride(cdims),size(x),1),
                               convdims(NNlib.dilation(cdims),size(x),1),
                               mode,
                               cudnnDataType(conv_compute_type(real(T))),
                               math_mode(),
                               CUDNN_DEFAULT_REORDER,
                               Cint(NNlib.groupcount(cdims)))
end

@inline function _complex!(y::DenseCuArray{T1}, yr::DenseCuArray{T2}, yi::DenseCuArray{T2}; bias=zero(T1), alpha=one(T1), beta=zero(T1), σ=identity) where {T1 <: CUDNNComplexFloat, T2<:CUDNNFloat}
    # if y is from similar(), it may have NaNs, and beta*NaN will propagate.
    if beta != 0
        @. y = σ(alpha*(yr + im*yi) + bias + beta*y)
    else
        @. y = σ(alpha*(yr + im*yi) + bias)
    end
    return y
end

function conv!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T}, cdims::DenseConvDims;
               alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    d, x, _ = cudnnConvolutionDescriptorAndPaddedInput(cdims, x)
    cudnnConvolutionForward!(y, w, x, d; alpha, beta, z=y)
end

# Complex convolution with Gauss's trick (1 complex mul === 3 real mul):
# Consider x = xr + im*xi, y = yr + im*yi,
# so x*y = (xr*yr - xi*yi) + im*(xr*yi + xi*yr).
# Let a = xr*yr,
#     b = xi*yi,
#     c = (xr + xi)*(yr + yi) = xr*yr + xr*yi + xi*yr + xi*yi.
# Then,
# x*y = (a - b) + im*(c - a - b).
# Convolution is linear so this multiplication trick translates to convolution.
function conv!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T}, cdims::DenseConvDims;
               alpha=1, beta=0, algo=-1) where T<:CUDNNComplexFloat
    xr, xi = reim(x)
    wr, wi = reim(w)
    a = conv!(similar(real(y)), xr, wr, cdims; algo=algo)
    b = conv!(similar(a), xi, wi, cdims; algo=algo)
    c = conv!(similar(a), xr + xi, wr + wi, cdims; algo=algo)
    return _complex!(y, a - b, c - a - b; alpha=alpha, beta=beta)
end

# (xr + im*xi) * w = xr*w + im*(xi*w)
function conv!(y::DenseCuArray{T1}, x::DenseCuArray{T1}, w::DenseCuArray{T2}, cdims::DenseConvDims;
               alpha=1, beta=0, algo=-1) where {T1<:CUDNNComplexFloat, T2<:CUDNNFloat}
    xr, xi = reim(x)
    yr = conv!(similar(real(y)), xr, w, cdims; algo=algo)
    yi = conv!(similar(yr), xi, w, cdims; algo=algo)
    return _complex!(y, yr, yi; alpha=alpha, beta=beta)
end

# x * (wr + im*wi) = x*wr + im*(x*wi)
function conv!(y::DenseCuArray{T1}, x::DenseCuArray{T2}, w::DenseCuArray{T1}, cdims::DenseConvDims;
               alpha=1, beta=0, algo=-1) where {T1<:CUDNNComplexFloat, T2<:CUDNNFloat}
    wr, wi = reim(w)
    yr = conv!(similar(real(y)), x, wr, cdims; algo=algo)
    yi = conv!(similar(yr), x, wi, cdims; algo=algo)
    return _complex!(y, yr, yi; alpha=alpha, beta=beta)
end

function conv_bias_act!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T},
                        cdims::DenseConvDims, bias::DenseCuArray{T}, σ=identity;
                        z::DenseCuArray{T}=y, alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    d, x, _ = cudnnConvolutionDescriptorAndPaddedInput(cdims, x)
    # only relu and identity are supported by cudnnConvolutionForward!
    activation = (σ == NNlib.relu ? CUDNN_ACTIVATION_RELU : CUDNN_ACTIVATION_IDENTITY)
    cudnnConvolutionForward!(y, w, x, d; z, bias, activation, alpha, beta)
    if activation === CUDNN_ACTIVATION_IDENTITY && σ ∉ (nothing, identity)
        @. y = σ(y)
    end
    return y
end

function conv_bias_act!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{T},
                        cdims::DenseConvDims, bias::DenseCuArray{T}, σ=identity;
                        z::DenseCuArray{T}=y, alpha=1, beta=0, algo=-1) where T<:CUDNNComplexFloat
    xr, xi = reim(x)
    wr, wi = reim(w)
    a = conv!(similar(real(y)), xr, wr, cdims; alpha=1, beta=0, algo=algo)
    b = conv!(similar(a), xi, wi, cdims; alpha=1, beta=0, algo=algo)
    c = conv!(similar(a), xr + xi, wr + wi, cdims; alpha=1, beta=0, algo=algo)
    return _complex!(y, a - b, c - a - b; bias=bias, alpha=alpha, beta=beta, σ=σ)
end

function ∇conv_data!(dx::DenseCuArray{T}, dy::DenseCuArray{T}, w::DenseCuArray{T},
                     cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
    convDesc, dx, depad = cudnnConvolutionDescriptorAndPaddedInput(cdims, dx)
    xDesc, yDesc, wDesc = cudnnTensorDescriptor(dx), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(w)
    p = cudnnConvolutionBwdDataAlgoPerf(wDesc, w, yDesc, dy, convDesc, xDesc, dx, beta!=0)
    with_workspace(p.memory) do workspace
        cudnnConvolutionBackwardData(handle(), alpha, wDesc, w, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, xDesc, dx)
    end
    return depad(dx)
end

function ∇conv_data!(dx::DenseCuArray{T}, dy::DenseCuArray{T}, w::DenseCuArray{T},
                     cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNComplexFloat
    dyr, dyi = reim(dy)
    wr, wi = reim(w)
    # note: w is conjugated, i.e. wi is negated below
    a = ∇conv_data!(similar(real(dx)), dyr, wr, cdims; alpha=1, beta=0, algo=algo)
    b = ∇conv_data!(similar(a), dyi, -wi, cdims; alpha=1, beta=0, algo=algo)
    c = ∇conv_data!(similar(a), dyr + dyi, wr - wi, cdims; alpha=1, beta=0, algo=algo)
    return _complex!(dx, a - b, c - a - b; alpha=alpha, beta=beta)
end

# dx = (dyr + im*dyi)*w = dyr*w + im*(dyi*w)
function ∇conv_data!(dx::DenseCuArray{T1}, dy::DenseCuArray{T1}, w::DenseCuArray{T2},
                     cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where {T1<:CUDNNComplexFloat, T2<:CUDNNFloat}
    dyr, dyi = reim(dy)
    dxr = ∇conv_data!(similar(real(dx)), dyr, w, cdims; alpha=1, beta=0, algo=algo)
    dxi = ∇conv_data!(similar(dxr), dyi, w, cdims; alpha=1, beta=0, algo=algo)
    return _complex!(dx, dxr, dxi; alpha=alpha, beta=beta)
end

function ∇conv_filter!(dw::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
                       cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNFloat
    if cudnnversion() < v"6"
        all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
    end
    if algo != -1
        @warn "The algo option has been deprecated, the fastest algo is computed automatically" maxlog=1
    end
    alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta);
    convDesc, x, _ = cudnnConvolutionDescriptorAndPaddedInput(cdims, x)
    xDesc, yDesc, wDesc = cudnnTensorDescriptor(x), cudnnTensorDescriptor(dy), cudnnFilterDescriptor(dw)
    p = cudnnConvolutionBwdFilterAlgoPerf(xDesc, x, yDesc, dy, convDesc, wDesc, dw, beta!=0);
    with_workspace(p.memory) do workspace
        cudnnConvolutionBackwardFilter(handle(), alpha, xDesc, x, yDesc, dy, convDesc, p.algo, workspace, sizeof(workspace), beta, wDesc, dw);
    end
    return dw
end

function ∇conv_filter!(dw::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
                       cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where T<:CUDNNComplexFloat
    xr, xi = reim(x)
    dyr, dyi = reim(dy)
    # note: x is conjugated, i.e. xi is negated below
    a = ∇conv_filter!(similar(real(dw)), xr, dyr, cdims; alpha=1, beta=0, algo=algo)
    b = ∇conv_filter!(similar(a), -xi, dyi, cdims; alpha=1, beta=0, algo=algo)
    c = ∇conv_filter!(similar(a), xr - xi, dyr + dyi, cdims; alpha=1, beta=0, algo=algo)
    return _complex!(dw, a - b, c - a - b; alpha=alpha, beta=beta)
end

# dw = x*(dyr + im*dyi) = x*dyr + im*(x*dyi)
function ∇conv_filter!(dw::DenseCuArray{T1}, x::DenseCuArray{T2}, dy::DenseCuArray{T1},
                       cdims::DenseConvDims; alpha=1, beta=0, algo=-1) where {T1<:CUDNNComplexFloat, T2<:CUDNNFloat}
    dyr, dyi = reim(dy)
    dwr = ∇conv_filter!(similar(real(dw)), x, dyr, cdims; alpha=1, beta=0, algo=algo)
    dwi = ∇conv_filter!(similar(dwr), x, dyi, cdims; alpha=1, beta=0, algo=algo)
    return _complex!(dw, dwr, dwi; alpha=alpha, beta=beta)
end

function ∇conv_bias!(db::DenseCuArray{T}, dy::DenseCuArray{T}; alpha=1, beta=0) where T<:CUDNNFloat
    alpha,beta = scalingParameter(T,alpha), scalingParameter(T,beta)
    bDesc, yDesc = cudnnTensorDescriptor.((db,dy))
    cudnnConvolutionBackwardBias(handle(), alpha, yDesc, dy, beta, bDesc, db)
    return db
end

function ∇conv_bias!(db::DenseCuArray{T}, dy::DenseCuArray{T}; alpha=1, beta=0) where T<:CUDNNComplexFloat
    dyr, dyi = reim(dy)
    dbr = ∇conv_bias!(similar(real(db)), dyr; alpha=1, beta=0)
    dbi = ∇conv_bias!(similar(dbr), dyi; alpha=1, beta=0)
    return _complex!(db, dbr, dbi; alpha=alpha, beta=beta)
end
