#pragma once

#include <Eigen/Eigen>
#include <EigenRand/EigenRand>
#include <utility>

namespace network {

enum class Score_Id { MSE, MAE, CrossEntropy };

struct Score_Database {
  using MatrixXd = Eigen::MatrixXd;
  using VectorXd = Eigen::VectorXd;

  static MatrixXd SoftMax(const MatrixXd &vec);

  template <Score_Id> static double score(const MatrixXd &, const MatrixXd &);

  template <Score_Id>
  static MatrixXd gradient(const MatrixXd &, const MatrixXd &);

  template <>
  inline double score<Score_Id::MSE>(const MatrixXd &x,
                                     const MatrixXd &reference) {
    return ((x - reference).transpose() * (x - reference)).trace() / x.cols();
  }

  template <>
  inline MatrixXd gradient<Score_Id::MSE>(const MatrixXd &x,
                                          const MatrixXd &reference) {
    return 2.0 * (x - reference);
  }

  template <>
  inline double score<Score_Id::MAE>(const MatrixXd &x,
                                     const MatrixXd &reference) {
    return (x - reference).array().abs().sum();
  }

  template <>
  inline MatrixXd gradient<Score_Id::MAE>(const MatrixXd &x,
                                          const MatrixXd &reference) {
    return (x - reference).unaryExpr([](double el) {
      return el > 0 ? 1.0 : -1.0;
    });
  }
  template <>
  inline double score<Score_Id::CrossEntropy>(const MatrixXd &x,
                                              const MatrixXd &reference) {
    double res = 0;
    for (int i = 0; i < x.cols(); ++i) {
      res += -(reference.col(i).transpose() / x.cols()) *
             VectorXd(SoftMax(x).col(i)).unaryExpr([](double el) {
               return log(el);
             });
    }
    return res;
  }
  template <>
  inline MatrixXd gradient<Score_Id::CrossEntropy>(const MatrixXd &x,
                                                   const MatrixXd &reference) {
    auto sf_x = SoftMax(x);
    auto sf_reference = SoftMax(reference);
    MatrixXd res(x.rows(), x.cols());
    double exp_sum;
    VectorXd const_vec;
    for (int i = 0; i < x.cols(); ++i) {
      exp_sum = VectorXd(sf_x.col(i)).array().exp().sum();
      const_vec = sf_x.unaryExpr([](double el) { return exp(el); }) -
                  exp_sum * VectorXd::Ones(x.size());
      res.col(i) = const_vec * sf_x / exp_sum;
    }
    return res;
  }
};

class Score_Func {

public:
  using MatrixXd = Eigen::MatrixXd;
  using VectorXd = Eigen::VectorXd;

private:
  using ScoreFuncType =
      std::function<double(const VectorXd &, const VectorXd &)>;
  using GradientFuncType =
      std::function<VectorXd(const VectorXd &, const VectorXd &)>;

public:
  Score_Func(ScoreFuncType score_func, GradientFuncType gradient_func);

  template <Score_Id Id> static Score_Func create() {
    return Score_Func(Score_Database::score<Id>, Score_Database::gradient<Id>);
  }

  static Score_Func create(Score_Id score);

  double score(const MatrixXd &x, const MatrixXd &reference) const;

  VectorXd gradient(const VectorXd &x, const VectorXd &reference) const;

private:
  ScoreFuncType score_func_;
  GradientFuncType gradient_func_;
};
} // namespace network