#include <gtest/gtest.h>
#include <iostream>
#include <list>

#include "builder/advanced.h"
#include "materials/graphene.h"

using namespace tbplas::base;
using namespace tbplas::builder;
using namespace tbplas::materials;

namespace {
double calc_hop(const Term& term)
{
    constexpr double a0 = 0.1418;
    constexpr double a1 = 0.3349;
    constexpr double r_c = 0.6140;
    constexpr double l_c = 0.0265;
    constexpr double gamma0 = 2.7;
    constexpr double gamma1 = 0.48;
    constexpr double decay = 22.18;
    constexpr double q_pi = decay * a0;
    constexpr double q_sigma = decay * a1;
    double dr = term.distance;
    double n = term.rij(2) / dr;
    double v_pp_pi = -gamma0 * exp(q_pi * (1 - dr / a0));
    double v_pp_sigma = gamma1 * exp(q_sigma * (1 - dr / a1));
    double fc = 1 / (1 + exp((dr - r_c) / l_c));
    double hop = (n * n * v_pp_sigma + (1 - n * n) * v_pp_pi) * fc;
    return hop;
}

void extend_hop(PrimitiveCell<complex_t>& model, const double& max_distance = 0.75)
{
    auto neighbors = find_neighbors(model, model, 1, 1, 0, max_distance);
    for (const auto& n : neighbors) {
        cell_index_t ra, rb, rc;
        orbital_index_t orb_i, orb_j;
        std::tie(ra, rb, rc) = n.rn;
        std::tie(orb_i, orb_j) = n.pair;
        model.add_hopping(ra, rb, rc, orb_i, orb_j, calc_hop(n));
    }
}

TEST(IntraHopping, All)
{
    std::set<orbital_index_t> mySet = { 0, 1 };
    IntraHopping<complex_t> intra_hop;
    InterHopping<complex_t> inter_hop;
    intra_hop.add_hopping(0, 0, 0, 1, 0, 1.0);
    intra_hop.add_hopping(0, 0, 0, 0, 1, 1.0);
    intra_hop.add_hopping(-1, 0, 0, 1, 1, { -2.0, 1.0 });
    intra_hop.add_hopping(0, -1, 0, 4, 1, { 3.0, -1.0 });
    intra_hop.add_hopping(1, -1, 0, 1, 3, -1.0);
    intra_hop.add_hopping(1, 1, 0, 2, 2, 5.0);
    intra_hop.add_hopping(1, 1, 0, 2, 3, 4.0);
    std::cout << "intra" << std::endl;
    intra_hop.print();
    intra_hop.remove_orbitals(mySet);
    std::cout << "after remove" << std::endl;
    intra_hop.print();
    //    intra_hop.purge();
    //    inter_hop.add_hopping(0, 0, 0, 1, 0, 1.0);
    //    inter_hop.add_hopping(0, 0, 0, 0, 1, 1.0);
    //    inter_hop.add_hopping(-1, 0, 0, 1, 1, { -2.0, 1.0 });
    //    inter_hop.add_hopping(0, -1, 0, 4, 1, { 3.0, -1.0 });
    //    inter_hop.add_hopping(1, -1, 0, 1, 3, -1.0);
    //    inter_hop.add_hopping(1, 1, 0, 2, 2, 5.0);
    //    inter_hop.add_hopping(1, 1, 0, 2, 3, 4.0);
    //    std::cout << "inter" << std::endl;
    //    inter_hop.print();
    //    bool f=intra_hop.remove_rn(1, 1, 0);
    //    bool f = intra_hop.remove_hopping(1, 1, 0,2,3);
    //    inter_hop.remove_orbitals(mySet);
    //    int i_ = inter_hop.get_num_hopping();
    //    std::cout << f_ << std::endl;
    //    std::cout << "i_=" << i_ << std::endl;
    //    inter_hop.print();
    //    inter_hop.purge();
    //    auto a = inter_hop.get_data();
}

TEST(PrimitiveCell, all)
{
    PrimitiveCell<double> prim_cell(10);
    std::list<Orbital> orbital_list;
    orbital_list.emplace_back(0.3, 0.3, 0.0, 0.0, 1);
    orbital_list.emplace_back(0.6, 0.6, 0.0, 0.0, -1);
    prim_cell.add_orbitals(orbital_list);
    std::cout << "before:\n";
    prim_cell.print();
    for (auto& orb : orbital_list) {
        std::cout << orb.label << "\n";
    }
    for (auto& orb : orbital_list) {
        orb.label = 2;
    }
    std::cout << "after:\n";
    prim_cell.print();
    for (auto& orb : orbital_list) {
        std::cout << orb.label << "\n";
    }
}

TEST(PrimitiveCell, save)
{
    auto prim_cell = make_graphene_diamond();
    prim_cell.save("graphene");
    PrimitiveCell<complex_t> extend_cell = extend_prim_cell(prim_cell, { 3, 3, 1 });
    extend_cell.remove_orbitals({ 9, 10 });
    extend_cell.trim();
    extend_cell.save("extend_cell");
    Eigen::Matrix3d lat_frac { { 1.0, 0.0, 0.0 }, { -1.0, 2.0, 0.0 }, { 0.0, 0.0, 1.0 } };
    PrimitiveCell<complex_t> res_cell = reshape_prim_cell(prim_cell, lat_frac);
    res_cell = extend_prim_cell(res_cell, { 3, 3, 1 });
    res_cell.apply_pbc({ true, false, false });
    res_cell.save("gnr");
    res_cell.rotate(0.5 * PI);
    res_cell.save("gnr_rot");
}

TEST(PrimitiveCell, check_dim)
{
    auto prim_cell = make_graphene_diamond();
    ASSERT_DEATH({ prim_cell.check_dim(); }, "Assertion");
    PrimitiveCell<complex_t> cell2 = extend_prim_cell(prim_cell, { 2, 1, 1 });
    ASSERT_DEATH({ cell2.check_dim(); }, "Assertion");
    cell2 = extend_prim_cell(prim_cell, { 1, 2, 1 });
    ASSERT_DEATH({ cell2.check_dim(); }, "Assertion");
    cell2 = extend_prim_cell(prim_cell, { 2, 2, 1 });
    ASSERT_NO_FATAL_FAILURE({ cell2.check_dim(); });
}

TEST(PrimitiveCell, tbg)
{
    // Evaluate twisting angle
    int tbg_idx = 5;
    double cos_ang = (3 * tbg_idx * tbg_idx + 3 * tbg_idx + 0.5) / (3 * tbg_idx * tbg_idx + 3 * tbg_idx + 1);
    double angle = acos(cos_ang);

    // Build primitive cells of each layer
    auto fixed_cell = make_graphene_diamond();
    auto twisted_cell = make_graphene_diamond();

    // Twist the primitive cell of top layer
    double shift = 0.3349;
    twisted_cell.rotate(angle, Eigen::Vector3d::Zero(), shift);

    // Evaluate hetero-lattice
    double i = tbg_idx * 1.0;
    Eigen::Matrix3d hetero_lattice { { i, i + 1, 0 }, { -(i + 1), 2 * i + 1, 0 }, { 0, 0, 1 } };
    hetero_lattice = frac2cart(fixed_cell.get_lattice(), hetero_lattice.transpose()).transpose();

    // Reshape primitive cells to produce layers
    PrimitiveCell<complex_t> layer_fixed = make_hetero_layer(fixed_cell, hetero_lattice, 1e-5);
    PrimitiveCell<complex_t> layer_twisted = make_hetero_layer(twisted_cell, hetero_lattice, 1e-5);

    // Align the layers for better appearance though it is not essential.
    // angle = -atan(hetero_lattice(0, 1) / hetero_lattice(0, 0));
    // layer_fixed.rotate(angle);
    // layer_twisted.rotate(angle);

    // Two algorithms to build the model
    int algo = 0;
    PrimitiveCell<complex_t> merged_cell;
    std::vector<const PrimitiveCell<complex_t>*> prim_cells = { &layer_fixed, &layer_twisted };
    if (algo == 0) {
        std::vector<const PCInterHopping<complex_t>*> inter_hops = {};
        merged_cell = merge_prim_cell(prim_cells, inter_hops);
        extend_hop(merged_cell, 0.75);
    } else {
        PCInterHopping<complex_t> inter_hop(layer_fixed, layer_twisted);
        auto neighbors = find_neighbors(layer_fixed, layer_twisted, 1, 1, 0, 0.75);
        for (const auto& n : neighbors) {
            cell_index_t ra, rb, rc;
            orbital_index_t orb_i, orb_j;
            std::tie(ra, rb, rc) = n.rn;
            std::tie(orb_i, orb_j) = n.pair;
            inter_hop.add_hopping(ra, rb, rc, orb_i, orb_j, calc_hop(n));
        }
        extend_hop(layer_fixed, 0.75);
        extend_hop(layer_twisted, 0.75);
        std::vector<const PCInterHopping<complex_t>*> inter_hops = {&inter_hop};
        merged_cell = merge_prim_cell(prim_cells, inter_hops);
    }
    std::cout << "prim_cell: " << merged_cell.get_num_orbitals() << " " << merged_cell.get_num_hopping() << std::endl;

    // Save model
    merged_cell = extend_prim_cell(merged_cell, { 4, 4, 1 });
    merged_cell.apply_pbc({ true, true, false });
    std::cout << "sample: " << merged_cell.get_num_orbitals() << " " << merged_cell.get_num_hopping() << std::endl;
    merged_cell.save("merged_cell");
}

TEST(SK, all)
{
    SK sk;
    std::vector<std::string> labels = { "s", "px", "py", "pz", "dxy", "dyz", "dzx", "dx2-y2", "dz2" };
    Eigen::VectorXcd coeff = Eigen::VectorXcd::Random(10);
    Eigen::Vector3d r = Eigen::Vector3d::Random();
    r = Eigen::Vector3d::Random();
    r = Eigen::Vector3d::Random();
    std::cout << coeff.transpose().real() << "\n";
    std::cout << coeff.transpose().imag() << "\n";
    std::cout << r.transpose() << "\n";
    complex_t v_sss = coeff(0);
    complex_t v_sps = coeff(1);
    complex_t v_sds = coeff(2);
    complex_t v_pps = coeff(3);
    complex_t v_ppp = coeff(4);
    complex_t v_pds = coeff(5);
    complex_t v_pdp = coeff(6);
    complex_t v_dds = coeff(7);
    complex_t v_ddp = coeff(8);
    complex_t v_ddd = coeff(9);
    for (const auto& s1 : labels) {
        for (const auto& s2 : labels) {
            complex_t t = sk.eval(r, s1, s2, v_sss, v_sps, v_sds, v_pps, v_ppp, v_pds, v_pdp, v_dds, v_ddp, v_ddd);
            std::cout << s1 << " " << s2 << " " << t.real() << " " << t.imag() << "\n";
        }
    }
}
}
