#include "LSM.h"
namespace DAA {

LSM::LSM(const vector<double>& X, const vector<double>& Y) : vX(X), vY(Y) {}

LSM::~LSM() {}

void LSM::reset(const vector<double>& X, const vector<double>& Y) {
  vX = arma::vec(X);
  vY = arma::vec(Y);
  polyfit_init();
}
void LSM::polyfit_init() {
  int m = vY.size();
  double Tss = m * arma::var(vY);
  double r2a_max = 0.4;
  for (int i = 0; i < max_orders; i++) {
    arma::vec res = arma::polyfit(vX, vY, i + 1);
    res.for_each([](arma::vec::elem_type& val) {
      val = limit_precision(val);
    });
    pnn[i] = res;
    fY = arma::polyval(pnn[i], vX);
    double Rssi = arma::accu(pow((fY - vY), 2));
    double r2i = 1 - (Rssi / Tss);
    r2a[i] = limit_precision(1 - (1 - r2i) * (m - 1) / (m - 2));
    if (r2a[i] > r2a_max) {
      order_best = i + 1;
      r2a_max = 1;
    }
  }
  is_init = true;
}

vector<double> LSM::polyfit(int orders) {
  if (order_now != orders) {
    pn = arma::polyfit(vX, vY, orders);
    order_now = orders;
  }
  return arma::conv_to<vector<double>>::from(pn);
}

vector<vector<double>> LSM::polyfit() {
  if (!is_init) {
    polyfit_init();
  }
  vector<vector<double>> pnn_v;
  for (int i = 0; i < max_orders; i++) {
    pnn_v.push_back(arma::conv_to<vector<double>>::from(pnn[i]));
  }
  return pnn_v;
}

vector<double> LSM::get_r2a() {
  if (!is_init) {
    polyfit_init();
  }
  return vector<double>(r2a.begin(), r2a.end());
}

vector<double> LSM::polyval(int orders) {
  if (order_now != orders) {
    pn = arma::polyfit(vX, vY, orders);
    order_now = orders;
  }
  fY = arma::polyval(pn, vX);

  return arma::conv_to<vector<double>>::from(fY);
}

double LSM::polyval(double x, int orders) {
  if (order_now != orders) {
    pn = arma::polyfit(vX, vY, orders);
    order_now = orders;
  }
  double res = 0;
  for (int i = 0; i < order_now + 1; i++) {
    res += std::pow(x, orders - i) * pn[i];
  }
  return res;
}

bool LSM::is_legal() { return vX.size() == vX.size() && !vX.empty(); }

double LSM::cor() {
  if (abs(cor_coef) > 1) {
    arma::vec res = arma::cor(vX, vY);
    cor_coef = res[0];
  }
  return cor_coef;
}

int LSM::get_order_best() {
  if (!is_init) {
    polyfit_init();
  }
  return this->order_best;
}

};  // namespace DAA