// 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 op_powmat
//! @{


template<typename T1>
inline
void
op_powmat::apply(Mat<typename T1::elem_type>& out, const Op<T1, op_powmat>& expr)
  {
  arma_debug_sigprint();
  
  const uword y     =  expr.aux_uword_a;
  const bool  y_neg = (expr.aux_uword_b == uword(1));
  
  const bool status = op_powmat::apply_direct(out, expr.m, y, y_neg);
  
  if(status == false)
    {
    out.soft_reset();
    arma_stop_runtime_error("powmat(): transformation failed");
    }
  }



template<typename T1>
inline
bool
op_powmat::apply_direct(Mat<typename T1::elem_type>& out, const Base<typename T1::elem_type,T1>& X, const uword y, const bool y_neg)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type eT;
  
  if(y_neg)
    {
    if(y == uword(1))
      {
      return op_inv_gen_default::apply_direct(out, X.get_ref(), "powmat()");
      }
    else
      {
      Mat<eT> X_inv;
      
      const bool inv_status = op_inv_gen_default::apply_direct(X_inv, X.get_ref(), "powmat()");
      
      if(inv_status == false)  { return false; }
      
      op_powmat::apply_direct_positive(out, X_inv, y);
      }
    }
  else
    {
    const quasi_unwrap<T1> U(X.get_ref());
    
    arma_conform_check( (U.M.is_square() == false), "powmat(): given matrix must be square sized" );
    
    op_powmat::apply_direct_positive(out, U.M, y);
    }
  
  return true;
  }



template<typename eT>
inline
void
op_powmat::apply_direct_positive(Mat<eT>& out, const Mat<eT>& X, const uword y)
  {
  arma_debug_sigprint();
  
  const uword N = X.n_rows;
  
  if(y == uword(0))  { out.eye(N,N); return; }
  if(y == uword(1))  { out = X;      return; }
  
  if(X.is_diagmat())
    {
    arma_debug_print("op_powmat: diag optimisation");
    
    podarray<eT> tmp(N);  // use temporary array in case we have aliasing
    
    for(uword i=0; i<N; ++i)  { tmp[i] = eop_aux::pow(X.at(i,i), int(y)); }
    
    out.zeros(N,N);
    
    for(uword i=0; i<N; ++i)  { out.at(i,i) = tmp[i]; }
    }
  else
    {
         if(y == uword(2))  {                          out = X*X;       }
    else if(y == uword(3))  { const Mat<eT> tmp = X*X; out = X*tmp;     }
    else if(y == uword(4))  { const Mat<eT> tmp = X*X; out =   tmp*tmp; }
    else if(y == uword(5))  { const Mat<eT> tmp = X*X; out = X*tmp*tmp; }
    else
      {
      Mat<eT> tmp = X;
      
      out = X;
      
      uword z = y-1;
      
      while(z > 0)
        {
        if(z & 1)  { out = tmp * out; }
        
        z /= uword(2);
        
        if(z > 0)  { tmp = tmp * tmp; }
        }
      }
    }
  }



template<typename T1>
inline
void
op_powmat_cx::apply(Mat< std::complex<typename T1::pod_type> >& out, const mtOp<std::complex<typename T1::pod_type>,T1,op_powmat_cx>& expr)
  {
  arma_debug_sigprint();
  
  typedef typename T1::pod_type in_T;
  
  const in_T y = std::real(expr.aux_out_eT);
  
  const bool status = op_powmat_cx::apply_direct(out, expr.m, y);
  
  if(status == false)
    {
    out.soft_reset();
    arma_stop_runtime_error("powmat(): transformation failed");
    }
  }



template<typename T1>
inline
bool
op_powmat_cx::apply_direct(Mat< std::complex<typename T1::pod_type> >& out, const Base<typename T1::elem_type,T1>& X, const typename T1::pod_type y)
  {
  arma_debug_sigprint();
  
  typedef typename T1::elem_type in_eT;
  typedef typename T1::pod_type  in_T;
  typedef std::complex<in_T>     out_eT;
  
  if( y == in_T(int(y)) )
    {
    arma_debug_print("op_powmat_cx::apply_direct(): integer exponent detected; redirecting to op_powmat");
    
    const uword y_val = (y < int(0)) ? uword(-y) : uword(y);
    const bool  y_neg = (y < int(0));
    
    Mat<in_eT> tmp;
    
    const bool status = op_powmat::apply_direct(tmp, X.get_ref(), y_val, y_neg);
    
    if(status == false)  { return false; }
    
    out = conv_to< Mat<out_eT> >::from(tmp);
    
    return true;
    }
  
  const quasi_unwrap<T1> U(X.get_ref());
  const Mat<in_eT>& A  = U.M;
  
  arma_conform_check( (A.is_square() == false), "powmat(): given matrix must be square sized" );
  
  const uword N = A.n_rows;
  
  if(A.is_diagmat())
    {
    arma_debug_print("op_powmat_cx: diag optimisation");
    
    podarray<out_eT> tmp(N);  // use temporary array in case we have aliasing
    
    for(uword i=0; i<N; ++i)  { tmp[i] = eop_aux::pow( std::complex<in_T>(A.at(i,i)), y ); }
    
    out.zeros(N,N);
    
    for(uword i=0; i<N; ++i)  { out.at(i,i) = tmp[i]; }
    
    return true;
    }
  
  const bool try_sym = arma_config::optimise_sym && sym_helper::is_approx_sym(A);
  
  if(try_sym)
    {
    arma_debug_print("op_powmat_cx: symmetric/hermitian optimisation");
    
    Col<in_T>  eigval;
    Mat<in_eT> eigvec;
    
    const bool eig_status = eig_sym(eigval, eigvec, A);
    
    if(eig_status)
      {
      bool all_pos = true;
      
      for(uword i=0; i<N; ++i)  { all_pos = (eigval[i] <= in_T(0)) ? false : all_pos; }
      
      if(all_pos)
        {
        arma_debug_print("op_powmat_cx: all_pos = true");
        
        eigval = pow(eigval, y);
        
        const Mat<in_eT> tmp = eigvec * diagmat(eigval);
        
        out = conv_to< Mat<out_eT> >::from(tmp  * eigvec.t());
        }
      else
        {
        arma_debug_print("op_powmat_cx: all_pos = false");
        
        Col<out_eT> cx_eigval_pow(N, arma_nozeros_indicator());
        
        for(uword i=0; i<N; ++i)  { cx_eigval_pow[i] = eop_aux::pow( std::complex<in_T>(eigval[i]), y ); }
        
        const Mat<out_eT> cx_eigvec = conv_to< Mat<out_eT> >::from(eigvec);
        
        const Mat<out_eT> tmp = cx_eigvec * diagmat(cx_eigval_pow);
        
        out = tmp * cx_eigvec.t();
        }
      
      return true;
      }
    
    arma_debug_print("op_powmat_cx: symmetric/hermitian optimisation failed");
    
    // fallthrough if optimisation failed
    }
  
  bool powmat_status = false;
  
  Col<out_eT> eigval;
  Mat<out_eT> eigvec;
  
  const bool eig_status = eig_gen(eigval, eigvec, A);
  
  if(eig_status)
    {
    eigval = pow(eigval, y);
    
    Mat<out_eT> eigvec_t = trans(eigvec);
    Mat<out_eT> tmp      = diagmat(conj(eigval)) * eigvec_t;
    
    const bool solve_status = auxlib::solve_square_fast(out, eigvec_t, tmp);
    
    if(solve_status)  { out = trans(out); powmat_status = true; }
    }
  
  return powmat_status;
  }



//! @}
