module NMatrix::LAPACK

# File lib/nmatrix/lapack_core.rb, line 77
def alloc_svd_result(matrix)
  [
    NMatrix.new(matrix.shape[0], dtype: matrix.dtype),
    NMatrix.new([[matrix.shape[0],matrix.shape[1]].min,1], dtype: matrix.abs_dtype),
    NMatrix.new(matrix.shape[1], dtype: matrix.dtype)
  ]
end

# File lib/nmatrix/lapack_core.rb, line 176
def clapack_getri(order, n, a, lda, ipiv)
  raise(NotImplementedError,"clapack_getri requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 164
def clapack_potrf(order, uplo, n, a, lda)
  raise(NotImplementedError,"clapack_potrf requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 168
def clapack_potri(order, uplo, n, a, lda)
  raise(NotImplementedError,"clapack_potri requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 172
def clapack_potrs(order, uplo, n, nrhs, a, lda, b, ldb)
  raise(NotImplementedError,"clapack_potrs requires the nmatrix-atlas gem")
end

# File lib/nmatrix/atlas.rb, line 77
def geev(matrix, which=:both)
  raise(StorageTypeError, "LAPACK functions only work on dense matrices") unless matrix.dense?
  raise(ShapeError, "eigenvalues can only be computed for square matrices") unless matrix.dim == 2 && matrix.shape[0] == matrix.shape[1]

  jobvl = (which == :both || which == :left) ? :t : false
  jobvr = (which == :both || which == :right) ? :t : false

  n = matrix.shape[0]

  # Outputs
  eigenvalues = NMatrix.new([n, 1], dtype: matrix.dtype) # For real dtypes this holds only the real part of the eigenvalues.
  imag_eigenvalues = matrix.complex_dtype? ? nil : NMatrix.new([n, 1], dtype: matrix.dtype) # For complex dtypes, this is unused.
  left_output      = jobvl ? matrix.clone_structure : nil
  right_output     = jobvr ? matrix.clone_structure : nil

  # lapack_geev is a pure LAPACK routine so it expects column-major matrices,
  # so we need to transpose the input as well as the output.
  temporary_matrix = matrix.transpose
  NMatrix::LAPACK::lapack_geev(jobvl, # compute left eigenvectors of A?
                               jobvr, # compute right eigenvectors of A? (left eigenvectors of A**T)
                               n, # order of the matrix
                               temporary_matrix,# input matrix (used as work)
                               n, # leading dimension of matrix
                               eigenvalues,# real part of computed eigenvalues
                               imag_eigenvalues,# imag part of computed eigenvalues
                               left_output,     # left eigenvectors, if applicable
                               n, # leading dimension of left_output
                               right_output,    # right eigenvectors, if applicable
                               n, # leading dimension of right_output
                               2*n)
  left_output = left_output.transpose if jobvl
  right_output = right_output.transpose if jobvr


  # For real dtypes, transform left_output and right_output into correct forms.
  # If the j'th and the (j+1)'th eigenvalues form a complex conjugate
  # pair, then the j'th and (j+1)'th columns of the matrix are
  # the real and imag parts of the eigenvector corresponding
  # to the j'th eigenvalue.
  if !matrix.complex_dtype?
    complex_indices = []
    n.times do |i|
      complex_indices << i if imag_eigenvalues[i] != 0.0
    end

    if !complex_indices.empty?
      # For real dtypes, put the real and imaginary parts together
      eigenvalues = eigenvalues + imag_eigenvalues*Complex(0.0,1.0)
      left_output = left_output.cast(dtype: NMatrix.upcast(:complex64, matrix.dtype)) if left_output
      right_output = right_output.cast(dtype: NMatrix.upcast(:complex64, matrix.dtype)) if right_output
    end

    complex_indices.each_slice(2) do |i, _|
      if right_output
        right_output[0...n,i] = right_output[0...n,i] + right_output[0...n,i+1]*Complex(0.0,1.0)
        right_output[0...n,i+1] = right_output[0...n,i].complex_conjugate
      end

      if left_output
        left_output[0...n,i] = left_output[0...n,i] + left_output[0...n,i+1]*Complex(0.0,1.0)
        left_output[0...n,i+1] = left_output[0...n,i].complex_conjugate
      end
    end
  end

  if which == :both
    return [eigenvalues, left_output, right_output]
  elsif which == :left
    return [eigenvalues, left_output]
  else
    return [eigenvalues, right_output]
  end
end

# File lib/nmatrix/atlas.rb, line 166
def gesdd(matrix, workspace_size=nil)
  min_workspace_size = matrix.shape.min * (6 + 4 * matrix.shape.min) + matrix.shape.max
  workspace_size = min_workspace_size if workspace_size.nil? || workspace_size < min_workspace_size

  result = alloc_svd_result(matrix)

  m = matrix.shape[0]
  n = matrix.shape[1]

  # This is a pure LAPACK function so it expects column-major functions.
  # So we need to transpose the input as well as the output.
  matrix = matrix.transpose
  NMatrix::LAPACK::lapack_gesdd(:a, m, n, matrix, m, result[1], result[0], m, result[2], n, workspace_size)
  result[0] = result[0].transpose
  result[2] = result[2].transpose
  result
end

# File lib/nmatrix/atlas.rb, line 151
def gesvd(matrix, workspace_size=1)
  result = alloc_svd_result(matrix)

  m = matrix.shape[0]
  n = matrix.shape[1]

  # This is a pure LAPACK function so it expects column-major functions.
  # So we need to transpose the input as well as the output.
  matrix = matrix.transpose
  NMatrix::LAPACK::lapack_gesvd(:a, :a, m, n, matrix, m, result[1], result[0], m, result[2], n, workspace_size)
  result[0] = result[0].transpose
  result[2] = result[2].transpose
  result
end

# File lib/nmatrix/lapack_core.rb, line 160
def lapack_geev(jobvl, jobvr, n, a, lda, w, wi, vl, ldvl, vr, ldvr, lwork)
  raise(NotImplementedError,"lapack_geev requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 156
def lapack_gesdd(jobz, m, n, a, lda, s, u, ldu, vt, ldvt, lwork)
  raise(NotImplementedError,"lapack_gesdd requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 152
def lapack_gesvd(jobu, jobvt, m, n, a, lda, s, u, ldu, vt, ldvt, lwork)
  raise(NotImplementedError,"lapack_gesvd requires the nmatrix-atlas gem")
end

# File lib/nmatrix/lapack_core.rb, line 69
def laswp(matrix, ipiv)
  raise(ArgumentError, "expected NMatrix for argument 0") unless matrix.is_a?(NMatrix)
  raise(StorageTypeError, "LAPACK functions only work on :dense NMatrix instances") unless matrix.stype == :dense
  raise(ArgumentError, "expected Array ipiv to have no more entries than NMatrix a has columns") if ipiv.size > matrix.shape[1]

  clapack_laswp(matrix.shape[0], matrix, matrix.shape[1], 0, ipiv.size-1, ipiv, 1)
end

# File lib/nmatrix/atlas.rb, line 59
def posv(uplo, a, b)
  raise(ShapeError, "a must be square") unless a.dim == 2 && a.shape[0] == a.shape[1]
  raise(ShapeError, "number of rows of b must equal number of cols of a") unless a.shape[1] == b.shape[0]
  raise(StorageTypeError, "only works with dense matrices") unless a.stype == :dense && b.stype == :dense
  raise(DataTypeError, "only works for non-integer, non-object dtypes") if 
    a.integer_dtype? || a.object_dtype? || b.integer_dtype? || b.object_dtype?

  x     = b.clone
  clone = a.clone
  n = a.shape[0]
  nrhs = b.shape[1]
  clapack_potrf(:row, uplo, n, clone, n)
  # Must transpose b before and after: http://math-atlas.sourceforge.net/faq.html#RowSolve
  x = x.transpose
  clapack_potrs(:row, uplo, n, nrhs, clone, n, x, n)
  x.transpose
end