#include "advanced.h"
#include "base/consts.h"
#include "builder/base.h"

namespace tbplas::builder {
OrbitalMap::OrbitalMap(orbital_index_t num_orb_sc) { sc2pc_vec_ = std::vector<id_pc_t>(num_orb_sc); }

void OrbitalMap::add(const id_pc_t& id_pc)
{
    sc2pc_vec_[counter_] = id_pc;
    pc2sc_map_[id_pc] = counter_;
    ++counter_;
}

orbital_index_t OrbitalMap::pc2sc(const id_pc_t& id_pc) const
{
    auto iter = pc2sc_map_.find(id_pc);
    assert(iter != pc2sc_map_.end());
    return (*iter).second;
}

const id_pc_t& OrbitalMap::sc2pc(const orbital_index_t& id_sc) const
{
    assert(sc2pc_vec_.size() > id_sc >= 0);
    return sc2pc_vec_[id_sc];
}

const orbital_index_t& OrbitalMap::get_counter() const { return counter_; }

RegularOrbitalMap::RegularOrbitalMap(cell_index_t dim_a, cell_index_t dim_b, cell_index_t dim_c, orbital_index_t num_orb_pc)
    : dim_a_(dim_a)
    , dim_b_(dim_b)
    , dim_c_(dim_c)
    , num_orb_pc_(num_orb_pc)
    , dim_bc_(dim_b * dim_c)
{
}

void RegularOrbitalMap::add() { ++counter_; }

orbital_index_t RegularOrbitalMap::pc2sc(
    const cell_index_t& i_a,
    const cell_index_t& i_b,
    const cell_index_t& i_c,
    const orbital_index_t& i_o) const
{
  return (static_cast<orbital_index_t>(i_a) * dim_bc_ +
          static_cast<orbital_index_t>(i_b) * dim_c_ + i_c) *
             num_orb_pc_ +
         static_cast<orbital_index_t>(i_o);
}

const orbital_index_t& RegularOrbitalMap::get_counter() const { return counter_; }

void SK::check_p_labels_(const std::string& label)
{
    assert(p_labels_.find(label) != p_labels_.end());
}

void SK::check_d_labels_(const std::string& label)
{
    assert(d_labels_.find(label) != d_labels_.end());
}

Eigen::Vector3d SK::perm_vector_(const Eigen::Vector3d& vector, const char x_new)
{
    Eigen::Vector3d vector_ = vector;
    assert((x_new == 'x') || (x_new == 'y') || (x_new == 'z'));
    if (x_new == 'x') {
        return vector_;
    }

    if (x_new == 'y') {
        vector_ << vector(1), vector(2), vector(0);
        return vector_;
    }

    if (x_new == 'z') {
        vector_ << vector(2), vector(0), vector(1);
        return vector_;
    }
}

std::string SK::remap_label_(const std::string& label, const char x_new)
{
    assert((x_new == 'x') || (x_new == 'y') || (x_new == 'z'));
    if (x_new == 'x') {
        return label;
    }

    std::map<char, char> map_table;

    if (x_new == 'y') {
        map_table = { { 'x', 'z' }, { 'y', 'x' }, { 'z', 'y' } };
    }
    if (x_new == 'z') {
        map_table = { { 'x', 'y' }, { 'y', 'z' }, { 'z', 'x' } };
    }
    std::string new_label;
    new_label.push_back(label[0]);
    for (size_t i = 1; i < label.size(); i++) {
        new_label.push_back(map_table[label[i]]);
    }
    return new_label;
}

std::tuple<double, double, double> SK::eval_dir_cos_(const Eigen::Vector3d& r)
{
    // double norm = std::sqrt(r.transpose() * r);
    double norm = r.norm();
    assert(norm > base::INFINITESIMAL);
    Eigen::Vector3d dir_cos = r / norm;
    return std::make_tuple(dir_cos(0), dir_cos(1), dir_cos(2));
}

complex_t SK::ss(const complex_t& v_sss)
{
    return v_sss;
}

complex_t SK::sp(const Eigen::Vector3d& r, const std::string& label_p, const complex_t& v_sps)
{

    Eigen::Vector3d r_ = r;
    check_p_labels_(label_p);
    double l, m, n;
    complex_t t;
    std::tie(l, m, n) = eval_dir_cos_(r_);
    if (label_p == px_) {
        t = l * v_sps;
    }
    if (label_p == py_) {
        t = m * v_sps;
    }
    if (label_p == pz_) {
        t = n * v_sps;
    }
    return t;
}

complex_t SK::sd(const Eigen::Vector3d& r, const std::string& label_d, const complex_t& v_sds)
{
    check_d_labels_(label_d);
    // Permute the coordinates
    Eigen::Vector3d r_ = r;
    char x_new;
    std::string label_d_ = label_d;
    if (label_d_ == dyz_) {
        x_new = 'y';
    } else if (label_d_ == dzx_) {
        x_new = 'z';
    } else {
        x_new = 'x';
    }
    r_ = perm_vector_(r_, x_new);
    label_d_ = remap_label_(label_d_, x_new);
    assert((label_d_ == dxy_) || (label_d_ == dx2_y2_) || (label_d_ == dz2_));
    // Evaluate the hopping integral
    double l, m, n;
    std::tie(l, m, n) = eval_dir_cos_(r_);
    double l2 = std::pow(l, 2), m2 = std::pow(m, 2), n2 = std::pow(n, 2);
    complex_t t;
    if (label_d_ == dxy_) {
        t = sqrt3_ * l * m * v_sds;
    }
    if (label_d_ == dx2_y2_) {
        t = half_sqrt3_ * (l2 - m2) * v_sds;
    }
    if (label_d_ == dz2_) {
        t = (n2 - 0.5 * (l2 + m2)) * v_sds;
    }
    return t;
}

complex_t SK::pp(const Eigen::Vector3d& r, const std::string& label_i, const std::string& label_j, const complex_t& v_pps, const complex_t& v_ppp)
{
    check_p_labels_(label_i);
    check_p_labels_(label_j);
    std::string label_i_ = label_i;
    std::string label_j_ = label_j;
    char x_new;
    x_new = label_i_[1];
    Eigen::Vector3d r_ = r;
    r_ = perm_vector_(r_, x_new);

    label_i_ = remap_label_(label_i_, x_new);
    label_j_ = remap_label_(label_j_, x_new);

    // After permutation, label_i will always be px.
    // Otherwise, something must be wrong.
    assert(label_i_ == px_);
    // The minimal hopping table in the reference.
    double l, m, n;
    std::tie(l, m, n) = eval_dir_cos_(r_);
    double l2 = std::pow(l, 2);
    complex_t t;
    if (label_j_ == px_) {
        t = l2 * v_pps + (1 - l2) * v_ppp;
    } else if (label_j_ == py_) {
        t = l * m * (v_pps - v_ppp);
    } else {
        t = l * n * (v_pps - v_ppp);
    }
    return t;
}

complex_t SK::pd(const Eigen::Vector3d& r, const std::string& label_p, const std::string& label_d, const complex_t& v_pds, const complex_t& v_pdp)
{
    check_p_labels_(label_p);
    check_d_labels_(label_d);
    std::string label_p_ = label_p;
    std::string label_d_ = label_d;
    // Permute coordinates

    auto perm_labels = std::set { dxy_, dyz_, dzx_ };
    char x_new;
    Eigen::Vector3d r_ = r;
    if (perm_labels.find(label_d_) != perm_labels.end()) {
        x_new = label_p_[1];
        r_ = perm_vector_(r_, x_new);
        label_p_ = remap_label_(label_p_, x_new);
        label_d_ = remap_label_(label_d_, x_new);
    }
    // Evaluate hopping integral
    double l, m, n;
    std::tie(l, m, n) = eval_dir_cos_(r_);
    double l2 = std::pow(l, 2), m2 = std::pow(m, 2), n2 = std::pow(n, 2);
    double l2_p_m2 = l2 + m2;
    double l2_m_m2 = l2 - m2;
    double sqrt3 = sqrt3_;
    double sqrt3_2 = half_sqrt3_;
    complex_t t;
    // 可改成无上下文依赖并行
    if (label_p_ == px_) {
        if (label_d_ == dxy_) {
            t = sqrt3 * l2 * m * v_pds + m * (1 - 2 * l2) * v_pdp;
        } else if (label_d_ == dyz_) {
            t = l * m * n * (sqrt3 * v_pds - 2.0 * v_pdp);
        } else if (label_d_ == dzx_) {
            t = sqrt3 * l2 * n * v_pds + n * (1 - 2 * l2) * v_pdp;
        } else if (label_d_ == dx2_y2_) {
            t = sqrt3_2 * l * l2_m_m2 * v_pds + l * (1 - l2_m_m2) * v_pdp;
        } else {
            t = l * (n2 - 0.5 * l2_p_m2) * v_pds - sqrt3 * l * n2 * v_pdp;
        }
    } else if (label_p_ == py_) {
        assert((label_d_ == dx2_y2_) || (label_d_ == dz2_));
        if (label_d_ == dx2_y2_) {
            t = sqrt3_2 * m * l2_m_m2 * v_pds - m * (1 + l2_m_m2) * v_pdp;
        }
        if (label_d_ == dz2_) {
            t = m * (n2 - 0.5 * l2_p_m2) * v_pds - sqrt3 * m * n2 * v_pdp;
        }
    } else {
        assert((label_d_ == dx2_y2_) || (label_d_ == dz2_));
        if (label_d_ == dx2_y2_) {
            t = sqrt3_2 * n * l2_m_m2 * v_pds - n * l2_m_m2 * v_pdp;
        }
        if ((label_d_ == dz2_)) {
            t = n * (n2 - 0.5 * l2_p_m2) * v_pds + sqrt3 * n * l2_p_m2 * v_pdp;
        }
    }
    return t;
}

complex_t SK::dd(const Eigen::Vector3d& r, const std::string& label_i, const std::string& label_j, const complex_t& v_dds, const complex_t& v_ddp, const complex_t& v_ddd)
{
    std::string label_i_ = label_i;
    std::string label_j_ = label_j;
    check_d_labels_(label_i_);
    check_d_labels_(label_j_);
    /*
        Number the orbitals such that we can filter the diagonal terms
        The order of the orbitals strictly follows the reference.
        DO NOT CHANGE IT UNLESS YOU KNOW WHAT YOU ARE DOING!
    */
    auto d_labels = std::vector { dxy_, dyz_, dzx_, dx2_y2_, dz2_ };
    auto it_i = std::find(d_labels.begin(), d_labels.end(), label_i_);
    auto it_j = std::find(d_labels.begin(), d_labels.end(), label_j_);
    orbital_index_t id_i = 0;
    orbital_index_t id_j = 0;
    assert(it_i != d_labels.end());
    id_i = std::distance(d_labels.begin(), it_i);
    assert(it_j != d_labels.end());
    id_j = std::distance(d_labels.begin(), it_j);
    complex_t t;
    Eigen::Vector3d r_ = r;
    if (id_i > id_j) {
        // Eigen::Vector3d mr = -r;
        t = base::conj_t(dd(-r_, label_j_, label_i_, v_dds, v_ddp, v_ddd));
    } else {
        char x_new;
        double l, m, n, l2, m2, n2, l2_p_m2, l2_m_m2, lm, mn, nl, l2m2, sqrt3, factor, t1, t2, t3;
        // Permute the coordinates if essential
        if ((label_i_ == dyz_) && ((label_j_ == dyz_) || (label_j_ == dzx_))) {
            x_new = 'y';
        } else if ((label_i_ == dzx_) && (label_j_ == dzx_)) {
            x_new = 'z';

        } else {
            x_new = 'x';
        }
        r_ = perm_vector_(r_, x_new);
        label_i_ = remap_label_(label_i_, x_new);
        label_j_ = remap_label_(label_j_, x_new);
        // # Evaluate hopping integral
        std::tie(l, m, n) = eval_dir_cos_(r_);
        l2 = std::pow(l, 2);
        m2 = std::pow(m, 2);
        n2 = std::pow(n, 2);
        l2_p_m2 = l2 + m2;
        l2_m_m2 = l2 - m2;
        lm = l * m;
        mn = m * n;
        nl = n * l;
        l2m2 = l2 * m2;
        sqrt3 = sqrt3_;
        if (label_i_ == dxy_) {
            if (label_j_ == dxy_) {
                factor = 1.0;
                t1 = 3 * l2m2;
                t2 = l2_p_m2 - 4 * l2m2;
                t3 = n2 + l2m2;
            } else if (label_j_ == dyz_) {
                factor = nl;
                t1 = 3 * m2;
                t2 = 1 - 4 * m2;
                t3 = m2 - 1;
            } else if (label_j_ == dzx_) {
                factor = mn;
                t1 = 3 * l2;
                t2 = 1 - 4 * l2;
                t3 = l2 - 1;
            } else if (label_j_ == dx2_y2_) {
                factor = lm * l2_m_m2;
                t1 = 1.5;
                t2 = -2;
                t3 = 0.5;
            } else {
                factor = sqrt3 * lm;
                t1 = n2 - 0.5 * l2_p_m2;
                t2 = -2 * n2;
                t3 = 0.5 * (1 + n2);
            }
        } else if (label_i_ == dyz_) {
            assert((label_j_ == dx2_y2_) || (label_j_ == dz2_));
            if (label_j_ == dx2_y2_) {
                factor = mn;
                t1 = 1.5 * l2_m_m2;
                t2 = -(1 + 2 * l2_m_m2);
                t3 = 1 + 0.5 * l2_m_m2;
            }
            if (label_j_ == dz2_) {
                factor = sqrt3 * mn;
                t1 = n2 - 0.5 * l2_p_m2;
                t2 = l2_p_m2 - n2;
                t3 = -0.5 * l2_p_m2;
            }
        } else if (label_i_ == dzx_) {
            assert((label_j_ == dx2_y2_) || (label_j_ == dz2_));
            if (label_j_ == dx2_y2_) {
                factor = nl;
                t1 = 1.5 * l2_m_m2;
                t2 = 1 - 2 * l2_m_m2;
                t3 = -1 * (1 - 0.5 * l2_m_m2);
            }
            if (label_j_ == dz2_) {
                factor = sqrt3 * nl;
                t1 = n2 - 0.5 * l2_p_m2;
                t2 = l2_p_m2 - n2;
                t3 = -0.5 * l2_p_m2;
            }
        } else if (label_i_ == dx2_y2_) {
            assert((label_j_ == dx2_y2_) || (label_j_ == dz2_));
            if (label_j_ == dx2_y2_) {
                factor = 1;
                t1 = 0.75 * std::pow(l2_m_m2, 2);
                t2 = l2_p_m2 - std::pow(l2_m_m2, 2);
                t3 = n2 + 0.25 * std::pow(l2_m_m2, 2);
            }
            if (label_j_ == dz2_) {
                factor = sqrt3 * l2_m_m2;
                t1 = 0.5 * (n2 - 0.5 * l2_p_m2);
                t2 = -n2;
                t3 = 0.25 * (1 + n2);
            }
        } else {
            assert(label_j_ == dz2_);
            factor = 1;
            t1 = std::pow(n2 - 0.5 * l2_p_m2, 2);
            t2 = 3 * n2 * l2_p_m2;
            t3 = 0.75 * std::pow(l2_p_m2, 2);
        }
        t = factor * (t1 * v_dds + t2 * v_ddp + t3 * v_ddd);
    }
    return t;
}
complex_t SK::ps(const Eigen::Vector3d& r, const std::string& label_p, const complex_t& v_sps)
{
    return base::conj_t(sp(-r, label_p, v_sps));
}

complex_t SK::ds(const Eigen::Vector3d& r, const std::string& label_d, const complex_t& v_sds)
{
    return base::conj_t(sd(-r, label_d, v_sds));
}

complex_t SK::dp(const Eigen::Vector3d& r, const std::string& label_p, const std::string& label_d, const complex_t& v_pds, const complex_t& v_pdp)
{
    // Eigen::Vector3d rm = -r;
    return base::conj_t(pd(-r, label_p, label_d, v_pds, v_pdp));
}

complex_t SK::eval(const Eigen::Vector3d& r, const std::string& label_i, const std::string& label_j,
    const complex_t& v_sss, const complex_t& v_sps, const complex_t& v_sds,
    const complex_t& v_pps, const complex_t& v_ppp, const complex_t& v_pds, const complex_t& v_pdp,
    const complex_t& v_dds, const complex_t& v_ddp, const complex_t& v_ddd)
{
    complex_t t;
    if (label_i == s_) {
        if (label_j == s_) {
            t = ss(v_sss);
        } else if (p_labels_.find(label_j) != p_labels_.end()) {
            t = sp(r, label_j, v_sps);
        } else {
            t = sd(r, label_j, v_sds);
        }
    } else if (p_labels_.find(label_i) != p_labels_.end()) {
        if (label_j == s_) {
            t = ps(r, label_i, v_sps);
        } else if (p_labels_.find(label_j) != p_labels_.end()) {
            t = pp(r, label_i, label_j, v_pps, v_ppp);
        } else {
            t = pd(r, label_i, label_j, v_pds, v_pdp);
        }
    } else {
        if (label_j == s_) {
            t = ds(r, label_i, v_sds);
        } else if (p_labels_.find(label_j) != p_labels_.end()) {
            t = dp(r, label_j, label_i, v_pds, v_pdp);
        } else {
            t = dd(r, label_i, label_j, v_dds, v_ddp, v_ddd);
        }
    }
    return t;
}

complex_t SOCTable::eval(const std::string& label_i, const std::string& spin_i,
    const std::string& label_j, const std::string& spin_j)
{
    std::tuple<std::string, std::string> label_idx = std::make_tuple(label_i, label_j);
    std::tuple<std::string, std::string> spin_idx = std::make_tuple(spin_i, spin_j);
    complex_t product;
    assert(orbital_labels_.find(label_i) != orbital_labels_.end());
    assert(orbital_labels_.find(label_j) != orbital_labels_.end());
    assert(spin_labels_.find(spin_i) != spin_labels_.end());
    assert(spin_labels_.find(spin_j) != spin_labels_.end());
    if (spin_idx == std::make_tuple("up", "up")) {
        product = up_up_[label_idx];
    } else if (spin_idx == std::make_tuple("up", "down")) {
        product = up_down_[label_idx];
    } else if (spin_idx == std::make_tuple("down", "up")) {
        product = down_up_[label_idx];
    } else {
        product = -1.0 * up_up_[label_idx];
    }
    return product;
}
} // namespace tbplas::builder