/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <helpers/EigenValsAndVecs.h>
#include <helpers/FullPivLU.h>
#include <helpers/HessenbergAndSchur.h>
#include <helpers/MmulHelper.h>
#include <helpers/Sqrtm.h>
#include <ops/declarable/helpers/lup.h>

namespace sd {
namespace ops {
namespace helpers {

//////////////////////////////////////////////////////////////////////////
template <typename T>
static void sqrtmQuasiTrianDiag(NDArray& matrixT, NDArray& sqrtT) {
  const int rows = matrixT.sizeAt(0);

  for (int i = 0; i < rows; i++) {
    if (i == rows - 1 || matrixT.t<T>(i + 1, i) == (T)0) {
      const auto elemT = matrixT.t<T>(i, i);
      if (elemT < (T)0)
        THROW_EXCEPTION(
            "ops::helpers::Sqrtm::sqrtmQuasiTrianDiag: can't take sqrt of negative diagonal element of T matrix !");
      sqrtT.r<T>(i, i) = math::sd_sqrt<T, T>(elemT);
    } else {
      NDArray *esViewPtr = matrixT({i, i + 2, i, i + 2}, true);
      EigenValsAndVecs<T> es(*esViewPtr);  // es._Vecs {2,2,2}, es._Vals{2,2}
      delete esViewPtr;

       NDArray& vecs = es._Vecs;
       NDArray& vals = es._Vals;

      const T& vecsReal00 = vecs.t<T>(0, 0, 0);
      const T& vecsImag00 = vecs.t<T>(0, 0, 1);
      const T& vecsReal01 = vecs.t<T>(0, 1, 0);
      const T& vecsImag01 = vecs.t<T>(0, 1, 1);
      const T& vecsReal10 = vecs.t<T>(1, 0, 0);
      const T& vecsImag10 = vecs.t<T>(1, 0, 1);
      const T& vecsReal11 = vecs.t<T>(1, 1, 0);
      const T& vecsImag11 = vecs.t<T>(1, 1, 1);

      // es.eigenvalues().cwiseSqrt().asDiagonal()
      T eigenValsSqrt[2][2];
      eigenValsSqrt[0][0] = vals.t<T>(0, 0);
      eigenValsSqrt[0][1] = vals.t<T>(0, 1);
      eigenValsSqrt[1][0] = vals.t<T>(1, 0);
      eigenValsSqrt[1][1] = vals.t<T>(1, 1);
      EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[0][0], eigenValsSqrt[0][1]);
      EigenValsAndVecs<T>::sqrtComplexNum(eigenValsSqrt[1][0], eigenValsSqrt[1][1]);

      // es.eigenvectors() * es.eigenvalues().cwiseSqrt().asDiagonal()
      T vecsElem[2][2][2];
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
                                               vecsElem[0][0][0], vecsElem[0][0][1]);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
                                               vecsElem[0][1][0], vecsElem[0][1][1]);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal10, vecsImag10, eigenValsSqrt[0][0], eigenValsSqrt[0][1],
                                               vecsElem[1][0][0], vecsElem[1][0][1]);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal11, vecsImag11, eigenValsSqrt[1][0], eigenValsSqrt[1][1],
                                               vecsElem[1][1][0], vecsElem[1][1][1]);

      // es.eigenvectors().inverse()
      T vecsElemInv[2][2][2];

      T tempReal, tempImag, divisorReal, divisorImag;
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal00, vecsImag00, vecsReal11, vecsImag11, divisorReal,
                                               divisorImag);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsReal01, vecsImag01, vecsReal10, vecsImag10, tempReal, tempImag);
      divisorReal -= tempReal;
      divisorImag -= tempImag;

      EigenValsAndVecs<T>::divideComplexNums(vecsReal11, vecsImag11, divisorReal, divisorImag, vecsElemInv[0][0][0],
                                             vecsElemInv[0][0][1]);
      EigenValsAndVecs<T>::divideComplexNums(-vecsReal01, -vecsImag01, divisorReal, divisorImag, vecsElemInv[0][1][0],
                                             vecsElemInv[0][1][1]);
      EigenValsAndVecs<T>::divideComplexNums(-vecsReal10, -vecsImag10, divisorReal, divisorImag, vecsElemInv[1][0][0],
                                             vecsElemInv[1][0][1]);
      EigenValsAndVecs<T>::divideComplexNums(vecsReal00, vecsImag00, divisorReal, divisorImag, vecsElemInv[1][1][0],
                                             vecsElemInv[1][1][1]);

      // result
      T result[2][2][2];

      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][0][0],
                                               vecsElemInv[0][0][1], tempReal, tempImag);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][0][0],
                                               vecsElemInv[1][0][1], result[0][0][0], result[0][0][1]);
      result[0][0][0] += tempReal;

      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][0][0], vecsElem[0][0][1], vecsElemInv[0][1][0],
                                               vecsElemInv[0][1][1], tempReal, tempImag);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[0][1][0], vecsElem[0][1][1], vecsElemInv[1][1][0],
                                               vecsElemInv[1][1][1], result[0][1][0], result[0][1][1]);
      result[0][1][0] += tempReal;

      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][0][0],
                                               vecsElemInv[0][0][1], tempReal, tempImag);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][0][0],
                                               vecsElemInv[1][0][1], result[1][0][0], result[1][0][1]);
      result[1][0][0] += tempReal;

      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][0][0], vecsElem[1][0][1], vecsElemInv[0][1][0],
                                               vecsElemInv[0][1][1], tempReal, tempImag);
      EigenValsAndVecs<T>::multiplyComplexNums(vecsElem[1][1][0], vecsElem[1][1][1], vecsElemInv[1][1][0],
                                               vecsElemInv[1][1][1], result[1][1][0], result[1][1][1]);
      result[1][1][0] += tempReal;

      sqrtT.r<T>(i, i) = result[0][0][0];
      sqrtT.r<T>(i, i + 1) = result[0][1][0];
      sqrtT.r<T>(i + 1, i) = result[1][0][0];
      sqrtT.r<T>(i + 1, i + 1) = result[1][1][0];

      ++i;
    }
  }
}

//////////////////////////////////////////////////////////////////////////
// all matrices are {2,2} here
template <typename T>
static void sqrtmQuasiTrianAuxEq(NDArray& A, NDArray& B, NDArray& C, NDArray& X) {
  std::vector<LongType> tempShape = {4,4};
  NDArray tempMatrix(A.ordering(),tempShape, A.dataType(), A.getContext());

  tempMatrix.r<T>(0, 0) = A.t<T>(0, 0) + B.t<T>(0, 0);
  tempMatrix.r<T>(1, 1) = A.t<T>(0, 0) + B.t<T>(1, 1);
  tempMatrix.r<T>(2, 2) = A.t<T>(1, 1) + B.t<T>(0, 0);
  tempMatrix.r<T>(3, 3) = A.t<T>(1, 1) + B.t<T>(1, 1);
  tempMatrix.r<T>(0, 1) = B.t<T>(1, 0);
  tempMatrix.r<T>(0, 2) = A.t<T>(0, 1);
  tempMatrix.r<T>(1, 0) = B.t<T>(0, 1);
  tempMatrix.r<T>(1, 3) = A.t<T>(0, 1);
  tempMatrix.r<T>(2, 0) = A.t<T>(1, 0);
  tempMatrix.r<T>(2, 3) = B.t<T>(1, 0);
  tempMatrix.r<T>(3, 1) = A.t<T>(1, 0);
  tempMatrix.r<T>(3, 2) = B.t<T>(0, 1);
  tempMatrix.r<T>(0, 3) = (T)0;
  tempMatrix.r<T>(1, 2) = (T)0;
  tempMatrix.r<T>(2, 1) = (T)0;
  tempMatrix.r<T>(3, 0) = (T)0;

  std::vector<LongType> resultShape = {4,1};
  NDArray result(A.ordering(), resultShape, A.dataType(), A.getContext());
  result.r<T>(0, 0) = C.t<T>(0, 0);
  result.r<T>(1, 0) = C.t<T>(0, 1);
  result.r<T>(2, 0) = C.t<T>(1, 0);
  result.r<T>(3, 0) = C.t<T>(1, 1);

  FullPivLU<T>::solve(tempMatrix, result, result);

  X.r<T>(0, 0) = result.t<T>(0);
  X.r<T>(0, 1) = result.t<T>(1);
  X.r<T>(1, 0) = result.t<T>(2);
  X.r<T>(1, 1) = result.t<T>(3);
}

//////////////////////////////////////////////////////////////////////////
template <typename T>
static void sqrtmQuasiTrianOffDiag(NDArray& matrixT, NDArray& sqrtT) {
  const int rows = matrixT.sizeAt(0);

  for (int j = 1; j < rows; j++) {
    if (matrixT.t<T>(j, j - 1) != (T)0) continue;

    for (int i = j - 1; i >= 0; i--) {
      if (i > 0 && matrixT.t<T>(i, i - 1) != (T)0) continue;

      const bool iBlockIs2x2 = (i < rows - 1) && (matrixT.t<T>(i + 1, i) != (T)0);
      const bool jBlockIs2x2 = (j < rows - 1) && (matrixT.t<T>(j + 1, j) != (T)0);

      if (iBlockIs2x2 && jBlockIs2x2) {
        NDArray *APtr = sqrtT({i, i + 2, i, i + 2}, true);
        NDArray A = *APtr;
        delete APtr;
        
        NDArray *BPtr = sqrtT({j, j + 2, j, j + 2}, true);
        NDArray B = *BPtr;
        delete BPtr;
        
        NDArray *XPtr = matrixT({i, i + 2, j, j + 2}, true);
        NDArray X = *XPtr;
        delete XPtr;

        if (j - i > 2) {
          NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
          NDArray *rightPtr = sqrtT({i + 2, j, j, j + 2}, true);
          auto mul = mmul(*leftPtr, *rightPtr);
          X -= *mul;
          delete leftPtr;
          delete rightPtr;
          delete mul;
        }

        sqrtmQuasiTrianAuxEq<T>(A, B, X, X);

        sqrtT.syncToDevice();
        NDArray *assignPtr = sqrtT({i, i + 2, j, j + 2}, true);
        assignPtr->assign(&X);
        delete assignPtr;
      } else if (iBlockIs2x2 && !jBlockIs2x2) {
        NDArray *rhsPtr = matrixT({i, i + 2, j, j + 1}, true);
        NDArray rhs = *rhsPtr;
        delete rhsPtr;

        if (j - i > 2) {
          NDArray *leftPtr = sqrtT({i, i + 2, i + 2, j}, true);
          NDArray *rightPtr = sqrtT({i + 2, j, j, j + 1}, true);
          auto mul = mmul(*leftPtr, *rightPtr);
          rhs -= *mul;
          delete leftPtr;
          delete rightPtr;
          delete mul;
        }

        std::vector<LongType> aShape = {2,2};
        NDArray A(matrixT.ordering(), aShape, matrixT.dataType(), matrixT.getContext());
        A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(j, j);
        A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
        
        NDArray *addPtr = sqrtT({i, i + 2, i, i + 2}, true);
        A += *addPtr;
        delete addPtr;

        FullPivLU<T>::solve(A, rhs, rhs);

        // sqrtT.syncToDevice();
        NDArray *assignPtr = sqrtT({i, i + 2, j, j + 1}, true);
        assignPtr->assign(&rhs);
        delete assignPtr;
      } else if (!iBlockIs2x2 && jBlockIs2x2) {
        NDArray *rhsPtr = matrixT({i, i + 1, j, j + 2}, true);
        NDArray rhs = *rhsPtr;
        delete rhsPtr;

        if (j - i > 1) {
          NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j}, true);
          NDArray *rightPtr = sqrtT({i + 1, j, j, j + 2}, true);
          auto mul =  mmul(*leftPtr, *rightPtr);
          rhs -= *mul;
          delete leftPtr;
          delete rightPtr;
          delete mul;
        }

        std::vector<LongType> aShape = {2,2};
        NDArray A(matrixT.ordering(),aShape, matrixT.dataType(), matrixT.getContext());
        A.r<T>(0, 0) = A.r<T>(1, 1) = sqrtT.t<T>(i, i);
        A.r<T>(0, 1) = A.r<T>(1, 0) = T(0);
        
        NDArray *addPtr = sqrtT({j, j + 2, j, j + 2}, true);
        NDArray *add = addPtr->transpose();
        delete addPtr;
        A += *add;
        delete add;

        NDArray *rhsT = rhs.transpose();
        FullPivLU<T>::solve(A, *rhsT, *rhsT);

        // sqrtT.syncToDevice();
        NDArray *assignPtr = sqrtT({i, i + 1, j, j + 2}, true);
        assignPtr->assign(&rhs);
        delete assignPtr;
        delete rhsT;
      } else if (!iBlockIs2x2 && !jBlockIs2x2) {
        NDArray *leftPtr = sqrtT({i, i + 1, i + 1, j});
        NDArray *rightPtr = sqrtT({i + 1, j, j, j + 1});
        auto mul = mmul(*leftPtr, *rightPtr);
        T temp = mul->t<T>(0);  // dot
        delete leftPtr;
        delete rightPtr;
        delete mul;
        
        sqrtT.r<T>(i, j) = (matrixT.t<T>(i, j) - temp) / (sqrtT.t<T>(i, i) + sqrtT.t<T>(j, j));
      }
    }
  }
}

//////////////////////////////////////////////////////////////////////////
template <typename T>
void Sqrtm<T>::calc(NDArray& in, NDArray& out) {
  if (in.rankOf() != 2 || in.sizeAt(0) != in.sizeAt(1))
    THROW_EXCEPTION("ops::helpers::Sqrtm::calc: input matrix must have rank 2 and be square !");
  if (!out.isSameShape(in))
    THROW_EXCEPTION("ops::helpers::Sqrtm::calc: output matrix must have the same shape as input one!");

  if (in.lengthOf() == 1) {
    out.r<T>(0) = math::sd_sqrt<T, T>(in.t<T>(0));
    return;
  }

  Schur<T> schur(in);


  NDArray *inULike = in.ulike();
  NDArray sqrtT = *inULike;
  sqrtT.nullify();

  sqrtmQuasiTrianDiag<T>(*schur.t, sqrtT);
  sqrtmQuasiTrianOffDiag<T>(*schur.t, sqrtT);

  NDArray *second = schur.u->transpose();
  // out = U * sqrtT * U^T;
  NDArray *temp = mmul(sqrtT, *second);
  MmulHelper::mmul(schur.u, temp, &out);
  delete inULike;
  delete second;
  delete temp;
}

BUILD_SINGLE_TEMPLATE( class Sqrtm, , SD_FLOAT_TYPES);

}  // namespace helpers
}  // namespace ops
}  // namespace sd