#include <eigen3/Eigen/Dense>

#include "base/consts.h"
#include "base/datatypes.h"
#include "base/kpoints.h"
#include "base/lattice.h"
#include "base/parallel.h"
#include "base/utils.h"
#include "builder/advanced.h"
#include "builder/primitive.h"
#include "diag/base.h"
#include "tbpm/solver.h"

using tbplas::base::complex_t;
using tbplas::base::gen_kpath;
using tbplas::base::gen_lattice_vectors;
using tbplas::base::NM;
using tbplas::base::Timer;
using tbplas::builder::extend_prim_cell;
using tbplas::builder::PrimitiveCell;
using tbplas::diag::DiagSolver;
using tbplas::tbpm::TBPMSolver;
using model_t = PrimitiveCell<complex_t>;

/**
 * @brief Make graphene primitive cell
 *
 * @return model_t
 */
model_t make_graphene_prim_cell()
{
    // Model parameters
    constexpr double t = -2.7;
    constexpr double lat = 0.246;
    constexpr double f = 1.0 / 3;
    constexpr double onsite = 0.0;

    // Generate lattice vectors
    Eigen::Matrix3d lat_vec = gen_lattice_vectors(lat, lat, 1.0, 90.0, 90.0, 60.0);
    Eigen::Vector3d origin(0.0, 0.0, 0.0);

    // Create the primitive cell
    model_t prim_cell(lat_vec, origin, NM);

    // Add orbitals, last 0 for labeling C_pz orbital
    prim_cell.add_orbital(0.0, 0.0, 0.0, onsite, 0);
    prim_cell.add_orbital(f, f, 0.0, onsite, 0);

    // Add hopping terms
    prim_cell.add_hopping(0, 0, 0, 0, 1, t);
    prim_cell.add_hopping(1, 0, 0, 1, 0, t);
    prim_cell.add_hopping(0, 1, 0, 1, 0, t);
    return prim_cell;
}

/**
 * @brief Make graphene sample with specific dimension.
 *
 * @param[in] dim dimension along a, b and c directions
 * @return model_t
 */
model_t make_graphene_sample(const std::tuple<int, int, int>& dim)
{
    // Create primitive cell
    model_t prim_cell = make_graphene_prim_cell();

    // For C++, the PrimitiveCell class is fast enough. And this is no need
    // to utilize the Sample class in most cases.
    model_t sample = extend_prim_cell(prim_cell, dim);
    return sample;
}

/**
 * @brief Calculate band structure using diagonalization.
 *
 */
void test_diag_bands()
{
    // Build the model
    model_t model = make_graphene_prim_cell();

    // Create a solver for the model
    DiagSolver<model_t> solver(model);

    // Set up parameters of the solver
    Eigen::MatrixX3d k_points {
        { 0.0, 0.0, 0.0 },
        { 2. / 3, 1. / 3, 0.0 },
        { 0.5, 0.0, 0.0 }, { 0.0, 0.0, 0.0 }
    };
    Eigen::Matrix3Xd k_path;
    Eigen::VectorXi k_idx;
    std::tie(k_path, k_idx) = gen_kpath(k_points.transpose(), { 100, 100, 100 });
    solver.config.prefix = "graphene";
    solver.config.k_points = k_path;

    // Call 'calc_bands' method of solver to evaluate band structure
    // Data files will be saved automatically.
    Timer timer;
    timer.tic("bands");
    auto data = solver.calc_bands();
    timer.toc("bands");

    // Report time usage
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

/**
 * @brief Calculate DOS using TBPM.
 *
 */
void test_tbpm_dos()
{
    // Build the model
    model_t model = make_graphene_sample({ 512, 512, 1 });

    // Create a solver for the model
    TBPMSolver<model_t> solver(model);

    // Set up parameters of the solver
    solver.config.prefix = "graphene";
    solver.config.num_random_samples = 1;
    solver.config.rescale = 9.0;
    solver.config.num_time_steps = 1024;

    // Call 'calc_corr_dos' method of solver to evaluate DOS correlation function
    // Data files will be saved automatically.
    Timer timer;
    timer.tic("dos");
    solver.calc_corr_dos();
    timer.toc("dos");

    // Report time usage
    if (solver.is_master()) {
        timer.report_total_time();
    }
}

int main(int argc, char* argv[])
{
    // Initialize MPI and openmp environments.
    tbplas::base::MPIENV_INIT(argc, argv);
    Eigen::initParallel();

    test_diag_bands();
    test_tbpm_dos();

    return 0;
}