/**
 * @brief Utilities for TBPM module.
 *
 * @author: Zewen Wu, Miao Zhang, Yunhai Li
 */

#ifndef TBPLAS_TBPM_UTILS_H
#define TBPLAS_TBPM_UTILS_H

#include <iostream>
#include <random>
#include <chrono>
#include <map>
#include <string>
#include <functional> 
#include <iomanip>   

#include <eigen3/Eigen/Dense>
#include <eigen3/unsupported/Eigen/Splines> 

#include "dense.h"
#include "base/consts.h"

namespace tbplas::tbpm {

/**
 * @brief Class for generating random numbers.
 *
 */
class RandomGenerator {
public:
    /**
     * @brief Construct a new Random Generator object
     *
     * @param[in] seed seed for generating random numbers
     * @param[in] r_min lower bound of random numbers
     * @param[in] r_max upper bound of random numbers
     */
    explicit RandomGenerator(int seed = 1337, double r_min = 0.0, double r_max = 1.0);

    /**
     * @brief Generate a random number without resetting seed
     *
     * @return random number
     */
    double random_number();

    /**
     * @brief Generate a random number after resetting seed
     *
     * @param[in] seed seed with which to reset the
     * @return random number
     */
    double random_number(int seed);

    /**
     * @brief Populate the Eigen vector with random complex numbers and normalize it after
     * without resetting the seed
     *
     * @param[inout] vec vector in which to put the random numbers
     */
    void random_state(Eigen::VectorXcd& vec);

    /**
     * @brief Populate the Eigen vector with random complex numbers and normalize it after
     * resetting the seed
     *
     * @param[inout] vec vector in which to put the random numbers
     * @param[in] seed_i seed value assigned to this process
     *
     * @note When MPI is enabled and there are more than one process, each process will have
     * a different seed.
     */
    void random_state(Eigen::VectorXcd& vec, int seed_i);

    /**
     * @brief Reset the random engine's seed
     *
     * @param[in] seed seed for the random engine
     */
    void reset_seed(int seed);

private:
    std::mt19937 engine_;
    std::uniform_real_distribution<double> uni_real_dist_;
};

/**
 * @brief Print the head and tail elements of a vector for debugging.
 *
 * @tparam vec_t datatype of vector
 * @param[in] msg message to print before the elements
 * @param[in] x vector to print
 * @param[in] num_element number of elements
 */
template <typename vec_t>
void dump(
    const std::string& msg,
    const Eigen::MatrixBase<vec_t>& x,
    int num_element = 5)
{
    std::streamsize prev_precision = std::cout.precision();
    std::cout.setf(std::ios::scientific);
    std::cout.precision(15);
    int i_min = 0;
    int i_max = num_element;
    std::cout << msg << "\n";
    for (int i = i_min; i < i_max; ++i) {
        std::cout << i << " " << x[i] << "\n";
    }
    std::cout << "... ...\n";
    i_min = x.size() - num_element;
    i_max = x.size();
    for (int i = i_min; i < i_max; ++i) {
        std::cout << i << " " << x[i] << "\n";
    }
    std::cout << "\n";
    std::cout.unsetf(std::ios::scientific);
    std::cout.precision(prev_precision);
}

// Macro for performing tasks on master process in parallel mode
#define PARA_EXEC(cmd) if(meta.is_master){cmd;}


// /**
//  * @brief Benchmarks the performance of math operations.
//  *
//  * This function measures the execution time of fundamental operations like `amxsy`
//  *  `axpy`, `xmy`, which are critical performance bottlenecks in many calculations.
//  *  It allows for benchmarking operations relevant to different calculation modes 
//  * (e.g., quasi-eigenstates, DOS, DC conductivity).
//  *
//  * @param[in] h_sparse The sparse Hamiltonian matrix used for `amxsy` and `xmy` operations.
//  * @param[in] n_tests The number of times each operation will be repeated for benchmarking.
//  * @param[in] mode A string indicating the calculation mode, which determines which operations to benchmark.
//  * - "qe": Benchmarks operations relevant to quasi-eigenstate calculation.
//  * - "dos": Benchmarks operations relevant to Density of States calculation.
//  * - "dc": Benchmarks operations relevant to DC conductivity calculation.
//  * @return A `std::map` containing the total time (in seconds) spent on each benchmarked operation.
//  */
// template <typename sparse_t>
// std::map<std::string, double> benchmark_ops(const sparse_t& h_sparse,
//                                              int n_tests = 1000) {
//   using std::chrono::duration;
//   using std::chrono::duration_cast;
//   using std::chrono::high_resolution_clock;
//   using time_point_t = std::chrono::time_point<high_resolution_clock>;

//   std::map<std::string, double> total_time;
//   size_t vec_size = h_sparse.get_dim();

//   RandomGenerator rand_gen(1337);

//   // Initialize vectors once outside the benchmarking loop to avoid including
//   // allocation time in the measurement.
//   Eigen::VectorXcd t0(vec_size), t1(vec_size);
//   rand_gen.random_state(t0, 1337); // Populate t0 with random complex numbers
//   rand_gen.random_state(t1, 1338); // Populate t1 with random complex numbers

//   std::complex<double> coeff(0.5, 0.5); // A general complex coefficient for axpy
//   std::complex<double> corr(0.0, 0.0);  // Accumulator for xmy's return value
//   time_point_t start, end;              // Timestamps for measuring execution time

//   // Run the benchmark tests for the specified number of iterations
//   total_time["amxsy"] = 0.;
//   total_time["axpy"] = 0. ;
//   for (int i = 0; i < n_tests; ++i) {
//       // Benchmark amxsy: z = a * M * x - y
//       start = high_resolution_clock::now();
//       h_sparse.amxsy(2.0, t1, t0);
//       end = high_resolution_clock::now();
//       total_time["amxsy"] +=
//           duration_cast<duration<double>>(end - start).count();

//       // Benchmark axpy: y = x * x + y
//       start = high_resolution_clock::now();
//       axpy(coeff, t1, t0);
//       end = high_resolution_clock::now();
//       total_time["axpy"] +=
//           duration_cast<duration<double>>(end - start).count();
//   }

//   return total_time;
// }


/**
 * @brief Benchmarks the performance of amxsy and axpy operations.
 *
 * This function measures the execution time of fundamental operations like `amxsy`
 * and `axpy`, which are critical performance bottlenecks in many calculations.
 *
 * @param[in] h_sparse The sparse Hamiltonian matrix used for `amxsy` operation.
 * @param[in] n_tests The number of times each operation will be repeated for benchmarking.
 * @return A `std::map` containing the total time (in seconds) spent on each benchmarked operation.
 */
template <typename sparse_t>
std::map<std::string, double> benchmark_ops(const sparse_t& h_sparse,
                                             int n_tests = 1000) { // Default n_tests can be adjusted
    using std::chrono::duration;
    using std::chrono::duration_cast;
    using std::chrono::high_resolution_clock;
    using time_point_t = std::chrono::time_point<high_resolution_clock>;

    std::map<std::string, double> total_time;
    size_t vec_size = h_sparse.get_dim();

    RandomGenerator rand_gen(1337);

    Eigen::VectorXcd t0(vec_size), t1(vec_size), wf_out(vec_size);
    rand_gen.random_state(t0, 1337);
    vec_copy(t0,t1);
    vec_copy(t0,wf_out);

    // Initialize pointers
    Eigen::VectorXcd* p0 = &t0;
    Eigen::VectorXcd* p1 = &t1;
    Eigen::VectorXcd* p2 = nullptr;


    std::complex<double> coeff(0.5, 0.5); // A general complex coefficient for axpy
    time_point_t start, end;              // Timestamps for measuring execution time

    // --- WARM-UP PHASE ---
    // Run a few iterations to warm up caches and CPU before actual timing
    int n_warmup_tests = 10; // A reasonable number for warm-up
    for (int i = 0; i < n_warmup_tests; ++i) {
        p2 = p0;
        // h_sparse->amxsy(2.0, *p1, *p0);
        h_sparse.amxsy(2.0, *p1, *p0);
        axpy(coeff, *p2, wf_out);
        p0 = p1;
        p1 = p2;
    }
    // --- END WARM-UP ---

    total_time["amxsy"] = 0.;
    total_time["axpy"] = 0. ;
    // Run the benchmark tests for the specified number of iterations
    for (int i = 0; i < n_tests; ++i) {
        // Benchmark amxsy: z = a * M * x - y
        p2 = p0;
        start = high_resolution_clock::now();
        // h_sparse->amxsy(2.0, *p1, *p0);
        h_sparse.amxsy(2.0, *p1, *p0);
        end = high_resolution_clock::now();
        total_time["amxsy"] +=
            duration_cast<duration<double>>(end - start).count();

        // Benchmark axpy: y = alpha * x + y
        start = high_resolution_clock::now();
        axpy(coeff, *p2, wf_out);
        end = high_resolution_clock::now();
        total_time["axpy"] +=
            duration_cast<duration<double>>(end - start).count();
        p0 = p1;
        p1 = p2;
    }

    double norm = vec_norm(wf_out);
    std::cout << "norm: " << norm << std::endl;

    return total_time;
}

// Helper function to calculate median
double calculate_median(std::vector<double>& data) {
    if (data.empty()) return 0.0;
    std::sort(data.begin(), data.end());
    size_t n = data.size();
    if (n % 2 == 1) {
        return data[n / 2];
    } else {
        return (data[n / 2 - 1] + data[n / 2]) / 2.0;
    }
}

/**
 * @brief Runs the benchmark_ops function multiple times and performs statistical analysis on the results.
 *
 * This function directly calls `benchmark_ops` with the provided parameters and aggregates
 * the results for statistical analysis.
 *
 * @tparam sparse_t The type of the sparse Hamiltonian matrix.
 * @param[in] h_sparse The sparse Hamiltonian matrix used for benchmarking.
 * @param[in] n_internal_iterations The number of times each operation is repeated within one benchmark run.
 * @param[in] n_external_runs The number of times to repeat the entire benchmark run.
 */
template <typename sparse_t>
void run_benchmarks_and_analyze(const sparse_t& h_sparse,
                                int n_internal_iterations,
                                int n_external_runs) {

    std::vector<std::map<std::string, double>> all_results;

    // Run the benchmark multiple times and collect results
    for (int run_idx = 0; run_idx < n_external_runs; ++run_idx) {
        all_results.push_back(benchmark_ops(h_sparse, n_internal_iterations));
        std::cout << run_idx << " amxsy: " << all_results.back()["amxsy"]
        << " axpy: " << all_results.back()["axpy"] 
        << " ratio: " << all_results.back()["amxsy"]/all_results.back()["axpy"] 
        << std::endl;
    }

    // Aggregate collected times for statistics
    std::map<std::string, std::vector<double>> collected_times;
    for (const auto& result_map : all_results) {
        for (const auto& pair : result_map) {
            collected_times[pair.first].push_back(pair.second);
        }
    }

    std::cout << "--- Aggregate Benchmark Results Across " << n_external_runs << " Runs ---\n";
    std::cout << std::fixed << std::setprecision(6);

    // Print aggregate statistics for each operation
    for (auto& op_pair : collected_times) {
        const std::string& op_name = op_pair.first;
        std::vector<double>& times = op_pair.second;

        if (times.empty()) continue;

        double sum = std::accumulate(times.begin(), times.end(), 0.0);
        double mean = sum / times.size();
        double median = calculate_median(times);

        double sq_diff_sum = 0.0;
        for (double t : times) {
            sq_diff_sum += (t - mean) * (t - mean);
        }
        double std_dev = std::sqrt(sq_diff_sum / times.size());

        std::cout << op_name << ":\n";
        std::cout << "  Mean     = " << mean << " s\n";
        std::cout << "  Median   = " << median << " s\n";
        std::cout << "  Std Dev  = " << std_dev << " s ";
        std::cout << "(Relative: " << std::setprecision(2) << (mean != 0 ? (std_dev / mean) * 100 : 0.0) << "%)\n";
    }

    // Calculate mean and standard deviation of ratios
    std::vector<double> ratios;
    for(const auto& result_map : all_results) {
        if (result_map.count("amxsy") && result_map.count("axpy") && result_map.at("axpy") != 0.0) {
            ratios.push_back(result_map.at("amxsy") / result_map.at("axpy"));
        }
    }
    if (!ratios.empty()) {
        double sum_ratios = std::accumulate(ratios.begin(), ratios.end(), 0.0);
        double mean_ratio = sum_ratios / ratios.size();
        double median_ratio = calculate_median(ratios);

        double sq_diff_sum_ratios = 0.0;
        for (double r : ratios) {
            sq_diff_sum_ratios += (r - mean_ratio) * (r - mean_ratio);
        }
        double std_dev_ratios = std::sqrt(sq_diff_sum_ratios / ratios.size());

        std::cout << "\nRatio (amxsy/axpy):\n";
        std::cout << "  Mean Ratio = " << std::setprecision(3) << mean_ratio << "\n";
        std::cout << "  Median Ratio   = " << std::setprecision(3) << median_ratio << "\n";
        std::cout << "  Std Dev Ratio = " << std::setprecision(3) << std_dev_ratios << " ";
        std::cout << "(Relative: " << std::setprecision(2) << (mean_ratio != 0 ? (std_dev_ratios / mean_ratio) * 100 : 0.0) << "%)\n";
    } else {
        std::cout << "\nNo valid ratios to calculate aggregate statistics.\n";
    }
}


// Define the spline type to be used.
// We'll use a 1D spline (Dim=1) as N is a single value for each t.
// The degree will be a constant, e.g., cubic (3).
template<typename Scalar, int Degree>
using MySpline = Eigen::Spline<Scalar, 1, Degree>;

/**
 * @brief A class to encapsulate spline interpolation for Nt values based on t.
 * All N and t values are hardcoded internally and used to build the spline
 * within the constructor.
 */
class NtSplineInterpolator {
public:
    // Define the spline degree as a static constant
    // This can be a static member, as its a property of the spline itself.
    static const int SPLINE_DEGREE = 3; // Commonly used cubic spline degree

    /**
     * @brief Constructs the NtSplineInterpolator and initializes the spline.
     * All necessary data (t and N values) are hardcoded directly within this constructor.
     */
    NtSplineInterpolator();

    /**
     * @brief Gets the interpolated N value for a given t.
     * @param t_input The t value for which to get the interpolated N.
     * @return The interpolated N value.
     * @note If t_input is outside the spline's domain, a warning is printed
     * and the value at the nearest domain boundary is returned.
     */
    double getInterpolatedNt(double t_input) const;

private:
    // The spline object itself is the only private member needed to store.
    MySpline<double, SPLINE_DEGREE> spline_;
    double t_start_;
    double t_step_;
    double t_end_;
};




    // // --- Time Evolution Optimizer Class (using SplineInterpolation1D) ---
    // class TimeEvolutionOptimizer {
    // private:
    //     SplineInterpolation1D N_interp_; // Changed to SplineInterpolation1D
    //     SplineInterpolation1D R_interp_; // Changed to SplineInterpolation1D
    //     double A_;
    //     double B_;

    // public:
    //     // Constructor: Takes your pre-benchmarked data (t_all, Nt_t, Rt) and costs (A, B)
    //     TimeEvolutionOptimizer(
    //         const std::vector<double>& t_all,
    //         const std::vector<double>& Nt_t,
    //         const std::vector<double>& Rt,
    //         double A_cost, double B_cost)
    //         : N_interp_(t_all, Nt_t), R_interp_(t_all, Rt), A_(A_cost), B_(B_cost) {
    //         // Nothing else needed here, interpolators are constructed
    //     }

    //     // ... (rest of compute_total_cost and optimize_block_size functions remain the same) ...
    //     /**
    //      * @brief Computes the total cost of time evolution for a given block size.
    //      *
    //      * @param T_total Total simulation time.
    //      * @param b Block size (number of time steps per block propagation).
    //      * @param delta_t Time step size.
    //      * @return A pair: (total_amxsy_cost, total_axpy_cost)
    //      */
    //     std::pair<double, double> compute_total_cost(double T_total, int b, double delta_t) const {
    //         if (delta_t <= 0) {
    //             throw std::invalid_argument("delta_t must be positive.");
    //         }
    //         if (b <= 0) {
    //             throw std::invalid_argument("block size 'b' must be positive.");
    //         }

    //         double t_b = static_cast<double>(b) * delta_t;
            
    //         double N_blocks = T_total / t_b; 

    //         double sum_N_interp_t_vals = 0.0;
    //         for (int j = 0; j < b; ++j) {
    //             double current_t_val = static_cast<double>(j + 1) * delta_t;
    //             sum_N_interp_t_vals += N_interp_(current_t_val); // Uses the spline interpolator
    //         }
            
    //         double total_axpy = N_blocks * B_ * sum_N_interp_t_vals;
            
    //         double total_amxpy = T_total * A_ * R_interp_(t_b); // Uses the spline interpolator
            
    //         return {total_amxpy, total_axpy};
    //     }

    //     /**
    //      * @brief Finds the optimal block size 'b' that minimizes the total computational cost.
    //      *
    //      * @param T_total Total simulation time.
    //      * @param delta_t Time step size.
    //      * @param b_min Minimum block size to test.
    //      * @param b_max Maximum block size to test.
    //      * @return A pair: {optimal_b, minimum_total_cost_ratio_to_b_min}
    //      */
    //     std::pair<int, double> optimize_block_size(double T_total, double delta_t, int b_min, int b_max) const {
    //         if (b_min <= 0 || b_max < b_min) {
    //             throw std::invalid_argument("b_min must be positive and b_max >= b_min.");
    //         }

    //         double min_total_cost = std::numeric_limits<double>::max();
    //         int optimal_b = -1;
    //         double cost_at_b_min = 0.0;

    //         for (int b = b_min; b <= b_max; ++b) {
    //             std::pair<double, double> costs = compute_total_cost(T_total, b, delta_t);
    //             double current_total_cost = costs.first + costs.second;

    //             if (b == b_min) {
    //                 cost_at_b_min = current_total_cost;
    //             }

    //             if (current_total_cost < min_total_cost) {
    //                 min_total_cost = current_total_cost;
    //                 optimal_b = b;
    //             }
    //         }
            
    //         double ratio_to_b_min = 1.0;
    //         if (cost_at_b_min > std::numeric_limits<double>::epsilon()) {
    //             ratio_to_b_min = min_total_cost / cost_at_b_min;
    //         }

    //         return {optimal_b, ratio_to_b_min};
    //     }
    // };


} // namespace tbplas::tbpm
#endif