rotgen/include/rotgen/functions/reshaper.hpp

//==================================================================================================
/*
  ROTGEN - Runtime Overlay for Eigen
  Copyright : CODE RECKONS
  SPDX-License-Identifier: BSL-1.0
*/
//==================================================================================================
#pragma once

#include <rotgen/config.hpp>
#include <rotgen/functions/functions.hpp>

#include <type_traits>

namespace rotgen
{
  template<typename Ref> struct rowwise_adaptor
  {
    using concrete_type = typename std::remove_cvref_t<Ref>::concrete_type;
    Ref& target_;

    concrete_type _sum() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::sum(r); });
      return res;
    }

    concrete_type _mean() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::mean(r); });
      return res;
    }

    concrete_type _prod() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::prod(r); });
      return res;
    }

    concrete_type _maxCoeff() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::maxCoeff(r); });
      return res;
    }

    concrete_type _minCoeff() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::minCoeff(r); });
      return res;
    }

    concrete_type _squaredNorm() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::squaredNorm(r); });
      return res;
    }

    concrete_type _norm() const
    {
      concrete_type res(rows(target_), 1);
      apply([&](auto r, auto i) { res(i) = rotgen::norm(r); });
      return res;
    }

  private:
    template<typename Func> void apply(Func f)
    {
      for (Index i = 0; i < rows(target_); ++i) { f(row(target_, i), i); }
    }

    template<typename Func> void apply(Func f) const
    {
      for (Index i = 0; i < rows(target_); ++i) { f(row(target_, i), i); }
    }
  };

  template<typename Ref> struct colwise_adaptor
  {
    using concrete_type = typename std::remove_cvref_t<Ref>::concrete_type;
    Ref& target_;

    concrete_type _sum() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::sum(r); });
      return res;
    }

    concrete_type _mean() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::mean(r); });
      return res;
    }

    concrete_type _prod() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::prod(r); });
      return res;
    }

    concrete_type _maxCoeff() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::maxCoeff(r); });
      return res;
    }

    concrete_type _minCoeff() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::minCoeff(r); });
      return res;
    }

    concrete_type _squaredNorm() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::squaredNorm(r); });
      return res;
    }

    concrete_type _norm() const
    {
      concrete_type res(1, cols(target_));
      apply([&](auto r, auto i) { res(i) = rotgen::norm(r); });
      return res;
    }

  private:
    template<typename Func> void apply(Func f)
    {
      for (Index i = 0; i < cols(target_); ++i) { f(col(target_, i), i); }
    }

    template<typename Func> void apply(Func f) const
    {
      for (Index i = 0; i < cols(target_); ++i) { f(col(target_, i), i); }
    }
  };

  template<typename T> auto rowwise(T&& t)
  {
    if constexpr (use_expression_templates) { return t.base().rowwise(); }
    else
    {
      return rowwise_adaptor<T>{t};
    }
  }

  template<typename T> auto colwise(T&& t)
  {
    if constexpr (use_expression_templates) { return t.base().colwise(); }
    else
    {
      return colwise_adaptor<T>{t};
    }
  }
}