#! /usr/bin/env python
"""
This example demonstrates how to evaluate the hopping integrals using the
Slater-Koster formulation as implemented in the 'SK' class. We show this
by setting the primitive cell and calculating the band structure of black
phosphorus.
"""

from math import exp

import numpy as np
from numpy.linalg import norm

import tbplas as tb


def calc_hop(sk: tb.SK, rij: np.ndarray, label_i: str,
            label_j: str) -> complex:
    """
    Evaluate the hopping integral <i,0|H|j,r> for single layer black phosphorus.

    Reference:
    https://www.sciencedirect.com/science/article/pii/S0927025617306705

    :param sk: SK instance
    :param rij: displacement vector from orbital i to j in nm
    :param label_i: label of orbital i
    :param label_j: label of orbital j
    :return: hopping integral in eV
    """
    r = norm(rij)
    r0 = 0.2224
    v_sss = -1.59 * exp(-(r - r0) / 0.033)
    v_sps = 2.39 * exp(-(r - r0) / 0.053)
    v_pps = 4.03 * exp(-(r - r0) / 0.058)
    v_ppp = -1.14 * exp(-(r - r0) / 0.053)
    return sk.eval(r=rij, label_i=label_i, label_j=label_j,
                   v_sss=v_sss, v_sps=v_sps,
                   v_pps=v_pps, v_ppp=v_ppp)


def test_bp():
    # Lattice vectors
    vectors = np.array([
        [4.623311297, 0.000000000, 0.000000000],
        [0.000000000, 3.299154095, 0.000000000],
        [0.000000000, 0.000000000, 20.478000000],
    ])

    # Orbital coordinates
    coord_p = np.array([
        [0.839001067, 0.250000000, 0.448627136],
        [0.339001577, 0.750000000, 0.551372864],
        [0.660998423, 0.250000000, 0.551372864],
        [0.160998933, 0.750000000, 0.448627136],
    ])
    orbital_coord = [row for row in coord_p for _ in range(4)]

    # Orbital labels
    orbital_label = ("s", "px", "py", "pz") * 4

    # Orbital energies
    orbital_energy = {"s": -8.80, "px": 0.0, "py": 0.0, "pz": 0.0}

    # Create primitive cell and add orbital
    cell = tb.PrimitiveCell(lat_vec=vectors, unit=tb.ANG)
    for i, label in enumerate(orbital_label):
        coord = orbital_coord[i]
        energy = orbital_energy[label]
        cell.add_orbital(coord, energy=energy, label=label)

    # Get hopping terms in the nearest approximation
    neighbors = tb.find_neighbors(cell, a_max=2, b_max=2, max_distance=1.0)

    # Add hopping terms
    sk = tb.SK()
    for term in neighbors:
        i, j = term.pair
        label_i = cell.get_orbital(i).label
        label_j = cell.get_orbital(j).label
        hop = calc_hop(sk, term.rij, label_i, label_j)
        cell.add_hopping(term.rn, i, j, hop)

    # Test band structure
    k_points = np.array([
        [0.0, 0.0, 0.0],
        [0.5, 0.0, 0.0],
        [0.5, 0.5, 0.0],
        [0.0, 0.5, 0.0],
        [0.0, 0.0, 0.0]
    ])
    k_label = ["G", "X", "S", "Y", "G"]
    k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40, 40])
    k_len, bands = tb.calc_bands(cell, k_path, prefix="bp")
    tb.Visualizer().plot_bands(k_len, bands, k_idx, k_label)


if __name__ == "__main__":
    test_bp()
