//
// Created by Denis Ryapolov on 06.01.2024.
//

#include "Network.h"
namespace network {
using VectorXd = Network::VectorXd;

Network::Network(std::initializer_list<int> dimensions,
                 std::initializer_list<Threshold_Id> threshold_id)
    : threshold_id_(threshold_id) {
  layers_.reserve(dimensions.size() - 1);
  auto dim_it = dimensions.begin();
  for (auto threshold_it = threshold_id.begin();
       threshold_it != threshold_id.end(); ++dim_it, ++threshold_it) {
    layers_.emplace_back(*threshold_it, *std::next(dim_it), *dim_it);
  }
}

Values Network::Forward_Prop(const MatrixXd &start_mat) {
  Values values;
  values.in.resize(layers_.size());
  values.out.resize(layers_.size());
  MatrixXd cur_mat = start_mat;
  for (size_t i = 0; i < layers_.size(); ++i) {
    values.in[i] = cur_mat;
    cur_mat = layers_[i].apply(cur_mat);
    values.out[i] = cur_mat;
    cur_mat = Threshold_Func::create(threshold_id_[i]).apply(cur_mat);
  }
  return values;
}

VectorXd Network::Back_Prop(const MatrixXd &start_vec,
                            const MatrixXd &reference,
                            const Score_Func &score_func, double step) {
  Values values = Forward_Prop(start_vec);
  const auto func = Threshold_Func::create(threshold_id_.back());
  VectorXd u = score_func.gradient(func.apply(values.out.back()), reference);
  for (int i = layers_.size() - 1; i >= 0; --i) {
    layers_[i].apply_gradA(layers_[i].gradA(values.in[i], u, values.out[i]),
                           step);
    layers_[i].apply_gradb(layers_[i].gradb(u, values.out[i]), step);
    u = layers_[i].gradx(u, values.out[i]);
  }
  return u;
}

VectorXd Network::Back_Prop_SGD(const MatrixXd &start_batch,
                                const MatrixXd &reference,
                                const Score_Func &score_func, int iter_num) {
  auto rand_ind = index_generator() % start_batch.cols();
  return Back_Prop(start_batch.col(rand_ind), reference.col(rand_ind),
                   score_func, 1.0 / iter_num);
}

void Network::TrainSGD(const MatrixXd &start_batch, const MatrixXd &reference,
                       const Score_Func &score_func, double needed_accuracy,
                       int max_epochs) {
  int epochs = 0;
  while (epochs != max_epochs &&
         abs(VectorXd::Ones(reference.rows()).transpose() *
             Back_Prop_SGD(start_batch, reference, score_func, epochs)) >
             needed_accuracy) {
    ++epochs;
  }
}

void Network::TrainBGD(const MatrixXd &start_batch, const MatrixXd &reference,
                        const Score_Func &score_func, int cols_in_minibatch,
                        double needed_accuracy, int max_epochs) {
  PermutationMatrix perm(start_batch.cols());
  perm.setIdentity();
  std::shuffle(perm.indices().data(),
               perm.indices().data() + perm.indices().size(), index_generator);
  MatrixXd mini_batch =
      start_batch * MatrixXd(perm).leftCols(cols_in_minibatch);
  MatrixXd mini_reference =
      reference * MatrixXd(perm).leftCols(cols_in_minibatch);
  int epochs = 0;
  while (epochs != max_epochs &&
         abs(VectorXd::Ones(reference.rows()).transpose() *
             Back_Prop(mini_batch, mini_reference, score_func, epochs)) >
             needed_accuracy) {
    ++epochs;
  }
}
VectorXd Network::Cols_Mean(const MatrixXd &x) {
  return x * VectorXd::Ones(x.cols()) / x.cols();
}
} // namespace network