module NMatrix::BLAS

# File lib/nmatrix/blas.rb, line 260
def asum(x, incx = 1, n = nil)
  n ||= x.size / incx
  raise(ArgumentError, "Expected dense NMatrix for arg 0") unless x.is_a?(NMatrix)
  raise(RangeError, "n out of range") if n*incx > x.size || n*incx <= 0 || n <= 0
  ::NMatrix::BLAS.cblas_asum(n, x, incx)
end

# File lib/nmatrix/blas.rb, line 301
def cblas_herk(order, uplo, trans, n, k, alpha, a, lda, beta, c, ldc)
  raise(NotImplementedError,"cblas_herk requires either the nmatrix-lapacke or nmatrix-atlas gem")
end

# File lib/nmatrix/blas.rb, line 297
def cblas_syrk(order, uplo, trans, n, k, alpha, a, lda, beta, c, ldc)
  raise(NotImplementedError,"cblas_syrk requires either the nmatrix-lapacke or nmatrix-atlas gem")
end

# File lib/nmatrix/blas.rb, line 293
def cblas_trmm(order, side, uplo, trans_a, diag, m, n, alpha, a, lda, b, ldb)
  raise(NotImplementedError,"cblas_trmm requires either the nmatrix-lapacke or nmatrix-atlas gem")
end

# File lib/nmatrix/blas.rb, line 71
def gemm(a, b, c = nil, alpha = 1.0, beta = 0.0, transpose_a = false, transpose_b = false, m = nil, n = nil, k = nil, lda = nil, ldb = nil, ldc = nil)
  raise(ArgumentError, 'Expected dense NMatrices as first two arguments.') unless a.is_a?(NMatrix) and b.is_a?(NMatrix) and a.stype == :dense and b.stype == :dense
  raise(ArgumentError, 'Expected nil or dense NMatrix as third argument.') unless c.nil? or (c.is_a?(NMatrix) and c.stype == :dense)
  raise(ArgumentError, 'NMatrix dtype mismatch.')                                                                                                    unless a.dtype == b.dtype and (c ? a.dtype == c.dtype : true)

  # First, set m, n, and k, which depend on whether we're taking the
  # transpose of a and b.
  if c
    m ||= c.shape[0]
    n ||= c.shape[1]
    k ||= transpose_a ? a.shape[0] : a.shape[1]

  else
    if transpose_a
      # Either :transpose or :complex_conjugate.
      m ||= a.shape[1]
      k ||= a.shape[0]

    else
      # No transpose.
      m ||= a.shape[0]
      k ||= a.shape[1]
    end

    n ||= transpose_b ? b.shape[0] : b.shape[1]
    c               = NMatrix.new([m, n], dtype: a.dtype)
  end

  # I think these are independent of whether or not a transpose occurs.
  lda ||= a.shape[1]
  ldb ||= b.shape[1]
  ldc ||= c.shape[1]

  # NM_COMPLEX64 and NM_COMPLEX128 both require complex alpha and beta.
  if a.dtype == :complex64 or a.dtype == :complex128
    alpha = Complex(1.0, 0.0) if alpha == 1.0
    beta  = Complex(0.0, 0.0) if beta  == 0.0
  end

  # For argument descriptions, see: http://www.netlib.org/blas/dgemm.f
  ::NMatrix::BLAS.cblas_gemm(:row, transpose_a, transpose_b, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc)

  return c
end

# File lib/nmatrix/blas.rb, line 143
def gemv(a, x, y = nil, alpha = 1.0, beta = 0.0, transpose_a = false, m = nil, n = nil, lda = nil, incx = nil, incy = nil)
  raise(ArgumentError, 'Expected dense NMatrices as first two arguments.') unless a.is_a?(NMatrix) and x.is_a?(NMatrix) and a.stype == :dense and x.stype == :dense
  raise(ArgumentError, 'Expected nil or dense NMatrix as third argument.') unless y.nil? or (y.is_a?(NMatrix) and y.stype == :dense)
  raise(ArgumentError, 'NMatrix dtype mismatch.')                                                                                                    unless a.dtype == x.dtype and (y ? a.dtype == y.dtype : true)

  m ||= transpose_a == :transpose ? a.shape[1] : a.shape[0]
  n ||= transpose_a == :transpose ? a.shape[0] : a.shape[1]
  raise(ArgumentError, "dimensions don't match") unless x.shape[0] == n && x.shape[1] == 1

  if y
    raise(ArgumentError, "dimensions don't match") unless y.shape[0] == m && y.shape[1] == 1
  else
    y = NMatrix.new([m,1], dtype: a.dtype)
  end

  lda               ||= a.shape[1]
  incx      ||= 1
  incy      ||= 1

  ::NMatrix::BLAS.cblas_gemv(transpose_a, m, n, alpha, a, lda, x, incx, beta, y, incy)

  return y
end

# File lib/nmatrix/blas.rb, line 283
def nrm2(x, incx = 1, n = nil)
  n ||= x.size / incx
  raise(ArgumentError, "Expected dense NMatrix for arg 0") unless x.is_a?(NMatrix)
  raise(RangeError, "n out of range") if n*incx > x.size || n*incx <= 0 || n <= 0
  ::NMatrix::BLAS.cblas_nrm2(n, x, incx)
end

# File lib/nmatrix/blas.rb, line 189
def rot(x, y, c, s, incx = 1, incy = 1, n = nil, in_place=false)
  raise(ArgumentError, 'Expected dense NMatrices as first two arguments.')    unless x.is_a?(NMatrix) and y.is_a?(NMatrix) and x.stype == :dense and y.stype == :dense
  raise(ArgumentError, 'NMatrix dtype mismatch.')                                                                                                       unless x.dtype == y.dtype
  raise(ArgumentError, 'Need to supply n for non-standard incx, incy values') if n.nil? && incx != 1 && incx != -1 && incy != 1 && incy != -1

  n ||= [x.size/incx.abs, y.size/incy.abs].min

  if in_place
    ::NMatrix::BLAS.cblas_rot(n, x, incx, y, incy, c, s)
    return [x,y]
  else
    xx = x.clone
    yy = y.clone

    ::NMatrix::BLAS.cblas_rot(n, xx, incx, yy, incy, c, s)

    return [xx,yy]
  end
end

# File lib/nmatrix/blas.rb, line 217
def rot!(x, y, c, s, incx = 1, incy = 1, n = nil)
  rot(x,y,c,s,incx,incy,n,true)
end

# File lib/nmatrix/blas.rb, line 237
def rotg(ab)
  raise(ArgumentError, "Expected dense NMatrix of shape [2,1] or [1,2]") unless ab.is_a?(NMatrix) && ab.stype == :dense && ab.size == 2

  ::NMatrix::BLAS.cblas_rotg(ab)
end