#include "model_sample.h"

#include "base/lattice.h"
#include "base/utils.h"
#include "builder/advanced.h"
#include "builder/base.h"
#include "builder/filters.h"
#include "builder/primitive.h"
#include "builder/sample.h"
#include "materials/graphene.h"

#include "mask.h"
#include "model.h"

namespace tbplas::tests {

using base::frac2cart;
using base::PI;
using base::rotate_coord;
using base::Timer;
using builder::AbstractHoppingModifier;
using builder::AbstractOrbitalFilter;
using builder::AbstractOrbitalModifier;
using builder::cell_index_t;
using builder::complex_t;
using builder::dim_t;
using builder::HoppingData;
using builder::Orbital;
using builder::orbital_index_t;
using builder::OrbitalData;
using builder::PrimitiveCell;
using builder::Sample;
using materials::make_graphene_diamond;

sample_t make_graphene_sample(
    double t,
    double a,
    bool with_onsite,
    const std::tuple<int, int, int>& dim)
{
    auto prim_cell = make_graphene_pc(t, a, with_onsite);
    Sample<complex_t> sample;
    auto sc = sample.add_super_cell(prim_cell, dim, { true, true, false });
    sample.add_sc_hopping(sc, sc, 0.15, t);
    return sample;
}

/**
 * @brief Filter to reserve orbitals within a circle.
 *
 */
class CircleFilter : public builder::AbstractOrbitalFilter {
public:
    explicit CircleFilter(
        const Eigen::Matrix3d& lattice,
        const Eigen::Vector3d& origin,
        const Eigen::Vector3d& center,
        const double& radius)
        : lattice_(lattice)
        , origin_(origin)
        , center_(center)
        , radius_(radius)
    {
    }

    void act(std::vector<Orbital>& full_orbitals) const final
    {
        std::vector<Orbital> new_orbitals;
        for (const auto& orb : full_orbitals) {
            Eigen::Vector3d cart_pos = (orb.get_position().transpose() * lattice_).transpose() + origin_;
            Eigen::Vector3d dr = cart_pos - center_;
            if (dr.norm() <= radius_) {
                new_orbitals.push_back(orb);
            }
        }
        full_orbitals = new_orbitals;
    }

private:
    Eigen::Matrix3d lattice_ = Eigen::Matrix3d::Identity();
    Eigen::Vector3d origin_ = Eigen::Vector3d::Zero();
    Eigen::Vector3d center_ = Eigen::Vector3d::Zero();
    double radius_ = 0.0;
};

/**
 * @brief Twisting operation to update orbital positions.
 *
 */
class TwistModifier : public AbstractOrbitalModifier {
public:
    explicit TwistModifier(
        const Eigen::Vector3d& center,
        const double& angle,
        const double& shift)
        : center_(center)
        , angle_(angle)
        , shift_(shift)
    {
    }

    void act(builder::OrbitalData& orb_data) const final
    {
        orb_data.orb_pos = rotate_coord(orb_data.orb_pos, angle_, "z", center_);
        orb_data.orb_pos.colwise() += Eigen::Vector3d(0.0, 0.0, shift_);
    }

private:
    Eigen::Vector3d center_ = Eigen::Vector3d::Zero();
    double angle_ = 0.0;
    double shift_ = 0.0;
};

/**
 * @brief Slater-Koster parameter calculator to update hopping terms.
 *
 */
class SKTable : public AbstractHoppingModifier {
public:
    void act(
        const OrbitalData& data_bra,
        const OrbitalData& data_ket,
        HoppingData& hop_data) const final
    {
        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;

        for (size_t i = 0; i < hop_data.get_num_hopping(); ++i) {
            double dr = hop_data.dr.col(i).norm();
            double n = hop_data.dr(2, i) / 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;
            hop_data.hop_eng[i] = hop;
        }
    }
};

/**
 * @brief Filter to reserve orbitals within a circle.
 *
 */
class FractalFilter : public AbstractOrbitalFilter {
public:
    explicit FractalFilter(
        const dim_t& dim,
        const int& start_width,
        const int& iteration,
        const int& extension)
        : dim_(dim)
        , start_width_(start_width)
        , iteration_(iteration)
        , extension_(extension)
    {
    }

    void act(std::vector<Orbital>& full_orbitals) const final
    {
        int final_width = std::get<0>(dim_);
        Box start_box(0, 0, final_width - 1, final_width - 1, false);
        Mask mask(start_box, extension_, iteration_);
        mask.etch_orbitals(full_orbitals, final_width);
    }

private:
    dim_t dim_ = { 1, 1, 1 };
    int start_width_ = 2;
    int iteration_ = 3;
    int extension_ = 3;
};

sample_t make_tbg_sample(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 cell
    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();

    // Make supercells
    Sample<complex_t> sample;
    auto sc_fixed = sample.add_super_cell(fixed_cell, hetero_lattice, { true, true, false });
    auto sc_twisted = sample.add_super_cell(twisted_cell, hetero_lattice, { true, true, false });

    // Make intra and inter hopping generators
    auto hop_fixed = sample.add_sc_hopping(sc_fixed, sc_fixed, 0.75);
    auto hop_twisted = sample.add_sc_hopping(sc_twisted, sc_twisted, 0.75);
    auto hop_inter = sample.add_sc_hopping(sc_fixed, sc_twisted, 0.75);
    SKTable sk;
    hop_fixed->add_modifier(sk);
    hop_twisted->add_modifier(sk);
    hop_inter->add_modifier(sk);

    sample.init_array();
    timer.toc("make_tbg");
    timer.report_total_time();
    std::cout << "sample: " << sample.get_num_orbitals() << " " << sample.get_num_hopping() << std::endl;
    // sample.save("tbg_sample");
    return sample;
}

sample_t make_quasi_crystal_sample(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();

    // 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();

    // Make supercells
    Sample<complex_t> sample;
    auto sc_fixed = sample.add_super_cell(prim_cell, dim, { false, false, false });
    auto sc_twisted = sample.add_super_cell(prim_cell, dim, { false, false, false });
    CircleFilter circle(sc_fixed->get_sc_lattice(), sc_fixed->get_sc_origin(), center, radius);
    TwistModifier fix(center, 0.0, 0.0);
    TwistModifier twist(center, angle, 0.3349);
    sc_fixed->add_filter(circle);
    sc_fixed->add_modifier(fix);
    sc_twisted->add_filter(circle);
    sc_twisted->add_modifier(twist);

    // Make intra and inter hopping generators
    auto hop_fixed = sample.add_sc_hopping(sc_fixed, sc_fixed, 0.75);
    auto hop_twisted = sample.add_sc_hopping(sc_twisted, sc_twisted, 0.75);
    auto hop_inter = sample.add_sc_hopping(sc_fixed, sc_twisted, 0.75);
    SKTable sk;
    hop_fixed->add_modifier(sk);
    hop_twisted->add_modifier(sk);
    hop_inter->add_modifier(sk);

    sample.init_array();
    timer.toc("make_quasi_crystal");
    timer.report_total_time();
    std::cout << "sample: " << sample.get_num_orbitals() << " " << sample.get_num_hopping() << std::endl;
    // sample.save("quasi_crystal_sample");
    return sample;
}

sample_t make_fractal_sample(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");
    int start_width = 2;
    int extension = 3;
    int final_width = start_width * static_cast<int>(std::pow(extension, iteration));
    dim_t dim = { final_width, final_width, 1 };
    Sample<complex_t> sample;
    auto fractal = sample.add_super_cell(prim_cell, dim, { false, false, false });
    FractalFilter filter(dim, start_width, iteration, extension);
    fractal->add_filter(filter);
    sample.add_sc_hopping(fractal, fractal, 0.15);
    sample.init_array();
    timer.toc("make_fractal");
    timer.report_total_time();
    std::cout << "sample: " << sample.get_num_orbitals() << " " << sample.get_num_hopping() << std::endl;
    // sample.save("fractal_sample");
    return sample;
}

} // namespace tbplas::tests