//
// Created by Jiandong Wang on 2026/6/18.
#include <RcppArmadillo.h>
#ifdef _OPENMP
#include <omp.h>
#endif

// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::plugins(openmp)]]

using namespace Rcpp;
using namespace arma;

// Log-sum-exp trick for numerical stability
inline double log_sum_exp(const rowvec& v) {
  double max_v = v.max();
  if (!std::isfinite(max_v)) return max_v;
  return max_v + std::log(sum(exp(v - max_v)));
}

// // Fit VII model for a specific K
// List fit_VII_single(const mat& X, int K, int max_iter = 1000, double tol = 1e-6,
//                     int threads = 1, int min_pts = 5, bool verbose = false) {
//    int n = X.n_rows;
//    int p = X.n_cols;
//    
//    // Initialization: k-means via Armadillo
//    mat means;
//    
//    // Explicitly force the boolean to avoid Armadillo template quirks
//    bool kmeans_success;
//    int current_threads = 1;
// #ifdef _OPENMP
//    current_threads = omp_get_max_threads();
//    omp_set_num_threads(1); 
// #endif
//    
//    if (verbose) {
//       kmeans_success = kmeans(means, X.t(), K, random_subset, 10, true);
//    } else {
//       kmeans_success = kmeans(means, X.t(), K, random_subset, 10, false);
//    }
//    
// #ifdef _OPENMP
//    omp_set_num_threads(current_threads);
// #endif
//    
//    if (!kmeans_success) {
//       // If k-means completely fails, disqualify this K
//       return List::create(Named("K") = K, Named("BIC") = -datum::inf);
//    }
//    means = means.t(); // Now K x p
//    
//    rowvec pro = ones<rowvec>(K) / K;
//    vec lambda = ones<vec>(K); // Spherical variances
//    
//    mat z = zeros<mat>(n, K);
//    double loglik_old = -datum::inf;
//    double loglik = 0;
//    
//    for (int iter = 0; iter < max_iter; ++iter) {
//       // --- E-step ---
//       loglik = 0;
//       mat log_rho = zeros<mat>(n, K);
//       
// #pragma omp parallel for num_threads(threads) if(threads > 1)
//       for (int i = 0; i < n; ++i) {
//          rowvec xi = X.row(i);
//          for (int k = 0; k < K; ++k) {
//             double dist2 = sum(square(xi - means.row(k)));
//             log_rho(i, k) = std::log(pro(k)) - (p / 2.0) * std::log(2 * datum::pi) -
//                (p / 2.0) * std::log(lambda(k)) - dist2 / (2 * lambda(k));
//          }
//       }
//       
//       // Calculate responsibilities and log-likelihood
//       for (int i = 0; i < n; ++i) {
//          double lse = log_sum_exp(log_rho.row(i));
//          z.row(i) = exp(log_rho.row(i) - lse);
//          loglik += lse;
//       }
//       
//       // Check convergence
//       if (std::abs(loglik - loglik_old) < tol) break;
//       loglik_old = loglik;
//       
//       // --- M-step ---
//       rowvec nk = sum(z, 0);
//       
//       // --- MINIMUM POINTS CHECK ---
//       for (int k = 0; k < K; ++k) {
//          if (nk(k) < min_pts) {
//             if (verbose) {
//                Rcpp::Rcout << "Cluster " << k + 1 << " collapsed to " << nk(k) 
//                            << " effective points. Disqualifying K=" << K << "\n";
//             }
//             // Immediately return a dead model so BIC rejects it
//             return List::create(
//                Named("K") = K,
//                Named("loglik") = -datum::inf,
//                Named("BIC") = -datum::inf
//             );
//          }
//       }
//       
//       pro = nk / n;
//       
//       for (int k = 0; k < K; ++k) {
//          if (nk(k) > 1e-10) { // Safety check, though the min_pts check above renders this mostly redundant
//             // Update means
//             rowvec weighted_sum = zeros<rowvec>(p);
//             for (int i = 0; i < n; ++i) {
//                weighted_sum += z(i, k) * X.row(i);
//             }
//             means.row(k) = weighted_sum / nk(k);
//             
//             // Update spherical variance (lambda)
//             double var_sum = 0;
// #pragma omp parallel for reduction(+:var_sum) num_threads(threads) if(threads > 1)
//             for (int i = 0; i < n; ++i) {
//                var_sum += z(i, k) * sum(square(X.row(i) - means.row(k)));
//             }
//             lambda(k) = var_sum / (nk(k) * p);
//          }
//       }
//    }
//    
//    // Calculate BIC (mclust convention: 2 * loglik - nu * log(n))
//    int nu = (K - 1) + (K * p) + K;
//    double bic = 2 * loglik - nu * std::log(n);
//    
//    return List::create(
//       Named("K") = K,
//       Named("loglik") = loglik,
//       Named("BIC") = bic,
//       Named("pro") = pro,
//       Named("mean") = means,
//       Named("variance") = lambda,
//       Named("z") = z
//    );
// }

List fit_VII_single(const mat& X, int K, int max_iter = 1000, double tol = 1e-6,
                    int threads = 1, int min_pts = 5, bool verbose = false) {
   int n = X.n_rows;
   int p = X.n_cols;
   
   // Initialization: k-means via Armadillo
   mat means;
   
   //bool kmeans_success = kmeans(means, X.t(), K, random_subset, 10, verbose);
   
   // 1. Save the current thread count and force OpenMP to 1 thread
#ifdef _OPENMP
   int current_threads = 1;
   current_threads = omp_get_max_threads(); // Remember the user's thread limit
   omp_set_num_threads(threads);                  // Clamp down to exactly 1 thread
#endif
   
   // 2. Execute kmeans (it will now strictly use 1 thread)
   bool kmeans_success;
   kmeans_success = kmeans(means, X.t(), K, random_subset, 10, verbose);
   
   // 3. Restore the original thread count so your E-step can use all cores
#ifdef _OPENMP
   omp_set_num_threads(current_threads);
#endif
   
   if (!kmeans_success) {
      stop("K-means initialization failed.");
   }
   means = means.t(); // Now K x p
   
   rowvec pro = ones<rowvec>(K) / K;
   vec lambda = ones<vec>(K); // Spherical variances
   
   mat z = zeros<mat>(n, K);
   double loglik_old = -datum::inf;
   double loglik = 0;
   
   for (int iter = 0; iter < max_iter; ++iter) {
      // --- E-step ---
      loglik = 0;
      mat log_rho = zeros<mat>(n, K);
      
#pragma omp parallel for num_threads(threads) if(threads > 1)
      for (int i = 0; i < n; ++i) {
         rowvec xi = X.row(i);
         for (int k = 0; k < K; ++k) {
            double dist2 = sum(square(xi - means.row(k)));
            log_rho(i, k) = std::log(pro(k)) - (p / 2.0) * std::log(2 * datum::pi) -
               (p / 2.0) * std::log(lambda(k)) - dist2 / (2 * lambda(k));
         }
      }
      
      // Calculate responsibilities and log-likelihood
      for (int i = 0; i < n; ++i) {
         double lse = log_sum_exp(log_rho.row(i));
         z.row(i) = exp(log_rho.row(i) - lse);
         loglik += lse;
      }
      
      // Check convergence
      if (std::abs(loglik - loglik_old) < tol) break;
      loglik_old = loglik;
      
      // --- M-step ---
      rowvec nk = sum(z, 0);
      pro = nk / n;
      
      for (int k = 0; k < K; ++k) {
         if (nk(k) > 1e-10) {
            // Update means
            rowvec weighted_sum = zeros<rowvec>(p);
            for (int i = 0; i < n; ++i) {
               weighted_sum += z(i, k) * X.row(i);
            }
            means.row(k) = weighted_sum / nk(k);
            
            // Update spherical variance (lambda)
            double var_sum = 0;
#pragma omp parallel for reduction(+:var_sum) num_threads(threads) if(threads > 1)
            for (int i = 0; i < n; ++i) {
               var_sum += z(i, k) * sum(square(X.row(i) - means.row(k)));
            }
            lambda(k) = var_sum / (nk(k) * p);
         }
      }
   }
   
   // Calculate BIC (mclust convention: 2 * loglik - nu * log(n))
   // Parameters: (K-1) for pro + K*p for means + K for lambda
   int nu = (K - 1) + (K * p) + K;
   double bic = 2 * loglik - nu * std::log(n);
   
   return List::create(
      Named("K") = K,
      Named("loglik") = loglik,
      Named("BIC") = bic,
      Named("pro") = pro,
      Named("mean") = means,
      Named("variance") = lambda,
      Named("z") = z
   );
}


// [[Rcpp::export]]
List fast_VII_mclust(const arma::mat& X, IntegerVector K_range, int threads = 1,
                     int min_pts = 5, bool verbose = false) {
  if(verbose){
     #ifdef _OPENMP
     Rcpp::Rcout << "OpenMP is ENABLED. Max threads available: " << omp_get_max_threads() << "\n";
     #else
     Rcpp::Rcout << "WARNING: OpenMP is DISABLED. Compiling serially.\n";
     #endif
     }

  List best_model;
  double best_bic = -datum::inf;
  NumericVector all_bics(K_range.size());

  for (int j = 0; j < K_range.size(); ++j) {
    int K = K_range[j];

    // Handle K=1 as a special edge case to avoid k-means crashing
    if (K == 1) {
       // ... K=1 logic (simple sample mean and global variance) ...
       // (Omitted for brevity, but essentially 1 iteration M-step)
       continue;
    }
    List current_model = fit_VII_single(X, K, 1000, 1e-6, threads, min_pts, verbose);
    double current_bic = current_model["BIC"];
    all_bics[j] = current_bic;

    if (current_bic > best_bic) {
      best_bic = current_bic;
      best_model = current_model;
    }
  }

  best_model["BIC_all"] = all_bics;
  best_model["K_range"] = K_range;
  return best_model;
}
