using NNlib
import NNlib: conv!, ∇conv_filter!, ∇conv_data!,
  maxpool!, meanpool!, ∇maxpool!, ∇meanpool!,
  softmax, softmax!, ∇softmax!, logsoftmax, logsoftmax!, ∇logsoftmax
using CUDAnative

# Activation functions
@cufunc σ(x) = ifelse(x < -80, zero(x), one(x) / (one(x) + exp(-x)))

@cufunc function logσ(x)
  max_v = max(zero(x), -x)
  z = exp(-max_v) + exp(-x-max_v)
  -(max_v + log(z))
end

@cufunc elu(x, α = one(x)) =
  ifelse(x ≥ 0, x/1, α * (exp(x) - one(x)))

# TODO: make @cufunc recognise its own definitions
cufunc(::typeof(swish)) = x -> x * cufunc(σ)(x)

@cufunc function selu(x)
  λ = oftype(x/1, 1.0507009873554804934193349852946)
  α = oftype(x/1, 1.6732632423543772848170429916717)
  λ * ifelse(x > 0, x/1, α * (exp(x) - 1))
end

@cufunc softplus(x) = log1p(exp(x))

if !@isdefined CUDNN
  function conv!(y::CuArray, x::CuArray, w::CuArray; kw...)
    error("CUDNN is not installed.")
  end
  function softmax!(out::CuVecOrMat, xs::CuVecOrMat)
    error("CUDNN is not installed.")
  end
  function logsoftmax!(out::CuVecOrMat, xs::CuVecOrMat)
    error("CUDNN is not installed.")
  end
end
