/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA 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.
 *
 * This program 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 MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <cmath>
#include <cassert>
#include <cstring>
#include <iostream>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_blas.h>

#include <mia/core/gsl_matrix.hh>
#include <mia/core/gsl_matrix_vector_ops.hh>
#include <algorithm>

namespace gsl
{

using std::swap;
using std::fill;
using std::ostream_iterator;
using std::sqrt;

Matrix::Matrix(): m_matrix(nullptr), m_const_matrix(nullptr), m_owner(false)
{
}

bool Matrix::is_valid() const
{
       return m_const_matrix != nullptr;
}


bool Matrix::is_writable() const
{
       return m_matrix != nullptr;
}

Matrix::Matrix(size_t rows, size_t columns, bool clean):
       m_matrix(NULL), m_owner(true)
{
       m_matrix = clean ?
                  gsl_matrix_calloc(rows, columns) :
                  gsl_matrix_alloc(rows, columns);
       m_const_matrix = m_matrix;
}

Matrix::Matrix(size_t rows, size_t columns, double init):
       m_matrix(NULL), m_owner(true)
{
       m_matrix = gsl_matrix_alloc(rows, columns);
       auto p = gsl_matrix_ptr (m_matrix, 0, 0);
       fill(p, p + rows * columns, init);
       m_const_matrix = m_matrix;
}

Matrix::Matrix(size_t rows, size_t columns, const double *init):
       m_matrix(NULL), m_owner(true)
{
       assert(init);
       m_matrix = gsl_matrix_alloc(rows, columns);
       memcpy(gsl_matrix_ptr (m_matrix, 0, 0), init, rows * columns * sizeof(double));
       m_const_matrix = m_matrix;
}

Matrix::Matrix(const Matrix& other):
       m_owner(true)
{
       m_matrix = gsl_matrix_alloc(other.rows(), other.cols());
       gsl_matrix_memcpy (m_matrix, other.m_matrix);
       m_const_matrix = m_matrix;
}

Matrix::Matrix(Matrix&& other):
       m_matrix(other.m_matrix),
       m_const_matrix(other.m_const_matrix),
       m_owner(other.m_owner)
{
       other.m_const_matrix = nullptr;
       other.m_matrix = nullptr;
       other.m_owner = false;
}



Matrix::Matrix(gsl_matrix *m):
       m_matrix(m),
       m_const_matrix(m),
       m_owner(false)
{
}

Matrix::Matrix(const gsl_matrix *m):
       m_matrix(nullptr),
       m_const_matrix(m),
       m_owner(false)
{
}

Matrix& Matrix::operator =(Matrix&& other)
{
       if (this == &other)
              return *this;

       m_const_matrix = other.m_const_matrix;
       m_matrix = other.m_matrix;
       m_owner = other.m_owner;
       other.m_const_matrix = nullptr;
       other.m_matrix = nullptr;
       other.m_owner = false;
       return *this;
}

Matrix& Matrix::operator =(const Matrix& other)
{
       if (this == &other)
              return *this;

       if (m_matrix && rows() == other.rows() && cols() == other.cols()) {
              gsl_matrix_memcpy (m_matrix, other.m_const_matrix);
              return *this;
       }

       gsl_matrix *help = gsl_matrix_alloc(other.rows(), other.cols());
       gsl_matrix_memcpy (help, other.m_matrix);
       swap(m_matrix, help);

       if (help && m_owner)
              gsl_matrix_free(help);

       m_const_matrix = m_matrix;
       m_owner = true;
       return *this;
}

void Matrix::reset(size_t rows, size_t columns, bool clean)
{
       gsl_matrix *help = clean ?
                          gsl_matrix_calloc(rows, columns) :
                          gsl_matrix_alloc(rows, columns);
       swap(help, m_matrix);

       if (help && m_owner)
              gsl_matrix_free(help);

       m_const_matrix = m_matrix;
       m_owner = true;
}

void Matrix::reset(size_t rows, size_t columns, double init)
{
       gsl_matrix *help = gsl_matrix_alloc(rows, columns);
       auto p = gsl_matrix_ptr (help, 0, 0);
       fill(p, p + rows * columns, init);
       swap(help, m_matrix);

       if (help && m_owner)
              gsl_matrix_free(help);

       m_const_matrix = m_matrix;
       m_owner = true;
}

Matrix::~Matrix()
{
       if (m_matrix && m_owner)
              gsl_matrix_free(m_matrix);
}

size_t Matrix::rows()const
{
       assert(m_const_matrix);
       return m_const_matrix->size1;
}

size_t Matrix::cols()const
{
       assert(m_const_matrix);
       return m_const_matrix->size2;
}

void Matrix::set_row(int r, const Vector& row)
{
       assert(row.size() == cols());
       auto mrow = gsl_matrix_row(m_matrix, r);
       gsl_vector_memcpy(&mrow.vector, row);
}

VectorView Matrix::get_row(int r)
{
       return VectorView(gsl_matrix_row(m_matrix, r));
}

ConstVectorView Matrix::get_row(int r) const
{
       return ConstVectorView(gsl_matrix_const_row(m_const_matrix, r));
}

void Matrix::set_column(int c, const Vector& col)
{
       assert(col.size() == rows());
       auto mcol = gsl_matrix_column(m_matrix, c);
       gsl_vector_memcpy(&mcol.vector, col);
}

VectorView Matrix::get_column(int c)
{
       return VectorView(gsl_matrix_column(m_matrix, c));
}

ConstVectorView Matrix::get_column(int c) const
{
       return ConstVectorView(gsl_matrix_const_column(m_const_matrix, c));
}



double Matrix::dot_row(int r, const Vector& row) const
{
       auto mrow = gsl_matrix_const_row(m_const_matrix, r);
       return dot(row, &mrow.vector);
}

double Matrix::dot_column(int c, const Vector& col) const
{
       auto mcol = gsl_matrix_const_column(m_const_matrix, c);
       return dot(col, &mcol.vector);
}

void Matrix::print(std::ostream& os) const
{
       os << "[";

       if (!m_matrix) {
              if (!m_const_matrix) {
                     os << "(null)]";
                     return;
              }

              os << "(const)";
       }

       os << "\n";

       for (unsigned r = 0; r < rows(); ++r) {
              auto mrow = get_row(r);
              os << "  ";
              copy(mrow.begin(), mrow.end(), ostream_iterator<double>(os, ", "));
              os << "\n";
       }

       os << "]";
}

void Matrix::set(size_t i, size_t j, double x)
{
       assert(m_matrix);
       gsl_matrix_set(m_matrix, i, j, x);
}

double Matrix::operator ()(size_t i, size_t j) const
{
       return gsl_matrix_get(m_const_matrix, i, j);
}

Matrix::operator gsl_matrix *()
{
       return m_matrix;
}

Matrix::operator const gsl_matrix *() const
{
       return m_const_matrix;
}


matrix_iterator Matrix::begin()
{
       assert(m_matrix);
       return matrix_iterator(m_matrix, true);
}

matrix_iterator Matrix::end()
{
       assert(m_matrix);
       return matrix_iterator(m_matrix, false);
}


const_matrix_iterator Matrix::begin() const
{
       assert(m_const_matrix);
       return const_matrix_iterator(m_const_matrix, true);
}

const_matrix_iterator Matrix::end() const
{
       assert(m_const_matrix);
       return const_matrix_iterator(m_const_matrix, false);
}

Matrix Matrix::transposed() const
{
       assert(m_matrix);
       Matrix result(cols(), rows(), false);
       gsl_matrix_transpose_memcpy (result.m_matrix, m_matrix);
       return result;
}

Matrix Matrix::row_covariance() const
{
       int d = rows();
       int n = cols();
       Matrix cov(d, d, true);
       Matrix help(d, n, false);

       // remove mean
       for (int r = 0; r < d; ++r)  {
              auto tmp = gsl_matrix_const_row (m_const_matrix, r);
              auto mean = gsl_stats_mean (tmp.vector.data, tmp.vector.stride, n);

              for (int c = 0; c < n; ++c)  {
                     help.set(r, c, gsl_matrix_get(m_const_matrix, r, c) - mean);
              }
       }

       for (int i = 0; i < d; i++) {
              for (int j = 0; j <= i; j++) {
                     auto a = gsl_matrix_row (help, i);
                     auto b = gsl_matrix_row (help, j);
                     double c = gsl_stats_covariance(a.vector.data, a.vector.stride,
                                                     b.vector.data, b.vector.stride, n);
                     cov.set(j, i, c);
                     cov.set(i, j, c);
              }
       }

       return cov;
}

Matrix Matrix::column_covariance() const
{
       int n = rows();
       int d = cols();
       Matrix cov(d, d, true);
       Matrix help(n, d, false);

       // remove mean
       for (int i = 0; i < d; ++i)  {
              auto tmp = gsl_matrix_const_column (m_const_matrix, i);
              auto mean = gsl_stats_mean (tmp.vector.data, tmp.vector.stride, n);

              for (int r = 0; r < n; ++r)  {
                     help.set(r, i, gsl_matrix_get(m_const_matrix, r, i) - mean);
              }
       }

       for (int i = 0; i < d; i++) {
              for (int j = 0; j <= i; j++) {
                     auto a = gsl_matrix_column (help, i);
                     auto b = gsl_matrix_column (help, j);
                     double c = gsl_stats_covariance(a.vector.data, a.vector.stride,
                                                     b.vector.data, b.vector.stride, n);
                     cov.set(j, i, c);
                     cov.set(i, j, c);
              }
       }

       return cov;
}


Matrix operator - (const Matrix& lhs, const Matrix& rhs)
{
       Matrix result(lhs);
       result -= rhs;
       return result;
}

Matrix operator + (const Matrix& lhs, const Matrix& rhs)
{
       Matrix result(lhs);
       result += rhs;
       return result;
}

Matrix operator * (const Matrix& lhs, const Matrix& rhs)
{
       Matrix result(lhs.rows(), rhs.cols(), false);
       multiply_m_m(result, lhs, rhs);
       return result;
}

CSymmvEvalEvec::CSymmvEvalEvec(Matrix m):
       evec(m.rows(), m.cols(), false),
       eval(m.rows(), false)
{
       assert(m.cols() == m.rows());
       gsl_eigen_symmv_workspace *ws = gsl_eigen_symmv_alloc (m.rows());
       gsl_eigen_symmv (m, eval, evec, ws);
       gsl_eigen_symmv_free (ws);
}

void matrix_inv_sqrt(Matrix& m )
{
       Matrix mTm(m.cols(), m.cols(), false);
       multiply_mT_m(mTm, m, m);
       CSymmvEvalEvec see(mTm);

       for (unsigned r = 0; r < see.evec.cols(); ++r) {
              auto wmr = see.evec.get_column(r);
              const double f = see.eval[r] > 0 ? 1.0 / sqrt(sqrt(see.eval[r])) : 0.0;
              std::transform(wmr.begin(), wmr.end(), wmr.begin(), [f](double x) {
                     return f * x;
              });
       }

       multiply_m_mT(mTm, see.evec, see.evec);
       m = m * mTm;
}


bool operator == (const matrix_iterator& lhs, const matrix_iterator& rhs)
{
       assert(lhs.m_matrix == rhs.m_matrix);
       return lhs.m_current == rhs.m_current;
}

bool operator != (const matrix_iterator& lhs, const matrix_iterator& rhs)
{
       assert(lhs.m_matrix == rhs.m_matrix);
       return lhs.m_current != rhs.m_current;
}

bool operator == (const const_matrix_iterator& lhs, const const_matrix_iterator& rhs)
{
       assert(lhs.m_matrix == rhs.m_matrix);
       return lhs.m_current == rhs.m_current;
}

bool operator != (const const_matrix_iterator& lhs, const const_matrix_iterator& rhs)
{
       assert(lhs.m_matrix == rhs.m_matrix);
       return lhs.m_current != rhs.m_current;
}


}
