/*

Copyright (C) 2001, 2003, 2005, 2006, 2007, 2008 Ross Lippert and Paul Kienzle

This file is part of Octave.

Octave is free software; you can redistribute it and/or modify it
under the terms of the GNU General Public License as published by the
Free Software Foundation; either version 3 of the License, or (at your
option) any later version.

Octave is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
for more details.

You should have received a copy of the GNU General Public License
along with Octave; see the file COPYING.  If not, see
<http://www.gnu.org/licenses/>.

*/

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <float.h>

#include "CmplxSCHUR.h"
#include "fCmplxSCHUR.h"
#include "lo-ieee.h"
#include "lo-mappers.h"

#include "defun-dld.h"
#include "error.h"
#include "gripes.h"
#include "utils.h"

template <class T>
static inline T
getmin (T x, T y)
{
  return x < y ? x : y;
}

template <class T>
static inline T
getmax (T x, T y)
{
  return x > y ? x : y;
}

static double
frobnorm (const ComplexMatrix& A)
{
  double sum = 0;

  for (octave_idx_type i = 0; i < A.rows (); i++)
    for (octave_idx_type j = 0; j < A.columns (); j++)
      sum += real (A(i,j) * conj (A(i,j)));

  return sqrt (sum);
}

static double
frobnorm (const Matrix& A)
{
  double sum = 0;
  for (octave_idx_type i = 0; i < A.rows (); i++)
    for (octave_idx_type j = 0; j < A.columns (); j++)
      sum += A(i,j) * A(i,j);

  return sqrt (sum);
}

static float
frobnorm (const FloatComplexMatrix& A)
{
  float sum = 0;

  for (octave_idx_type i = 0; i < A.rows (); i++)
    for (octave_idx_type j = 0; j < A.columns (); j++)
      sum += real (A(i,j) * conj (A(i,j)));

  return sqrt (sum);
}

static float
frobnorm (const FloatMatrix& A)
{
  float sum = 0;
  for (octave_idx_type i = 0; i < A.rows (); i++)
    for (octave_idx_type j = 0; j < A.columns (); j++)
      sum += A(i,j) * A(i,j);

  return sqrt (sum);
}

static ComplexMatrix
sqrtm_from_schur (const ComplexMatrix& U, const ComplexMatrix& T)
{
  const octave_idx_type n = U.rows ();

  ComplexMatrix R (n, n, 0.0);

  for (octave_idx_type j = 0; j < n; j++)
    R(j,j) = sqrt (T(j,j));

  const double fudge = sqrt (DBL_MIN);

  for (octave_idx_type p = 0; p < n-1; p++)
    {
      for (octave_idx_type i = 0; i < n-(p+1); i++)
        {
          const octave_idx_type j = i + p + 1;

          Complex s = T(i,j);

          for (octave_idx_type k = i+1; k < j; k++)
            s -= R(i,k) * R(k,j);

          // dividing
          //     R(i,j) = s/(R(i,i)+R(j,j));
          // screwing around to not / 0

          const Complex d = R(i,i) + R(j,j) + fudge;
          const Complex conjd = conj (d);

          R(i,j) =  (s*conjd)/(d*conjd);
        }
    }

  return U * R * U.hermitian ();
}

static FloatComplexMatrix
sqrtm_from_schur (const FloatComplexMatrix& U, const FloatComplexMatrix& T)
{
  const octave_idx_type n = U.rows ();

  FloatComplexMatrix R (n, n, 0.0);

  for (octave_idx_type j = 0; j < n; j++)
    R(j,j) = sqrt (T(j,j));

  const float fudge = sqrt (FLT_MIN);

  for (octave_idx_type p = 0; p < n-1; p++)
    {
      for (octave_idx_type i = 0; i < n-(p+1); i++)
        {
          const octave_idx_type j = i + p + 1;

          FloatComplex s = T(i,j);

          for (octave_idx_type k = i+1; k < j; k++)
            s -= R(i,k) * R(k,j);

          // dividing
          //     R(i,j) = s/(R(i,i)+R(j,j));
          // screwing around to not / 0

          const FloatComplex d = R(i,i) + R(j,j) + fudge;
          const FloatComplex conjd = conj (d);

          R(i,j) =  (s*conjd)/(d*conjd);
        }
    }

  return U * R * U.hermitian ();
}

DEFUN_DLD (sqrtm, args, nargout,
 "-*- texinfo -*-\n\
@deftypefn {Loadable Function} {[@var{result}, @var{error_estimate}] =} sqrtm (@var{a})\n\
Compute the matrix square root of the square matrix @var{a}.\n\
\n\
Ref: Nicholas J. Higham.  A new sqrtm for @sc{matlab}.  Numerical Analysis\n\
Report No. 336, Manchester Centre for Computational Mathematics,\n\
Manchester, England, January 1999.\n\
@seealso{expm, logm}\n\
@end deftypefn")
{
  octave_value_list retval;

  int nargin = args.length ();

  if (nargin != 1)
    {
      print_usage ();
      return retval;
    }

  octave_value arg = args(0);

  octave_idx_type n = arg.rows ();
  octave_idx_type nc = arg.columns ();

  int arg_is_empty = empty_arg ("sqrtm", n, nc);

  if (arg_is_empty < 0)
    return retval;
  else if (arg_is_empty > 0)
    return octave_value (Matrix ());

  if (n != nc)
    {
      gripe_square_matrix_required ("sqrtm");
      return retval;
    }

  retval(1) = lo_ieee_inf_value ();
  retval(0) = lo_ieee_nan_value ();


  if (arg.is_single_type ())
    {
      if (arg.is_real_scalar ())
        {
          float d = arg.float_value ();
          if (d > 0.0)
            {
              retval(0) = sqrt (d);
              retval(1) = 0.0;
            }
          else
            {
              retval(0) = FloatComplex (0.0, sqrt (d));
              retval(1) = 0.0;
            }
        }
      else if (arg.is_complex_scalar ())
        {
          FloatComplex c = arg.float_complex_value ();
          retval(0) = sqrt (c);
          retval(1) = 0.0;
        }
      else if (arg.is_matrix_type ())
        {
          float err, minT;

          if (arg.is_real_matrix ())
            {
              FloatMatrix A = arg.float_matrix_value();

              if (error_state)
                return retval;

              // FIXME -- eventually, FloatComplexSCHUR will accept a
              // real matrix arg.

              FloatComplexMatrix Ac (A);

              const FloatComplexSCHUR schur (Ac, std::string ());

              if (error_state)
                return retval;

              const FloatComplexMatrix U (schur.unitary_matrix ());
              const FloatComplexMatrix T (schur.schur_matrix ());
              const FloatComplexMatrix X (sqrtm_from_schur (U, T));

              // Check for minimal imaginary part
              float normX = 0.0;
              float imagX = 0.0;
              for (octave_idx_type i = 0; i < n; i++)
                for (octave_idx_type j = 0; j < n; j++)
                  {
                    imagX = getmax (imagX, imag (X(i,j)));
                    normX = getmax (normX, abs (X(i,j)));
                  }

              if (imagX < normX * 100 * FLT_EPSILON)
                retval(0) = real (X);
              else
                retval(0) = X;

              // Compute error
              // FIXME can we estimate the error without doing the
              // matrix multiply?

              err = frobnorm (X*X - FloatComplexMatrix (A)) / frobnorm (A);

              if (xisnan (err))
                err = lo_ieee_float_inf_value ();

              // Find min diagonal
              minT = lo_ieee_float_inf_value ();
              for (octave_idx_type i=0; i < n; i++)
                minT = getmin(minT, abs(T(i,i)));
            }
          else
            {
              FloatComplexMatrix A = arg.float_complex_matrix_value ();

              if (error_state)
                return retval;

              const FloatComplexSCHUR schur (A, std::string ());

              if (error_state)
                return retval;

              const FloatComplexMatrix U (schur.unitary_matrix ());
              const FloatComplexMatrix T (schur.schur_matrix ());
              const FloatComplexMatrix X (sqrtm_from_schur (U, T));

              retval(0) = X;

              err = frobnorm (X*X - A) / frobnorm (A);

              if (xisnan (err))
                err = lo_ieee_float_inf_value ();

              minT = lo_ieee_float_inf_value ();
              for (octave_idx_type i = 0; i < n; i++)
                minT = getmin (minT, abs (T(i,i)));
            }

          retval(1) = err;

          if (nargout < 2)
            {
              if (err > 100*(minT+FLT_EPSILON)*n)
                {
                  if (minT == 0.0)
                    error ("sqrtm: A is singular, sqrt may not exist");
                  else if (minT <= sqrt (FLT_MIN))
                    error ("sqrtm: A is nearly singular, failed to find sqrt");
                  else
                    error ("sqrtm: failed to find sqrt");
                }
            }
        }
    }
  else
    {
      if (arg.is_real_scalar ())
        {
          double d = arg.double_value ();
          if (d > 0.0)
            {
              retval(0) = sqrt (d);
              retval(1) = 0.0;
            }
          else
            {
              retval(0) = Complex (0.0, sqrt (d));
              retval(1) = 0.0;
            }
        }
      else if (arg.is_complex_scalar ())
        {
          Complex c = arg.complex_value ();
          retval(0) = sqrt (c);
          retval(1) = 0.0;
        }
      else if (arg.is_matrix_type ())
        {
          double err, minT;

          if (arg.is_real_matrix ())
            {
              Matrix A = arg.matrix_value();

              if (error_state)
                return retval;

              // FIXME -- eventually, ComplexSCHUR will accept a
              // real matrix arg.

              ComplexMatrix Ac (A);

              const ComplexSCHUR schur (Ac, std::string ());

              if (error_state)
                return retval;

              const ComplexMatrix U (schur.unitary_matrix ());
              const ComplexMatrix T (schur.schur_matrix ());
              const ComplexMatrix X (sqrtm_from_schur (U, T));

              // Check for minimal imaginary part
              double normX = 0.0;
              double imagX = 0.0;
              for (octave_idx_type i = 0; i < n; i++)
                for (octave_idx_type j = 0; j < n; j++)
                  {
                    imagX = getmax (imagX, imag (X(i,j)));
                    normX = getmax (normX, abs (X(i,j)));
                  }

              if (imagX < normX * 100 * DBL_EPSILON)
                retval(0) = real (X);
              else
                retval(0) = X;

              // Compute error
              // FIXME can we estimate the error without doing the
              // matrix multiply?

              err = frobnorm (X*X - ComplexMatrix (A)) / frobnorm (A);

              if (xisnan (err))
                err = lo_ieee_inf_value ();

              // Find min diagonal
              minT = lo_ieee_inf_value ();
              for (octave_idx_type i=0; i < n; i++)
                minT = getmin(minT, abs(T(i,i)));
            }
          else
            {
              ComplexMatrix A = arg.complex_matrix_value ();

              if (error_state)
                return retval;

              const ComplexSCHUR schur (A, std::string ());

              if (error_state)
                return retval;

              const ComplexMatrix U (schur.unitary_matrix ());
              const ComplexMatrix T (schur.schur_matrix ());
              const ComplexMatrix X (sqrtm_from_schur (U, T));

              retval(0) = X;

              err = frobnorm (X*X - A) / frobnorm (A);

              if (xisnan (err))
                err = lo_ieee_inf_value ();

              minT = lo_ieee_inf_value ();
              for (octave_idx_type i = 0; i < n; i++)
                minT = getmin (minT, abs (T(i,i)));
            }

          retval(1) = err;

          if (nargout < 2)
            {
              if (err > 100*(minT+DBL_EPSILON)*n)
                {
                  if (minT == 0.0)
                    error ("sqrtm: A is singular, sqrt may not exist");
                  else if (minT <= sqrt (DBL_MIN))
                    error ("sqrtm: A is nearly singular, failed to find sqrt");
                  else
                    error ("sqrtm: failed to find sqrt");
                }
            }
        }
      else
        gripe_wrong_type_arg ("sqrtm", arg);
    }

  return retval;
}

/*
;;; Local Variables: ***
;;; mode: C++ ***
;;; End: ***
*/