#ifndef __KOKKOSBATCHED_TRSM_SERIAL_IMPL_HPP__
#define __KOKKOSBATCHED_TRSM_SERIAL_IMPL_HPP__


/// \author Kyungjoo Kim (kyukim@sandia.gov)

#include "KokkosBatched_Util.hpp"
#include "KokkosBatched_Trsm_Serial_Internal.hpp"

namespace KokkosBatched {

  ///
  /// L/L/NT
  ///
  /// B := inv(tril(A)) (alpha*B)
  /// A(m x m), B(m x n)

#if                                                     \
  defined(__KOKKOSBATCHED_INTEL_MKL__) &&               \
  defined(__KOKKOSBATCHED_INTEL_MKL_BATCHED__) &&       \
  defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)    
  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::NoTranspose,ArgDiag,Algo::Trsm::CompactMKL> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      typedef typename BViewType::value_type vector_type;
      //typedef typename vector_type::value_type value_type;
        
      const int
        m = B.extent(0),
        n = B.extent(1);

      static_assert(is_vector<vector_type>::value, "value type is not vector type");      
      static_assert(vector_type::vector_length == 4 || vector_type::vector_length == 8, 
                    "AVX, AVX2 and AVX512 is supported");
      const MKL_COMPACT_PACK format = vector_type::vector_length == 8 ?  MKL_COMPACT_AVX512 : MKL_COMPACT_AVX;
        
      // no error check
      int r_val = 0;
      if (A.stride_0() == 1 && B.stride_0() == 1) {
        mkl_dtrsm_compact(MKL_COL_MAJOR, 
                          MKL_LEFT, MKL_LOWER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_1(), 
                          (double*)B.data(), B.stride_1(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else if (A.stride_1() == 1 && B.stride_1() == 1) {
        mkl_dtrsm_compact(MKL_ROW_MAJOR, 
                          MKL_LEFT, MKL_LOWER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_0(), 
                          (double*)B.data(), B.stride_0(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else {
        r_val = -1;
      }
      return r_val;
    }
  };
#endif    

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::NoTranspose,ArgDiag,Algo::Trsm::Unblocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Unblocked>::invoke(ArgDiag::use_unit_diag,
                                                                        B.extent(0), B.extent(1),
                                                                        alpha, 
                                                                        A.data(), A.stride_0(), A.stride_1(),
                                                                        B.data(), B.stride_0(), B.stride_1());
    }
  };

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::NoTranspose,ArgDiag,Algo::Trsm::Blocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Blocked>::invoke(ArgDiag::use_unit_diag,
                                                                      B.extent(0), B.extent(1),
                                                                      alpha, 
                                                                      A.data(), A.stride_0(), A.stride_1(),
                                                                      B.data(), B.stride_0(), B.stride_1());
    }
  };

  ///
  /// R/U/NT
  ///
  /// B := (alpha*B) inv(triu(A))
  /// A(n x n), B(m x n)
#if                                                     \
  defined(__KOKKOSBATCHED_INTEL_MKL__) &&               \
  defined(__KOKKOSBATCHED_INTEL_MKL_BATCHED__) &&       \
  defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)
  template<typename ArgDiag>
  struct SerialTrsm<Side::Right,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::CompactMKL> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      typedef typename BViewType::value_type vector_type;
      //typedef typename vector_type::value_type value_type;
        
      const int
        m = B.extent(0),
        n = B.extent(1);

      static_assert(is_vector<vector_type>::value, "value type is not vector type");      
      static_assert(vector_type::vector_length == 4 || vector_type::vector_length == 8, 
                    "AVX, AVX2 and AVX512 is supported");
      const MKL_COMPACT_PACK format = vector_type::vector_length == 8 ?  MKL_COMPACT_AVX512 : MKL_COMPACT_AVX;

      // no error check
      int r_val = 0;
      if (A.stride_0() == 1 && B.stride_0() == 1) {
        mkl_dtrsm_compact(MKL_COL_MAJOR, 
                          MKL_RIGHT, MKL_UPPER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_1(), 
                          (double*)B.data(), B.stride_1(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else if (A.stride_1() == 1 && B.stride_1() == 1) {
        mkl_dtrsm_compact(MKL_ROW_MAJOR, 
                          MKL_RIGHT, MKL_UPPER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_0(), 
                          (double*)B.data(), B.stride_0(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else {
        r_val = -1;
      }
      return r_val;
    }
  };
#endif

  template<typename ArgDiag>
  struct SerialTrsm<Side::Right,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::Unblocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Unblocked>::invoke(ArgDiag::use_unit_diag,
                                                                        B.extent(1), B.extent(0),
                                                                        alpha, 
                                                                        A.data(), A.stride_1(), A.stride_0(),
                                                                        B.data(), B.stride_1(), B.stride_0());
    }
  };

  template<typename ArgDiag>
  struct SerialTrsm<Side::Right,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::Blocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Blocked>::invoke(ArgDiag::use_unit_diag,
                                                                      B.extent(1), B.extent(0),
                                                                      alpha, 
                                                                      A.data(), A.stride_1(), A.stride_0(),
                                                                      B.data(), B.stride_1(), B.stride_0());
    }      
  };
    
  ///
  /// L/U/NT
  ///
  /// B := inv(triu(A)) (alpha*B) 
  /// A(m x m), B(m x n)
#if                                                     \
  defined(__KOKKOSBATCHED_INTEL_MKL__) &&               \
  defined(__KOKKOSBATCHED_INTEL_MKL_BATCHED__) &&       \
  defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)
  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::CompactMKL> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      typedef typename BViewType::value_type vector_type;
      //typedef typename vector_type::value_type value_type;
        
      const int
        m = B.extent(0),
        n = B.extent(1);

      static_assert(is_vector<vector_type>::value, "value type is not vector type");      
      static_assert(vector_type::vector_length == 4 || vector_type::vector_length == 8, 
                    "AVX, AVX2 and AVX512 is supported");
      const MKL_COMPACT_PACK format = vector_type::vector_length == 8 ?  MKL_COMPACT_AVX512 : MKL_COMPACT_AVX;

      // no error check
      int r_val = 0;
      if (A.stride_0() == 1 && B.stride_0() == 1) {
        mkl_dtrsm_compact(MKL_COL_MAJOR, 
                          MKL_LEFT, MKL_UPPER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_1(), 
                          (double*)B.data(), B.stride_1(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else if (A.stride_1() == 1 && B.stride_1() == 1) {
        mkl_dtrsm_compact(MKL_ROW_MAJOR, 
                          MKL_LEFT, MKL_UPPER, MKL_NOTRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_0(), 
                          (double*)B.data(), B.stride_0(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else {
        r_val = -1;
      }
      return r_val;
    }
  };
#endif

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::Unblocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftUpper<Algo::Trsm::Unblocked>::invoke(ArgDiag::use_unit_diag,
                                                                        B.extent(0), B.extent(1),
                                                                        alpha, 
                                                                        A.data(), A.stride_0(), A.stride_1(),
                                                                        B.data(), B.stride_0(), B.stride_1());
    }
  };

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::NoTranspose,ArgDiag,Algo::Trsm::Blocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftUpper<Algo::Trsm::Blocked>::invoke(ArgDiag::use_unit_diag,
                                                                      B.extent(0), B.extent(1),
                                                                      alpha, 
                                                                      A.data(), A.stride_0(), A.stride_1(),
                                                                      B.data(), B.stride_0(), B.stride_1());
    }
  };

  ///
  /// L/L/T
  ///
  /// B := inv(tril(AT)) (alpha*B)
  /// A(m x m), B(m x n)

#if                                                     \
  defined(__KOKKOSBATCHED_INTEL_MKL__) &&               \
  defined(__KOKKOSBATCHED_INTEL_MKL_BATCHED__) &&       \
  defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)    
  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::Transpose,ArgDiag,Algo::Trsm::CompactMKL> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      typedef typename BViewType::value_type vector_type;
      //typedef typename vector_type::value_type value_type;
        
      const int
        m = B.extent(0),
        n = B.extent(1);

      static_assert(is_vector<vector_type>::value, "value type is not vector type");      
      static_assert(vector_type::vector_length == 4 || vector_type::vector_length == 8, 
                    "AVX, AVX2 and AVX512 is supported");
      const MKL_COMPACT_PACK format = vector_type::vector_length == 8 ?  MKL_COMPACT_AVX512 : MKL_COMPACT_AVX;
        
      // no error check
      int r_val = 0;
      if (A.stride_0() == 1 && B.stride_0() == 1) {
        mkl_dtrsm_compact(MKL_COL_MAJOR, 
                          MKL_LEFT, MKL_LOWER, MKL_TRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_1(), 
                          (double*)B.data(), B.stride_1(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else if (A.stride_1() == 1 && B.stride_1() == 1) {
        mkl_dtrsm_compact(MKL_ROW_MAJOR, 
                          MKL_LEFT, MKL_LOWER, MKL_TRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_0(), 
                          (double*)B.data(), B.stride_0(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else {
        r_val = -1;
      }
      return r_val;
    }
  };
#endif    

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::Transpose,ArgDiag,Algo::Trsm::Unblocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftUpper<Algo::Trsm::Unblocked>::invoke(ArgDiag::use_unit_diag,
                                                                        B.extent(0), B.extent(1),
                                                                        alpha, 
                                                                        A.data(), A.stride_1(), A.stride_0(),
                                                                        B.data(), B.stride_0(), B.stride_1());
    }
  };

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Lower,Trans::Transpose,ArgDiag,Algo::Trsm::Blocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftUpper<Algo::Trsm::Blocked>::invoke(ArgDiag::use_unit_diag,
                                                                      B.extent(0), B.extent(1),
                                                                      alpha, 
                                                                      A.data(), A.stride_1(), A.stride_0(),
                                                                      B.data(), B.stride_0(), B.stride_1());
    }
  };
  ///
  /// L/U/NT
  ///
  /// B := inv(triu(AT)) (alpha*B) 
  /// A(m x m), B(m x n)
#if                                                     \
  defined(__KOKKOSBATCHED_INTEL_MKL__) &&               \
  defined(__KOKKOSBATCHED_INTEL_MKL_BATCHED__) &&       \
  defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)
  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::Transpose,ArgDiag,Algo::Trsm::CompactMKL> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      typedef typename BViewType::value_type vector_type;
      //typedef typename vector_type::value_type value_type;
        
      const int
        m = B.extent(0),
        n = B.extent(1);

      static_assert(is_vector<vector_type>::value, "value type is not vector type");      
      static_assert(vector_type::vector_length == 4 || vector_type::vector_length == 8, 
                    "AVX, AVX2 and AVX512 is supported");
      const MKL_COMPACT_PACK format = vector_type::vector_length == 8 ?  MKL_COMPACT_AVX512 : MKL_COMPACT_AVX;

      // no error check
      int r_val = 0;
      if (A.stride_0() == 1 && B.stride_0() == 1) {
        mkl_dtrsm_compact(MKL_COL_MAJOR, 
                          MKL_LEFT, MKL_UPPER, MKL_TRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_1(), 
                          (double*)B.data(), B.stride_1(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else if (A.stride_1() == 1 && B.stride_1() == 1) {
        mkl_dtrsm_compact(MKL_ROW_MAJOR, 
                          MKL_LEFT, MKL_UPPER, MKL_TRANS, 
                          ArgDiag::use_unit_diag ? MKL_UNIT : MKL_NONUNIT,
                          m, n, 
                          alpha, 
                          (const double*)A.data(), A.stride_0(), 
                          (double*)B.data(), B.stride_0(), 
                          format, (MKL_INT)vector_type::vector_length);
      } else {
        r_val = -1;
      }
      return r_val;
    }
  };
#endif

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::Transpose,ArgDiag,Algo::Trsm::Unblocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Unblocked>::invoke(ArgDiag::use_unit_diag,
                                                                        B.extent(0), B.extent(1),
                                                                        alpha, 
                                                                        A.data(), A.stride_1(), A.stride_0(),
                                                                        B.data(), B.stride_0(), B.stride_1());
    }
  };

  template<typename ArgDiag>
  struct SerialTrsm<Side::Left,Uplo::Upper,Trans::Transpose,ArgDiag,Algo::Trsm::Blocked> {
    template<typename ScalarType,
             typename AViewType,
             typename BViewType>
    KOKKOS_INLINE_FUNCTION
    static int
    invoke(const ScalarType alpha,
           const AViewType &A,
           const BViewType &B) {
      return SerialTrsmInternalLeftLower<Algo::Trsm::Blocked>::invoke(ArgDiag::use_unit_diag,
                                                                      B.extent(0), B.extent(1),
                                                                      alpha, 
                                                                      A.data(), A.stride_1(), A.stride_0(),
                                                                      B.data(), B.stride_0(), B.stride_1());
    }
  };

}


#endif
