// SPDX-License-Identifier: Apache-2.0
// 
// Copyright 2008-2016 Conrad Sanderson (https://conradsanderson.id.au)
// Copyright 2008-2016 National ICT Australia (NICTA)
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// https://www.apache.org/licenses/LICENSE-2.0
// 
// 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.
// ------------------------------------------------------------------------


//! \addtogroup spop_diagmat
//! @{



template<typename T1>
inline
void
spop_diagmat::apply(SpMat<typename T1::elem_type>& out, const SpOp<T1, spop_diagmat>& in)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  if(in.is_alias(out) == false)
    {
    spop_diagmat::apply_noalias(out, in.m);
    }
  else
    {
    SpMat<eT> tmp;
    
    spop_diagmat::apply_noalias(tmp, in.m);
    
    out.steal_mem(tmp);
    }
  }



template<typename T1>
inline
void
spop_diagmat::apply_noalias(SpMat<typename T1::elem_type>& out, const SpBase<typename T1::elem_type, T1>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const SpProxy<T1> P(expr.get_ref());
  
  const uword P_n_rows = P.get_n_rows();
  const uword P_n_cols = P.get_n_cols();
  const uword P_n_nz   = P.get_n_nonzero();
  
  const bool P_is_vec = (P_n_rows == 1) || (P_n_cols == 1);
  
  if(P_is_vec)    // generate a diagonal matrix out of a vector
    {
    const uword N = (P_n_rows == 1) ? P_n_cols : P_n_rows;
    
    out.zeros(N, N);
    
    if(P_n_nz == 0)  { return; }
    
    typename SpProxy<T1>::const_iterator_type it = P.begin();
    
    if(P_n_cols == 1)
      {
      for(uword i=0; i < P_n_nz; ++i)
        {
        const uword row = it.row();
        
        out.at(row,row) = (*it);
        
        ++it;
        }
      }
    else
    if(P_n_rows == 1)
      {
      for(uword i=0; i < P_n_nz; ++i)
        {
        const uword col = it.col();
        
        out.at(col,col) = (*it);
        
        ++it;
        }
      }
    }
  else   // generate a diagonal matrix out of a matrix
    {
    out.zeros(P_n_rows, P_n_cols);
    
    const uword N = (std::min)(P_n_rows, P_n_cols);
    
    if( (is_SpMat<typename SpProxy<T1>::stored_type>::value) && (P_n_nz >= 5*N) )
      {
      const unwrap_spmat<typename SpProxy<T1>::stored_type> U(P.Q);
      
      const SpMat<eT>& X = U.M;
      
      for(uword i=0; i < N; ++i)
        {
        const eT val = X.at(i,i);  // use binary search
        
        if(val != eT(0))  { out.at(i,i) = val; }
        }
      }
    else
      {
      if(P_n_nz == 0)  { return; }
      
      typename SpProxy<T1>::const_iterator_type it = P.begin();
      
      for(uword i=0; i < P_n_nz; ++i)
        {
        const uword row = it.row();
        const uword col = it.col();
        
        if(row == col)  { out.at(row,row) = (*it); }
        
        ++it;
        }
      }
    }
  }



template<typename T1, typename T2>
inline
void
spop_diagmat::apply_noalias(SpMat<typename T1::elem_type>& out, const SpGlue<T1,T2,spglue_plus>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const unwrap_spmat<T1> UA(expr.A);
  const unwrap_spmat<T2> UB(expr.B);
  
  const SpMat<eT>& A = UA.M;
  const SpMat<eT>& B = UB.M;
  
  arma_conform_assert_same_size(A.n_rows, A.n_cols, B.n_rows, B.n_cols, "addition");
  
  const bool is_vec = (A.n_rows == 1) || (A.n_cols == 1);
  
  if(is_vec)    // generate a diagonal matrix out of a vector
    {
    const uword N = (A.n_rows == 1) ? A.n_cols : A.n_rows;
    
    out.zeros(N,N);
    
    if(A.n_rows == 1)
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(0,i) + B.at(0,i); }
      }
    else
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(i,0) + B.at(i,0); }
      }
    }
  else   // generate a diagonal matrix out of a matrix
    {
    SpMat<eT> AA;  spop_diagmat::apply_noalias(AA, A);
    SpMat<eT> BB;  spop_diagmat::apply_noalias(BB, B);
    
    out = AA + BB;
    }
  }



template<typename T1, typename T2>
inline
void
spop_diagmat::apply_noalias(SpMat<typename T1::elem_type>& out, const SpGlue<T1,T2,spglue_minus>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const unwrap_spmat<T1> UA(expr.A);
  const unwrap_spmat<T2> UB(expr.B);
  
  const SpMat<eT>& A = UA.M;
  const SpMat<eT>& B = UB.M;
  
  arma_conform_assert_same_size(A.n_rows, A.n_cols, B.n_rows, B.n_cols, "subtraction");
  
  const bool is_vec = (A.n_rows == 1) || (A.n_cols == 1);
  
  if(is_vec)    // generate a diagonal matrix out of a vector
    {
    const uword N = (A.n_rows == 1) ? A.n_cols : A.n_rows;
    
    out.zeros(N,N);
    
    if(A.n_rows == 1)
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(0,i) - B.at(0,i); }
      }
    else
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(i,0) - B.at(i,0); }
      }
    }
  else   // generate a diagonal matrix out of a matrix
    {
    SpMat<eT> AA;  spop_diagmat::apply_noalias(AA, A);
    SpMat<eT> BB;  spop_diagmat::apply_noalias(BB, B);
    
    out = AA - BB;
    }
  }



template<typename T1, typename T2>
inline
void
spop_diagmat::apply_noalias(SpMat<typename T1::elem_type>& out, const SpGlue<T1,T2,spglue_schur>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const unwrap_spmat<T1> UA(expr.A);
  const unwrap_spmat<T2> UB(expr.B);
  
  const SpMat<eT>& A = UA.M;
  const SpMat<eT>& B = UB.M;
  
  arma_conform_assert_same_size(A.n_rows, A.n_cols, B.n_rows, B.n_cols, "element-wise multiplication");
  
  const bool is_vec = (A.n_rows == 1) || (A.n_cols == 1);
  
  if(is_vec)    // generate a diagonal matrix out of a vector
    {
    const uword N = (A.n_rows == 1) ? A.n_cols : A.n_rows;
    
    out.zeros(N,N);
    
    if(A.n_rows == 1)
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(0,i) * B.at(0,i); }
      }
    else
      {
      for(uword i=0; i < N; ++i) { out.at(i,i) = A.at(i,0) * B.at(i,0); }
      }
    }
  else   // generate a diagonal matrix out of a matrix
    {
    SpMat<eT> AA;  spop_diagmat::apply_noalias(AA, A);
    SpMat<eT> BB;  spop_diagmat::apply_noalias(BB, B);
    
    out = AA % BB;
    }
  }



template<typename T1, typename T2>
inline
void
spop_diagmat::apply_noalias(SpMat<typename T1::elem_type>& out, const SpGlue<T1,T2,spglue_times>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const unwrap_spmat<T1> UA(expr.A);
  const unwrap_spmat<T2> UB(expr.B);
  
  const SpMat<eT>& A = UA.M;
  const SpMat<eT>& B = UB.M;
  
  arma_conform_assert_mul_size(A.n_rows, A.n_cols, B.n_rows, B.n_cols, "matrix multiplication");
  
  const uword C_n_rows = A.n_rows;
  const uword C_n_cols = B.n_cols;
  
  const bool is_vec = (C_n_rows == 1) || (C_n_cols == 1);
  
  if(is_vec)    // generate a diagonal matrix out of a vector
    {
    const SpMat<eT> C = A*B;
    
    spop_diagmat::apply_noalias(out, C);
    }
  else   // generate a diagonal matrix out of a matrix
    {
    const uword N = (std::min)(C_n_rows, C_n_cols);
    
    if( (A.n_nonzero >= 5*N) || (B.n_nonzero >= 5*N) )
      {
      out.zeros(C_n_rows, C_n_cols);
      
      for(uword k=0; k < N; ++k)
        {
        typename SpMat<eT>::const_col_iterator B_it     = B.begin_col_no_sync(k);
        typename SpMat<eT>::const_col_iterator B_it_end = B.end_col_no_sync(k);
        
        eT acc = eT(0);
        
        while(B_it != B_it_end)
          {
          const eT    B_val = (*B_it);
          const uword i     = B_it.row();
          
          acc += A.at(k,i) * B_val;
          
          ++B_it;
          }
        
        out(k,k) = acc;
        }
      }
    else
      {
      const SpMat<eT> C = A*B;
      
      spop_diagmat::apply_noalias(out, C);
      }
    }
  }



//
//



template<typename T1>
inline
void
spop_diagmat2::apply(SpMat<typename T1::elem_type>& out, const SpOp<T1, spop_diagmat2>& in)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  const uword row_offset = in.aux_uword_a;
  const uword col_offset = in.aux_uword_b;
  
  const unwrap_spmat<T1> U(in.m);
  
  if(U.is_alias(out))
    {
    SpMat<eT> tmp;
    
    spop_diagmat2::apply_noalias(tmp, U.M, row_offset, col_offset);
    
    out.steal_mem(tmp);
    }
  else
    {
    spop_diagmat2::apply_noalias(out, U.M, row_offset, col_offset);
    }
  }



template<typename eT>
inline
void
spop_diagmat2::apply_noalias(SpMat<eT>& out, const SpMat<eT>& X, const uword row_offset, const uword col_offset)
  {
  arma_debug_sigprint();
  
  const uword n_rows = X.n_rows;
  const uword n_cols = X.n_cols;
  const uword n_elem = X.n_elem;
  
  if(n_elem == 0)  { out.reset(); return; }
  
  const bool X_is_vec = (n_rows == 1) || (n_cols == 1);
  
  if(X_is_vec)    // generate a diagonal matrix out of a vector
    {
    const uword n_pad = (std::max)(row_offset, col_offset);
    
    out.zeros(n_elem + n_pad, n_elem + n_pad);
    
    const uword X_n_nz = X.n_nonzero;
    
    if(X_n_nz == 0)  { return; }
    
    typename SpMat<eT>::const_iterator it = X.begin();
      
    if(n_cols == 1)
      {
      for(uword i=0; i < X_n_nz; ++i)
        {
        const uword row = it.row();
        
        out.at(row_offset + row, col_offset + row) = (*it);
        
        ++it;
        }
      }
    else
    if(n_rows == 1)
      {
      for(uword i=0; i < X_n_nz; ++i)
        {
        const uword col = it.col();
        
        out.at(row_offset + col, col_offset + col) = (*it);
        
        ++it;
        }
      }
    }
  else   // generate a diagonal matrix out of a matrix
    {
    arma_conform_check_bounds
      (
      ((row_offset > 0) && (row_offset >= n_rows)) || ((col_offset > 0) && (col_offset >= n_cols)),
      "diagmat(): requested diagonal out of bounds"
      );
    
    out.zeros(n_rows, n_cols);
    
    if(X.n_nonzero == 0)  { return; }
    
    const uword N = (std::min)(n_rows - row_offset, n_cols - col_offset);
    
    for(uword i=0; i<N; ++i)
      {
      const uword row = i + row_offset;
      const uword col = i + col_offset;
      
      const eT val = X.at(row,col);
      
      if(val != eT(0))  { out.at(row,col) = val; }
      }
    }
  }



//! @}