// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>

using namespace arma;
using namespace Rcpp;


// [[Rcpp::export]]
double denssr(double volume, int havlkm, int n, String method = "loglik", double mix = 0){
  double vastaus;
  if(havlkm == 0){
    vastaus = 0;  
  }else if(method == "loglik"){
    vastaus = havlkm * log(havlkm / (n*volume));  
  }else{
    if(mix == 0){
      vastaus = pow(havlkm, 2) / (n*volume);
    }else{
      vastaus = mix*(2-mix)*pow(havlkm, 2) / (n*volume);      
    }   
  }
  return(vastaus);
}

// [[Rcpp::export]]
vec denssrMat(vec volume, vec havlkm, int n, String method = "loglik", double mix = 0){
  vec vastaus = zeros<vec>(volume.n_elem);
  if(method == "loglik"){
    vastaus = havlkm % log(havlkm / (n*volume));  
  }else{
    if(mix == 0){
      vastaus = pow(havlkm, 2) / (n*volume);
    }else{
      vastaus = mix*(2-mix) * pow(havlkm, 2) / (n*volume);      
    }   
  }
  uvec bol = find_nonfinite(vastaus);
  vastaus(bol) = -datum::inf*ones<vec>(bol.n_elem);
  return vastaus;
}

// [[Rcpp::export]]
double massone(vec rec) {
  int d = rec.n_elem/2;
  uvec i = 2*linspace<uvec>(0, d-1, d);
  double vol = prod(rec(i+1) - rec(i));
  return(vol);
}

// [[Rcpp::export]]
vec massoneMat(mat rec){
  int d = rec.n_cols/2;
  uvec i = 2*linspace<uvec>(0, d-1, d);
  vec vol = prod(rec.cols(i+1) - rec.cols(i), 1);
  return vol;
}


// [[Rcpp::export]]
List findsplitter2(mat grid, mat x, vec rec, int n, String method = "loglik", int minobs = 0, vec recdown = 0, vec rechigh = 0){
  int l = x.n_rows;
  int d = x.n_cols;
  int jpislkm;
  int jakopisteint;
  int leftobslkm;
  int rightobslkm;
  int minvali;
  double jakopiste;
  double volumeleft;
  double volumeright;
  vec leftrec;
  vec rightrec;
  vec ssrvec1 = ones<vec>(d);
  vec ssrvec2;
  vec valvec = ssrvec1;
  vec valvecint = ssrvec1;
  
  int i = 0;
  while(i < d){
    jpislkm = rechigh(i) - recdown(i)-1;
    if(jpislkm >= 1){
      ssrvec2 = ones<vec>(jpislkm);  
      for(int j = 0; j < jpislkm; j++){
        jakopisteint = recdown(i)+j;
        jakopiste = grid(jakopisteint, i);  
        
        leftrec = rec;
        leftrec(2*i+1) = jakopiste;
        rightrec = rec;
        rightrec(2*i) = jakopiste;
                
        leftobslkm = sum(x.col(i) < jakopiste);
        rightobslkm = sum(x.col(i) > jakopiste);
        
        volumeleft = massone(leftrec);
        volumeright = massone(rightrec);
        if( (leftobslkm == 0) || (rightobslkm == 0)){
          ssrvec2(j) = -datum::inf;
        }else{
          ssrvec2(j) = denssr(volumeleft, leftobslkm, n, method) + denssr(volumeright, rightobslkm, n, method);
        }
      } 
      //ssrvec2 = ssrvec2(find(is_finite(ssrvec2)));
      minvali = which_min(as<NumericVector>(wrap(-ssrvec2)));
      if(!is_finite(minvali)){
        minvali = 0;
      }
      valvecint(i) = recdown(i) + minvali;
      valvec(i) = grid(valvecint(i), i);      
      ssrvec1(i) = ssrvec2(minvali);      
    }else{
      ssrvec1(i) = -datum::inf;  
    }
    i++;  
  }
  //return List::create(ssrvec2, ssrvec1, valvec, valvecint);
  int ind = which_min(as<NumericVector>(wrap(-ssrvec1)));
  if(!is_finite(ind)){
    ind = 0;
  }
  double val = valvec(ind);
  double valio = valvecint(ind);
  return List::create(Named("val") = val, Named("ind") = ind, Named("valio") = valio);
}  


// [[Rcpp::export]]
List findsplitter3(mat grid, mat x, vec rec, int n, String method = "loglik", int minobs = 0, vec recdown = 0, vec rechigh = 0){
  int l = x.n_rows;
  int d = x.n_cols;
  int minvali;
  int jpislkm;
  mat leftrec;
  mat rightrec;
  mat bol;
  vec ssrvec1 = ones<vec>(d);
  vec ssrvec2;
  vec valvec = ssrvec1;
  vec valvecint = ssrvec1;
  vec leftobslkm;
  vec rightobslkm;
  vec volumeleft;
  vec volumeright;
  vec jpisvec;
  
  uvec j;
  uvec jakopisteint;
  mat jakopiste;

  jpisvec = rechigh - recdown -1;
  vec nonZero = jpisvec(find(jpisvec != 0));
  for(int i = 0; i < nonZero.n_elem; i++){
    jpislkm = nonZero(i);
    j = linspace<uvec>(0, jpislkm-1, jpislkm);
    jakopisteint = recdown(i)+j;
    jakopiste= grid(jakopisteint, as<uvec>(wrap(i)));  
      
    leftrec = zeros<mat>(jpislkm, rec.n_elem);
    leftrec.each_row() = trans(rec);
    leftrec.col(2*i+1) = jakopiste;
      
    rightrec = zeros<mat>(jpislkm, rec.n_elem);
    rightrec.each_row() = trans(rec);
    rightrec.col(2*i) = jakopiste;
      
    leftobslkm = zeros<vec>(jakopiste.n_elem);
    rightobslkm = zeros<vec>(jakopiste.n_elem);
    for(int j = 0; j < jakopiste.n_elem; j++){      
      leftobslkm(j) = sum(x.col(i) < jakopiste(j));
      rightobslkm(j) = sum(x.col(i) > jakopiste(j));  
    }
    volumeleft = massoneMat(leftrec);
    volumeright = massoneMat(rightrec);
    ssrvec2 = denssrMat(volumeleft, leftobslkm, n, method) + denssrMat(volumeright, rightobslkm, n, method);
      
    minvali = which_max(as<NumericVector>(wrap(ssrvec2)));
    if(!is_finite(minvali)){
      minvali = 0;
    }
    valvecint(i) = recdown(i) + minvali;
    valvec(i) = grid(valvecint(i), i);      
    ssrvec1(i) = ssrvec2(minvali);           
  }
  ssrvec1(find(jpisvec == 0)) = -datum::inf*ones<vec>(jpisvec.n_elem - nonZero.n_elem);
    
  //return List::create(ssrvec2, ssrvec1, valvec, valvecint);
  int ind = which_max(as<NumericVector>(wrap(ssrvec1)));
  if(!is_finite(ind)){
    ind = 0;
  }
  double val = valvec(ind);
  double valio = valvecint(ind);
  return List::create(Named("val") = val, Named("ind") = ind, Named("valio") = valio);
}  


// [[Rcpp::export]]
List densplitter2(mat dendat, int minobs = 1, String method = "loglik", bool newSplit = true){
  int n = dendat.n_rows;
  int d = dendat.n_cols;
  int m;
  int maksi;
  int pinin = 0;
  int saatu = -1;
  int gridpoint;
  int direc;
  int lkmleft;
  int lkmright;
  List fs;
  double point;
  mat currecs;
  mat curobs;
  mat curdown;
  mat curhigh;
  mat finrecs;
  mat findown;
  mat finhigh;
  mat grid;
  mat inside;
  mat x;
  vec suppo;
  uvec indet;
  uvec wmin;
  uvec wmax;
  uvec i;
  uvec j;
  uvec leftobs;
  uvec rightobs;
  NumericVector notindet;
  uvec ordi;
  vec ala;
  vec yla;
  rowvec rec;
  rowvec leftrec;
  rowvec rightrec;
  rowvec recdown;
  rowvec leftdown;
  rowvec rightdown;
  rowvec rechigh;
  rowvec lefthigh;
  rowvec righthigh;
  vec obs;

  suppo = zeros<vec>(d*2, 1);
  indet = zeros<uvec>(d*2, 1);  
  i = 2*linspace<uvec>(0, d-1, d);
  suppo(i) = min(dendat);
  suppo(i+1) = max(dendat); 
  for(int i = 0; i < d; i++){
    wmin = find(dendat.col(i) == suppo(2*i));
    indet(2*i) = wmin(0);
    wmax = find(dendat.col(i) == suppo(2*i+1));
    indet(2*i+1) = wmax(0);
  }   
  
  i = linspace<uvec>(0, n-1, n);
  notindet = setdiff(as<NumericVector>(wrap(i)), as<NumericVector>(wrap(indet)));
  inside = dendat.rows(as<uvec>(wrap(notindet)));
  m = inside.n_rows;
  grid = zeros<mat>(m+1, d);
  
  for(int i = 0; i < d; i++){
    ordi = sort_index(inside.col(i));
    grid(0, i) = min(dendat.col(i));
    grid(m, i) = max(dendat.col(i));
    j = linspace<uvec>(0, m-2, m-1);    
    ala = inside.col(i);
    ala = ala(ordi(j));
    j = linspace<uvec>(1, m-1, m-1);    
    yla = inside.col(i);
    yla = yla(ordi(j));
    grid.submat(1, i, m-1, i) = (ala + yla)/2; 
  }
  
  maksi = 2*n;
  currecs = zeros<mat>(maksi, 2*d);
  finrecs = currecs;
  currecs.row(0) = trans(suppo);
  curobs = zeros<mat>(maksi, m);
  curobs.row(0) = ones<mat>(1, m);
  curdown =  zeros<mat>(maksi, d);
  curhigh = curdown;
  findown = curdown;
  finhigh = curdown;
  curdown.row(0) = ones<mat>(1, d);
  curhigh.row(0) = (m+1)*ones<mat>(1, d);
  while(pinin > -1){
    //Rcout << pinin << " " << saatu << std::endl;
    rec = currecs.row(pinin);
    obs = trans(curobs.row(pinin));
    recdown = curdown.row(pinin);
    rechigh = curhigh.row(pinin);  
    pinin --;
    x = inside.rows(find(obs == 1));    
    //if(p == 2) return(List::create(grid, x, rec, n, minobs, recdown, rechigh));    
    //fs = findsplitter2(grid, x, trans(rec), n, method, minobs, trans(recdown), trans(rechigh));
    if(newSplit){
      fs = findsplitter3(grid, x, trans(rec), n, method, minobs, trans(recdown), trans(rechigh));
    }else{
      fs = findsplitter2(grid, x, trans(rec), n, method, minobs, trans(recdown), trans(rechigh));      
    }
    //return fs;
    direc = fs["ind"];
    point = fs["val"];
    gridpoint = fs["valio"];
    gridpoint++;
    leftobs = (obs == 1) && (inside.col(direc) < point);
    rightobs = (obs == 1) && (inside.col(direc) > point);
    
    leftrec = rec;
    leftrec(2*direc+1) = point;
    rightrec = rec;
    rightrec(2*direc) = point;
    
    leftdown = recdown;
    lefthigh = rechigh;
    lefthigh(direc) = gridpoint;
    
    rightdown = recdown;
    righthigh = rechigh;
    rightdown(direc) = gridpoint;
        
    lkmleft = sum(leftobs);
    lkmright = sum(rightobs); 
    //if(p == 2) return List::create(direc, point, gridpoint, leftrec, rightrec, leftdown, lefthigh, rightdown, righthigh);
    //Rcout << "Doing ifs" << std::endl;
    if((lkmleft > minobs) && (lkmright > minobs)){
      //Rcout << pinin << " " << saatu <<  " eka" << std::endl;
      currecs.row(pinin+1) = leftrec;
      currecs.row(pinin+2) = rightrec;
      curobs.row(pinin+1) = as<rowvec>(wrap(leftobs));
      curobs.row(pinin+2) = as<rowvec>(wrap(rightobs));
      curdown.row(pinin+1) = leftdown;
      curhigh.row(pinin+1) = lefthigh;
      curdown.row(pinin+2) = rightdown;
      curhigh.row(pinin+2) = righthigh;
      pinin = pinin + 2;
    //return List::create(currecs, curobs, curdown, curhigh);
    }else if((lkmleft > minobs) && (lkmright <= minobs)){
      //Rcout << pinin << " " << saatu <<  " toka" << std::endl;
      pinin = pinin+1;
      saatu = saatu+1;
      currecs.row(pinin) = leftrec;
      curobs.row(pinin) = as<rowvec>(wrap(leftobs));
      curdown.row(pinin) = leftdown;
      curhigh.row(pinin) = as<rowvec>(wrap(lefthigh));
      finrecs.row(saatu) = rightrec;
      findown.row(saatu) = rightdown;
      finhigh.row(saatu) = righthigh;
      //return List::create(currecs, curobs, curdown, curhigh, finrecs, findown, finhigh);
    }else if((lkmleft <= minobs) && (lkmright > minobs)){
      //Rcout << pinin << " " << saatu <<  " kolmas" << std::endl;
      pinin = pinin+1;
      saatu = saatu+1;
      currecs.row(pinin) = rightrec;
      curobs.row(pinin) = as<rowvec>(wrap(rightobs));
      curdown.row(pinin) = rightdown;
      curhigh.row(pinin) = as<rowvec>(wrap(righthigh));
      finrecs.row(saatu) = leftrec;
      findown.row(saatu) = leftdown;
      finhigh.row(saatu) = lefthigh;
     // return List::create(currecs, curobs, curdown, curhigh, finrecs, findown, finhigh);
    }else{
      //Rcout << pinin << " " << saatu <<  " neljäs" << std::endl;
      finrecs.row(saatu+1) = leftrec;
      finrecs.row(saatu+2) = rightrec;
      findown.row(saatu+1) = leftdown;
      findown.row(saatu+2) = rightdown;
      finhigh.row(saatu+1) = lefthigh;
      finhigh.row(saatu+2) = righthigh;
      saatu = saatu+2;
      //return List::create(finrecs, findown, finhigh);
    }
  }
  //return List::create(finrecs, findown, finhigh, saatu);
  mat recs = finrecs.rows(0, saatu);
  mat down = findown.rows(0, saatu);
  mat high = finhigh.rows(0, saatu);
  
  return List::create(
    Named("recs") = recs, 
    Named("support") = suppo, 
    Named("grid") = grid, 
    Named("down") = down, 
    Named("high") = high);  
}


// [[Rcpp::export]]
int armaWhichMin(vec x){
  return which_min(as<NumericVector>(wrap(x)));
}

// [[Rcpp::export]]
int armaWhichMax(vec x){
  return which_max(as<NumericVector>(wrap(x)));
}