package breeze.stats.regression

import breeze.generic.UFunc
import breeze.linalg._
import org.netlib.util.intW
import dev.ludovic.netlib.lapack.LAPACK.{getInstance => lapack}
import java.util.Arrays

private object leastSquaresSingleImplementation {
  def doLeastSquaresSingle(
      data: DenseMatrix[Float],
      outputs: DenseVector[Float],
      workArray: Array[Float]): LeastSquaresRegressionSingleResult = {
    require(data.rows == outputs.size)
    require(data.rows > data.cols + 1)
    require(workArray.length >= 2 * data.rows * data.cols)

    val info = new intW(0)
    lapack.sgels(
      "N",
      data.rows,
      data.cols,
      1,
      data.data,
      data.rows,
      outputs.data,
      data.rows,
      workArray,
      workArray.length,
      info)
    if (info.`val` < 0) {
      throw new ArithmeticException("Least squares did not converge.")
    }

    val coefficients = new DenseVector[Float](Arrays.copyOf(outputs.data, data.cols))
    var r2 = 0.0
    for (i <- 0 until (data.rows - data.cols)) {
      r2 = r2 + math.pow(outputs.data(data.cols + i), 2)
    }
    LeastSquaresRegressionSingleResult(coefficients, r2.toFloat)
  }
}

case class LeastSquaresRegressionSingleResult(coefficients: DenseVector[Float], rSquared: Float)
    extends RegressionResult[DenseVector[Float], Float] {
  def apply(x: DenseVector[Float]): Float = coefficients.dot(x)

  def apply(X: DenseMatrix[Float]): DenseVector[Float] = X * coefficients
}

object leastSquaresSingle extends UFunc {
  implicit val matrixVectorWithWorkArray
    : Impl3[DenseMatrix[Float], DenseVector[Float], Array[Float], LeastSquaresRegressionSingleResult] =
    new Impl3[DenseMatrix[Float], DenseVector[Float], Array[Float], LeastSquaresRegressionSingleResult] {
      def apply(
          data: DenseMatrix[Float],
          outputs: DenseVector[Float],
          workArray: Array[Float]): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(data.copy, outputs.copy, workArray)
    }

  implicit val matrixVectorSpecifiedWork
    : Impl3[DenseMatrix[Float], DenseVector[Float], Int, LeastSquaresRegressionSingleResult] =
    new Impl3[DenseMatrix[Float], DenseVector[Float], Int, LeastSquaresRegressionSingleResult] {
      def apply(data: DenseMatrix[Float], outputs: DenseVector[Float], workSize: Int): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(data.copy, outputs.copy, new Array[Float](workSize))
    }

  implicit val matrixVector: Impl2[DenseMatrix[Float], DenseVector[Float], LeastSquaresRegressionSingleResult] =
    new Impl2[DenseMatrix[Float], DenseVector[Float], LeastSquaresRegressionSingleResult] {
      def apply(data: DenseMatrix[Float], outputs: DenseVector[Float]): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(
          data.copy,
          outputs.copy,
          new Array[Float](math.max(1, data.rows * data.cols * 2)))
    }
}

object leastSquaresSingleDestructive extends UFunc {
  implicit val matrixVectorWithWorkArray
    : Impl3[DenseMatrix[Float], DenseVector[Float], Array[Float], LeastSquaresRegressionSingleResult] =
    new Impl3[DenseMatrix[Float], DenseVector[Float], Array[Float], LeastSquaresRegressionSingleResult] {
      def apply(
          data: DenseMatrix[Float],
          outputs: DenseVector[Float],
          workArray: Array[Float]): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(data, outputs, workArray)
    }

  implicit val matrixVectorSpecifiedWork
    : Impl3[DenseMatrix[Float], DenseVector[Float], Int, LeastSquaresRegressionSingleResult] =
    new Impl3[DenseMatrix[Float], DenseVector[Float], Int, LeastSquaresRegressionSingleResult] {
      def apply(data: DenseMatrix[Float], outputs: DenseVector[Float], workSize: Int): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(data, outputs, new Array[Float](workSize))
    }

  implicit val matrixVector: Impl2[DenseMatrix[Float], DenseVector[Float], LeastSquaresRegressionSingleResult] =
    new Impl2[DenseMatrix[Float], DenseVector[Float], LeastSquaresRegressionSingleResult] {
      def apply(data: DenseMatrix[Float], outputs: DenseVector[Float]): LeastSquaresRegressionSingleResult =
        leastSquaresSingleImplementation.doLeastSquaresSingle(
          data,
          outputs,
          new Array[Float](math.max(1, data.rows * data.cols * 2)))
    }
}