#include <RcppArmadillo.h>

using namespace arma;
using namespace Rcpp;

// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::export]]
List l0araC(arma::mat x, arma::vec y, String family, double lam, int maxit, double eps) {
  const arma::uword n = x.n_rows;
  const arma::uword m = x.n_cols;
  arma::vec w(m, fill::zeros);
  arma::vec old_w(m, fill::zeros);
  arma::vec eta(n, fill::zeros);
  arma::vec mu(n, fill::zeros);
  arma::vec a(n, fill::zeros);
  arma::vec z(n, fill::zeros);
  arma::mat xt_weighted = x;
  arma::mat weighted_x;

  int iter = 1;
  const double ybar = mean(y);
  if (family == "gaussian") {
    w(0) = ybar;
  }
  if (family == "inv.gaussian") {
    w(0) = 1 / std::pow(ybar, 2.0);
  }
  if (family == "gamma") {
    w(0) = 1 / ybar;
  }
  if (family == "logit") {
    w(0) = std::log(ybar / (1 - ybar));
  }
  if (family == "poisson") {
    w(0) = std::log(ybar);
  }

  while (iter < maxit) {
    old_w = w;

    eta = x * w;
    if (family == "gaussian") {
      mu = eta;
      a.ones();
      z = y;
    }
    if (family == "poisson") {
      mu = exp(eta);
      a = mu;
      z = a % eta + y - mu;
    }
    if (family == "logit") {
      mu = 1 / (1 + exp(-eta));
      a = mu % (1 - mu);
      z = a % eta + y - mu;
    }
    if (family == "gamma") {
      mu = 1 / eta;
      a = -mu % mu;
      z = a % eta + y - mu;
    }
    if (family == "inv.gaussian") {
      mu = 1 / sqrt(eta);
      a = -pow(mu, 3) / 2;
      z = a % eta + y - mu;
    }

    weighted_x = x.each_col() % a;
    if (n > m) {
      arma::mat system = xt_weighted.t() * weighted_x;
      system.diag() += lam;
      w = solve(system, xt_weighted.t() * z);
    } else {
      arma::mat system = weighted_x * xt_weighted.t();
      system.diag() += lam;
      w = xt_weighted.t() * solve(system, z);
    }

    xt_weighted = x.each_row() % trans(square(w));
    iter += 1;

    if (iter >= maxit) {
      warning("Did not converge. Increase maxit.");
    }
    if (norm(w - old_w, 2) < eps) {
      break;
    }
  }

  w.elem(find(abs(w) < 1e-3)).zeros();
  return List::create(Named("beta") = w, Named("iter") = iter);
}
