/* =================================================================================================

(c - MIT) T.W.J. de Geus (Tom) | www.geus.me | github.com/tdegeus/GMatNonLinearElastic

================================================================================================= */

#ifndef GMATNONLINEARELASTIC_CARTESIAN3D_MATRIX_HPP
#define GMATNONLINEARELASTIC_CARTESIAN3D_MATRIX_HPP

#include "Cartesian3d.h"

namespace GMatNonLinearElastic {
namespace Cartesian3d {

// -------------------------------------------------------------------------------------------------

inline Matrix::Matrix(size_t nelem, size_t nip) : m_nelem(nelem), m_nip(nip)
{
  m_type   = xt::ones <size_t>({m_nelem, m_nip}) * Type::Unset;
  m_index  = xt::empty<size_t>({m_nelem, m_nip});
  m_allSet = false;
}

// -------------------------------------------------------------------------------------------------

inline Matrix::Matrix(size_t nelem, size_t nip, double kappa, double sig0, double eps0, double m) :
  m_nelem(nelem), m_nip(nip)
{
  m_type   = xt::ones<size_t>({m_nelem, m_nip}) * Type::NonLinearElastic;
  m_index  = xt::ones<size_t>({m_nelem, m_nip}) * m_NonLinearElastic.size();
  m_allSet = true;

  m_NonLinearElastic.push_back(NonLinearElastic(kappa, sig0, eps0, m));
}

// -------------------------------------------------------------------------------------------------

inline size_t Matrix::ndim() const
{
  return m_ndim;
}

// -------------------------------------------------------------------------------------------------

inline size_t Matrix::nelem() const
{
  return m_nelem;
}

// -------------------------------------------------------------------------------------------------

inline size_t Matrix::nip() const
{
  return m_nip;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<size_t,2> Matrix::type() const
{
  return m_type;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,2> Matrix::kappa() const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);

  xt::xtensor<double,2> out = xt::empty<double>({m_nelem, m_nip});
  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: out(e,q) = m_NonLinearElastic[m_index(e,q)].kappa(); break;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,2> Matrix::sig0() const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);

  xt::xtensor<double,2> out = xt::empty<double>({m_nelem, m_nip});
  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: out(e,q) = m_NonLinearElastic[m_index(e,q)].sig0(); break;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,2> Matrix::eps0() const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);

  xt::xtensor<double,2> out = xt::empty<double>({m_nelem, m_nip});
  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: out(e,q) = m_NonLinearElastic[m_index(e,q)].eps0(); break;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,2> Matrix::m() const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);

  xt::xtensor<double,2> out = xt::empty<double>({m_nelem, m_nip});
  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: out(e,q) = m_NonLinearElastic[m_index(e,q)].m(); break;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,4> Matrix::I2() const
{
  xt::xtensor<double,4> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor2 unit = Cartesian3d::I2();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0), xt::xshape<m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,6> Matrix::II() const
{
  xt::xtensor<double,6> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor4 unit = Cartesian3d::II();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,6> Matrix::I4() const
{
  xt::xtensor<double,6> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor4 unit = Cartesian3d::I4();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,6> Matrix::I4rt() const
{
  xt::xtensor<double,6> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor4 unit = Cartesian3d::I4rt();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,6> Matrix::I4s() const
{
  xt::xtensor<double,6> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor4 unit = Cartesian3d::I4s();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,6> Matrix::I4d() const
{
  xt::xtensor<double,6> out = xt::empty<double>({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim});
  #pragma omp parallel
  {
    Tensor4 unit = Cartesian3d::I4d();
    #pragma omp for
    for (size_t e = 0; e < m_nelem; ++e) {
      for (size_t q = 0; q < m_nip; ++q) {
        auto view = xt::adapt(&out(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
        xt::noalias(view) = unit;
      }
    }
  }
  return out;
}

// -------------------------------------------------------------------------------------------------

inline void Matrix::check() const
{
  if (xt::any(xt::equal(m_type, Type::Unset)))
    throw std::runtime_error("Points without material found");
}

// -------------------------------------------------------------------------------------------------

inline void Matrix::checkAllSet()
{
  if (xt::any(xt::equal(m_type, Type::Unset)))
    m_allSet = false;
  else
    m_allSet = true;
}

// -------------------------------------------------------------------------------------------------

inline void Matrix::setNonLinearElastic(
  const xt::xtensor<size_t,2>& I,
  double kappa,
  double sig0,
  double eps0,
  double m)
{
  GMATNONLINEARELASTIC_ASSERT(m_type.shape() == I.shape());
  GMATNONLINEARELASTIC_ASSERT(xt::all(xt::equal(I,0ul) || xt::equal(I,1ul)));
  GMATNONLINEARELASTIC_ASSERT(\
    xt::all(xt::equal(xt::where(xt::equal(I,1ul), m_type, Type::Unset), Type::Unset)));

  m_type = xt::where(xt::equal(I, 1ul), Type::NonLinearElastic, m_type);
  m_index = xt::where(xt::equal(I, 1ul), m_NonLinearElastic.size(), m_index);
  this->checkAllSet();
  m_NonLinearElastic.push_back(NonLinearElastic(kappa, sig0, eps0, m));
}

// -------------------------------------------------------------------------------------------------

inline void Matrix::stress(const xt::xtensor<double,4>& a_Eps, xt::xtensor<double,4>& a_Sig) const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);
  GMATNONLINEARELASTIC_ASSERT(a_Eps.shape() ==\
    std::decay_t<decltype(a_Eps)>::shape_type({m_nelem, m_nip, m_ndim, m_ndim}));
  GMATNONLINEARELASTIC_ASSERT(a_Eps.shape() == a_Sig.shape());

  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      auto Eps = xt::adapt(&a_Eps(e,q,0,0), xt::xshape<m_ndim,m_ndim>());
      auto Sig = xt::adapt(&a_Sig(e,q,0,0), xt::xshape<m_ndim,m_ndim>());
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: m_NonLinearElastic[m_index(e,q)].stress(Eps, Sig); break;
      }
    }
  }
}

// -------------------------------------------------------------------------------------------------

inline void Matrix::tangent(
  const xt::xtensor<double,4>& a_Eps,
        xt::xtensor<double,4>& a_Sig,
        xt::xtensor<double,6>& a_Tangent) const
{
  GMATNONLINEARELASTIC_ASSERT(m_allSet);
  GMATNONLINEARELASTIC_ASSERT(a_Eps.shape() == \
    std::decay_t<decltype(a_Eps)>::shape_type({m_nelem, m_nip, m_ndim, m_ndim}));
  GMATNONLINEARELASTIC_ASSERT(a_Eps.shape() == a_Sig.shape());
  GMATNONLINEARELASTIC_ASSERT(a_Tangent.shape() == \
    std::decay_t<decltype(a_Tangent)>::shape_type({m_nelem, m_nip, m_ndim, m_ndim, m_ndim, m_ndim}));

  #pragma omp parallel for
  for (size_t e = 0; e < m_nelem; ++e) {
    for (size_t q = 0; q < m_nip; ++q) {
      auto Eps = xt::adapt(&a_Eps(e,q,0,0), xt::xshape<m_ndim,m_ndim>());
      auto Sig = xt::adapt(&a_Sig(e,q,0,0), xt::xshape<m_ndim,m_ndim>());
      auto C = xt::adapt(&a_Tangent(e,q,0,0,0,0), xt::xshape<m_ndim,m_ndim,m_ndim,m_ndim>());
      switch (m_type(e,q)) {
        case Type::NonLinearElastic: m_NonLinearElastic[m_index(e,q)].tangent(Eps, Sig, C); break;
      }
    }
  }
}

// -------------------------------------------------------------------------------------------------

inline xt::xtensor<double,4> Matrix::Stress(const xt::xtensor<double,4>& Eps) const
{
  xt::xtensor<double,4> Sig = xt::empty<double>(Eps.shape());
  this->stress(Eps, Sig);
  return Sig;
}

// -------------------------------------------------------------------------------------------------

inline std::tuple<xt::xtensor<double,4>,xt::xtensor<double,6>> Matrix::Tangent(
  const xt::xtensor<double,4>& Eps) const
{
  xt::xtensor<double,4> Sig = xt::empty<double>({m_nelem,m_nip,m_ndim,m_ndim});
  xt::xtensor<double,6> C = xt::empty<double>({m_nelem,m_nip,m_ndim,m_ndim,m_ndim,m_ndim});
  this->tangent(Eps, Sig, C);
  return std::make_tuple(Sig, C);
}

// -------------------------------------------------------------------------------------------------

}} // namespace ...

#endif
