#include "model.h"

#include "base/consts.h"
#include "base/lattice.h"
#include "base/utils.h"
#include "builder/advanced.h"
#include "builder/base.h"
#include "builder/kdtree.h"
#include "builder/primitive.h"
#include "materials/graphene.h"

#include "mask.h"

namespace tbplas::tests {

using base::df_dk_cart2frac;
using base::frac2cart;
using base::gen_lattice_vectors;
using base::gen_reciprocal_vectors;
using base::NM;
using base::PI;
using base::Timer;
using base::TWOPI;
using base::UIMAG;
using builder::cell_index_t;
using builder::extend_prim_cell;
using builder::make_hetero_layer;
using builder::merge_prim_cell;
using builder::orbital_index_t;
using builder::PCInterHopping;
using builder::Term;
using materials::make_graphene_diamond;

model_t make_graphene_pc(
    double t,
    double a,
    bool with_onsite)
{
    double lat = a * sqrt(3);
    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);
    PrimitiveCell<complex_t> prim_cell(2, lat_vec, origin, NM);
    prim_cell.set_orbital_position(0, 0., 0., 0.0);
    prim_cell.set_orbital_position(1, 1. / 3, 1. / 3, 0.0);
    if (with_onsite) {
        prim_cell.set_orbital_energy(0, -0.1);
        prim_cell.set_orbital_energy(1, 0.1);
    }
    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;
}

model_t make_graphene(
    double t,
    double a,
    bool with_onsite,
    const std::tuple<int, int, int>& dim)
{
    model_t prim_cell = make_graphene_pc(t, a, with_onsite);
    prim_cell = extend_prim_cell(prim_cell, dim);
    return prim_cell;
}

model_t make_overlap(const model_t& model)
{
    // As a demo, overlap has only diagonal (on-site) terms which equals 1.0.
    model_t overlap(model.get_num_orbitals(), model.get_lattice(), model.get_origin(), 1.0);
    for (size_t i = 0; i < model.get_num_orbitals(); ++i) {
        const auto& orbital = model.get_orbital(i);
        overlap.set_orbital_position(i, orbital.fa, orbital.fb, orbital.fc);
        overlap.set_orbital_energy(i, 1.0);
    }
    // Add a fake hopping term to avoid empty hopping arrays,
    // to be consistent with Python version.
    overlap.add_hopping(0, 0, 0, 0, 1, 0.0);
    return overlap;
}

model_t make_haldane()
{
    double lat = 1.0;
    double delta = 0.2;
    double t = -1.0;
    complex_t t2 = 0.15 * std::exp(UIMAG * PI / 2.);
    complex_t t2c = std::conj(t2);

    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);
    PrimitiveCell<complex_t> prim_cell(2, lat_vec, origin, NM);
    prim_cell.set_orbital_position(0, 1. / 3, 1. / 3, 0.0);
    prim_cell.set_orbital_position(1, 2. / 3, 2. / 3, 0.0);
    prim_cell.set_orbital_energy(0, -delta);
    prim_cell.set_orbital_energy(1, delta);
    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);
    prim_cell.add_hopping(1, 0, 0, 0, 0, t2);
    prim_cell.add_hopping(1, -1, 0, 1, 1, t2);
    prim_cell.add_hopping(0, 1, 0, 1, 1, t2);
    prim_cell.add_hopping(1, 0, 0, 1, 1, t2c);
    prim_cell.add_hopping(1, -1, 0, 0, 0, t2c);
    prim_cell.add_hopping(0, 1, 0, 0, 0, t2c);
    return prim_cell;
}

//******************************************************************************//
//                           Complex model functions                            //
//******************************************************************************//
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(model_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));
    }
}

void cutoff_pc(
    model_t& model,
    const Eigen::Vector3d& center,
    const double& radius = 3.0)
{
    std::set<size_t> idx_remove;
    Eigen::Matrix3Xd orb_pos = model.get_orbital_positions_nm();
    Eigen::Vector3d dr;
    for (size_t i = 0; i < orb_pos.cols(); ++i) {
        dr = orb_pos.col(i) - center;
        if (dr.norm() > radius) {
            idx_remove.insert(i);
        }
    }
    model.remove_orbitals(idx_remove);
    // model.trim();
}

model_t top_down(
    const model_t& prim_cell,
    const int& start_width,
    const int& iteration,
    const int& extension)
{
    // Create the extended cell
    int final_width = start_width * static_cast<int>(std::pow(extension, iteration));
    model_t extended_cell = extend_prim_cell(prim_cell, { final_width, final_width, 1 });
    extended_cell.apply_pbc({ false, false, false });

    // Create the mask
    Box start_box(0, 0, final_width - 1, final_width - 1, false);
    Mask mask(start_box, extension, iteration);

    // Remove orbitals
    mask.etch_prim_cell(extended_cell, final_width);
    return extended_cell;
}

model_t make_tbg(int tbg_idx)
{
    Timer timer;
    timer.tic("make_tbg");

    // Evaluate twisting angle
    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;
    double r_cut = 0.55; //60:0.55 120:0.75 230:1.0
    model_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, r_cut);
    } else {
        PCInterHopping<complex_t> inter_hop(layer_fixed, layer_twisted);
        auto neighbors = find_neighbors(layer_fixed, layer_twisted, 1, 1, 0, r_cut);
        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, r_cut);
        extend_hop(layer_twisted, r_cut);
        std::vector<const PCInterHopping<complex_t>*> inter_hops = { &inter_hop };
        merged_cell = merge_prim_cell(prim_cells, inter_hops);
    }

    timer.toc("make_tbg");
    timer.report_total_time();
    std::cout << "prim_cell: " << merged_cell.get_num_orbitals() << " " << merged_cell.get_num_hopping() << std::endl;
    // merged_cell.save("tbg_prim_cell");
    return merged_cell;
}

model_t make_quasi_crystal(double radius)
{
    Timer timer;
    timer.tic("make_quasi_crystal");

    // Geometric parameters
    double shift = 0.3349;
    double a = 0.142;
    double angle = 30.0 / 180.0 * PI;
    Eigen::Vector3d center { { 2.0 / 3 }, { 2.0 / 3 }, { 0.0 } };

    // Estimate dim for diamond-shaped prim_cell
    // The radius of inscribed circle of a diamond is r = sqrt(3) / 4 * lat, with
    // lat being edge length. Since lat = a * sqrt(3), then r = 0.75 * a, with a
    // being C-C bond length.
    int rmin_dia = static_cast<int>(std::ceil(radius / (0.75 * a))) + 1;
    std::tuple<int, int, int> dim = { rmin_dia, rmin_dia, 1 };

    // Build fixed and twisted layers
    model_t prim_cell = make_graphene_diamond();
    model_t layer_fixed = extend_prim_cell(prim_cell, dim);
    model_t layer_twisted = extend_prim_cell(prim_cell, dim);

    // Get the Cartesian coordinate of rotation center
    center[0] += static_cast<int>(std::get<0>(dim) / 2);
    center[1] += static_cast<int>(std::get<1>(dim) / 2);
    center = (center.transpose() * prim_cell.get_lattice()).transpose();

    // Remove unnecessary orbitals
    cutoff_pc(layer_fixed, center, radius);
    cutoff_pc(layer_twisted, center, radius);

    // Rotate and shift twisted layer
    spiral_prim_cell(layer_twisted, angle, center, shift);

    // Reset the lattice of twisted layer
    layer_twisted.reset_lattice(layer_fixed.get_lattice(), layer_fixed.get_origin(),
        1.0, true);

    // Common procedure of making tbg
    int algo = 0;
    model_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);
    }

    timer.toc("make_quasi_crystal");
    timer.report_total_time();
    std::cout << "prim_cell: " << merged_cell.get_num_orbitals() << " " << merged_cell.get_num_hopping() << std::endl;
    // merged_cell.save("quasi_crystal_prim_cell");
    return merged_cell;
}

model_t make_fractal(int iteration)
{
    // Create a square lattice
    model_t prim_cell;
    prim_cell.add_orbital(0, 0, 0, 0.0, 0);
    prim_cell.add_hopping(1, 0, 0, 0, 0, 1.0);
    prim_cell.add_hopping(0, 1, 0, 0, 0, 1.0);
    prim_cell.add_hopping(1, 1, 0, 0, 0, 1.0);
    prim_cell.add_hopping(1, -1, 0, 0, 0, 1.0);

    // Create the fractal
    Timer timer;
    timer.tic("make_fractal");
    model_t fractal = top_down(prim_cell, 2, iteration, 3);
    timer.toc("make_fractal");
    timer.report_total_time();
    std::cout << "prim_cell: " << fractal.get_num_orbitals() << " " << fractal.get_num_hopping() << std::endl;
    // fractal.save("fractal_prim_cell");
    return fractal;
}

//******************************************************************************//
//                              Analytical models                               //
//******************************************************************************//
QAH::QAH()
    : AnalyticalModel(2, 0, Eigen::MatrixXd::Identity(3, 3))
{
}

void QAH::build_ham_dense(
    const Eigen::Vector3d& k_point,
    Eigen::MatrixXcd& ham,
    int convention) const
{
    assert((convention == 1) || (convention == 2));
    assert((ham.rows() == num_orb_) && (ham.cols() == num_orb_));
    Eigen::Vector3d k_cart = (k_point.transpose() * reciprocal_lattice_).transpose();
    double kx = k_cart[0];
    double ky = k_cart[1];
    constexpr double t1 = 1.0;
    constexpr double t2 = 1.0;
    constexpr double t3 = 0.5;
    constexpr double m = -1.0;
    ham(0, 0) = m + 2 * t2 * std::cos(kx + ky) + 2 * t3 * (std::sin(kx) + std::sin(ky));
    ham(0, 1) = 2 * t1 * (std::cos(kx) - UIMAG * std::cos(ky));
    ham(1, 0) = std::conj(ham(0, 1));
    ham(1, 1) = -ham(0, 0);
}

void QAH::build_ham_der_dense(
    const Eigen::Vector3d& k_point,
    Eigen::MatrixXcd& dh_dka,
    Eigen::MatrixXcd& dh_dkb,
    Eigen::MatrixXcd& dh_dkc,
    int convention) const
{
    assert((convention == 1) || (convention == 2));
    orbital_index_t num_orb = get_num_orbitals();
    assert((dh_dka.rows() == num_orb) && (dh_dka.cols() == num_orb));
    assert((dh_dkb.rows() == num_orb) && (dh_dkb.cols() == num_orb));
    assert((dh_dkc.rows() == num_orb) && (dh_dkc.cols() == num_orb));
    Eigen::Matrix3d recip_vec = get_reciprocal_lattice();
    Eigen::Vector3d k_cart = (k_point.transpose() * recip_vec).transpose();
    double kx = k_cart[0];
    double ky = k_cart[1];
    constexpr double t1 = 1.0;
    constexpr double t2 = 1.0;
    constexpr double t3 = 0.5;
    constexpr double m = -1.0;
    Eigen::MatrixXcd dh_dkx(num_orb, num_orb);
    Eigen::MatrixXcd dh_dky(num_orb, num_orb);
    Eigen::MatrixXcd dh_dkz(num_orb, num_orb);
    // dh_dkx
    dh_dkx(0, 0) = 2 * std::cos(kx) * t3 - 2 * std::sin(ky + kx) * t2;
    dh_dkx(0, 1) = -2 * std::sin(kx) * t1;
    dh_dkx(1, 0) = std::conj(dh_dkx(0, 1));
    dh_dkx(1, 1) = -dh_dkx(0, 0);
    // dh_dky
    dh_dky(0, 0) = 2 * std::cos(ky) * t3 - 2 * std::sin(ky + kx) * t2;
    dh_dky(0, 1) = 2.0 * UIMAG * std::sin(ky) * t1;
    dh_dky(1, 0) = std::conj(dh_dky(0, 1));
    dh_dky(1, 1) = -dh_dky(0, 0);
    // dh_dkz
    dh_dkz.setZero();
    // convert to fractional coordinates
    df_dk_cart2frac(recip_vec, dh_dkx, dh_dky, dh_dkz, dh_dka, dh_dkb, dh_dkc);
}

CheckerBoard::CheckerBoard(int ny)
    : AnalyticalModel(ny, 0, Eigen::MatrixXd::Identity(3, 3))
    , ny_(ny)
{
}

void CheckerBoard::build_ham_dense(
    const Eigen::Vector3d& k_point,
    Eigen::MatrixXcd& ham,
    int convention) const
{
    assert((convention == 1) || (convention == 2));
    assert((ham.rows() == num_orb_) && (ham.cols() == num_orb_));
    Eigen::Matrix3d recip_vec = get_reciprocal_lattice();
    Eigen::Vector3d k_cart = (k_point.transpose() * recip_vec).transpose();
    double kx = k_cart[0];
    double ky = k_cart[1];
    double t = 1.0;
    double b = 1.0 / ny_;
    Eigen::MatrixXcd h00 = Eigen::MatrixXcd::Zero(ny_, ny_);
    Eigen::MatrixXcd h01 = Eigen::MatrixXcd::Zero(ny_, ny_);
    for (int iy = 0; iy < ny_ - 1; ++iy) {
        h00(iy, iy + 1) = t;
        h00(iy + 1, iy) = t;
    }
    h00(ny_ - 1, 0) = t * std::exp(ky * UIMAG);
    h00(0, ny_ - 1) = t * std::exp(-ky * UIMAG);
    for (int iy = 0; iy < ny_; ++iy) {
        h01(iy, iy) = t * std::exp(-TWOPI * UIMAG * b * static_cast<double>(iy));
    }
    ham = h00 + h01 * std::exp(kx * UIMAG) + h01.transpose().conjugate() * std::exp(-kx * UIMAG);
}

Haldane::Haldane()
    : AnalyticalModel(2, 0, Eigen::MatrixXd::Identity(3, 3))
{
    double sqrt3 = std::sqrt(3);
    Eigen::Matrix3d lat_vec {
        { 0.5, 0.5 * sqrt3, 0.0 },
        { -0.5, 0.5 * sqrt3, 0.0 },
        { 0.0, 0.0, 1.0 }
    };
    lattice_ = lat_vec;
    reciprocal_lattice_ = gen_reciprocal_vectors(lat_vec);
}

void Haldane::build_ham_dense(
    const Eigen::Vector3d& k_point,
    Eigen::MatrixXcd& ham,
    int convention) const
{
    assert((convention == 1) || (convention == 2));
    assert((ham.rows() == num_orb_) && (ham.cols() == num_orb_));
    Eigen::Vector3d k_cart = (k_point.transpose() * reciprocal_lattice_).transpose();
    double k1 = k_cart[0];
    double k2 = k_cart[1];
    double sqrt3 = std::sqrt(3);
    double a = 1 / sqrt3;
    double m = 2.0 / 3.0;
    double t1 = 1.0;
    double t2 = 1.0 / 3.0;
    double phi = PI / 4;

    Eigen::MatrixXcd h0 = Eigen::MatrixXcd::Zero(2, 2);
    Eigen::MatrixXcd h1 = Eigen::MatrixXcd::Zero(2, 2);
    Eigen::MatrixXcd h2 = Eigen::MatrixXcd::Zero(2, 2);

    // mass term
    h0(0, 0) = m;
    h0(1, 1) = -m;

    // 1st nearest neighbor
    h1(1, 0) = t1 * (std::exp(UIMAG * k2 * a) + std::exp(UIMAG * sqrt3 / 2.0 * k1 * a - UIMAG / 2.0 * k2 * a) + std::exp(-UIMAG * sqrt3 / 2.0 * k1 * a - UIMAG / 2.0 * k2 * a));
    h1(0, 1) = std::conj(h1(1, 0));

    // 2nd nearest neighbor
    h2(0, 0) = t2 * std::exp(-UIMAG * phi) * (std::exp(UIMAG * sqrt3 * k1 * a) + std::exp(-UIMAG * sqrt3 / 2.0 * k1 * a + UIMAG * 3.0 / 2.0 * k2 * a) + std::exp(-UIMAG * sqrt3 / 2.0 * k1 * a - UIMAG * 3.0 / 2.0 * k2 * a));
    h2(1, 1) = t2 * std::exp(UIMAG * phi) * (std::exp(UIMAG * sqrt3 * k1 * a) + std::exp(-UIMAG * sqrt3 / 2.0 * k1 * a + UIMAG * 3.0 / 2.0 * k2 * a) + std::exp(-UIMAG * sqrt3 / 2.0 * k1 * a - UIMAG * 3.0 / 2.0 * k2 * a));
    ham = h0 + h1 + h2 + h2.transpose().conjugate();
}
} // namespace tbplas::tests
