"""Utilities for constructing graphene primitive cells."""

import math

import numpy as np

from tbplas.base.consts import NM
from tbplas.base.lattice import gen_lattice_vectors
from tbplas.builder.primitive import PrimitiveCell
from tbplas.builder.advanced import reshape_prim_cell, find_neighbors, SK


__all__ = ["make_graphene_rect", "make_graphene_diamond", "make_graphene_sp",
           "make_graphene_soc"]


def make_graphene_diamond(c: float = 1.0, t: complex = -2.7) -> PrimitiveCell:
    """
    Make graphene primitive cell in diamond shape.

    :param c: length of c-axis in NANOMETER
    :param t: hopping integral in eV
    :return: graphene primitive cell
    """
    vectors = gen_lattice_vectors(a=0.246, b=0.246, c=c, gamma=60)
    cell = PrimitiveCell(vectors, unit=NM)
    cell.add_orbital((0.0, 0.0), label="C_pz")
    cell.add_orbital((1/3., 1/3.), label="C_pz")
    cell.add_hopping((0, 0), 0, 1, t)
    cell.add_hopping((1, 0), 1, 0, t)
    cell.add_hopping((0, 1), 1, 0, t)
    return cell


def make_graphene_rect(from_scratch: bool = True,
                       c: float = 1.0,
                       t: complex = -2.7) -> PrimitiveCell:
    """
    Make graphene cell in rectangular shape.

    :param c: length of c-axis in NANOMETER
    :param from_scratch: method to build the cell
        If true, build the cell from scratch.
        Otherwise, build it by reshaping a primitive cell.
    :param t: hopping integral in eV
    :return: rectangular graphene cell
    """
    if from_scratch:
        # Calculate lattice vectors
        sqrt3 = math.sqrt(3)
        a = 0.246
        cc_bond = sqrt3 / 3 * a
        vectors = gen_lattice_vectors(sqrt3 * cc_bond, 3 * cc_bond, c)

        # Create cell and add orbitals
        cell = PrimitiveCell(vectors, unit=NM)
        cell.add_orbital((0, 0), label="C_pz")
        cell.add_orbital((0, 2/3.), label="C_pz")
        cell.add_orbital((1/2., 1/6.), label="C_pz")
        cell.add_orbital((1/2., 1/2.), label="C_pz")

        # Add hopping terms
        cell.add_hopping((0, 0), 0, 2, t)
        cell.add_hopping((0, 0), 2, 3, t)
        cell.add_hopping((0, 0), 3, 1, t)
        cell.add_hopping((0, 1), 1, 0, t)
        cell.add_hopping((1, 0), 3, 1, t)
        cell.add_hopping((1, 0), 2, 0, t)
    else:
        cell = make_graphene_diamond(t=t)
        lat_frac = np.array([[1, 0, 0], [-1, 2, 0], [0, 0, 1]])
        cell = reshape_prim_cell(cell, lat_frac)
    return cell


def make_graphene_sp(c: float = 1.0) -> PrimitiveCell:
    """
    Make 8-band graphene primitive cell.

    Reference:
    https://journals.aps.org/prb/abstract/10.1103/PhysRevB.82.245412

    :param c: length of c-axis in NANOMETER
    :return: graphene primitive cell
    """
    # Lattice vectors and orbital info.
    vectors = gen_lattice_vectors(a=0.246, b=0.246, c=c, gamma=60)
    orb_pos = np.array([[0.0, 0.0, 0.0], [1/3., 1/3., 0.0]])
    orb_label = ("s", "px", "py", "pz")
    orb_eng = {"s": -8.370, "px": 0.0, "py": 0.0, "pz": 0.0}

    # Create cell and add orbitals
    cell = PrimitiveCell(vectors, unit=NM)
    for pos in orb_pos:
        for label in orb_label:
            cell.add_orbital(pos, energy=orb_eng[label], label=label)

    # Add hopping terms
    neighbors = find_neighbors(cell, a_max=5, b_max=5, max_distance=0.25)
    d1 = 0.14202816622064793  # 1st nearest hopping distance in NM
    d2 = 0.24599999999999997  # 2nd neatest hopping distance in NM
    sk = SK()
    for term in neighbors:
        # Extract info.
        i, j = term.pair
        label_i = cell.get_orbital(i).label
        label_j = cell.get_orbital(j).label
        distance = term.distance
        rij = term.rij

        # Evaluate hopping integral using Slater-Koster formulation
        if abs(distance - d1) < 1.0e-5:
            v_sss, v_sps, v_pps, v_ppp = -5.729, 5.618, 6.050, -3.070
        elif abs(distance - d2) < 1.0e-5:
            v_sss, v_sps, v_pps, v_ppp = 0.102, -0.171, -0.377, 0.070
        else:
            raise ValueError(f"Too large distance {distance}")
        hop = sk.eval(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)

        # Add hopping term
        cell.add_hopping(term.rn, i, j, hop)
    return cell


def make_graphene_soc(is_qsh: bool = True) -> PrimitiveCell:
    """
    Set up the primitive cell of graphene according to Kane-Mele model.

    Ref: https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.95.146802

    :param is_qsh: whether is the model is in quantum spin Hall phase
    :return: the primitive cell of graphene
    """
    # Parameters
    lat = 1.
    t = -1.
    lamb_so = 0.06 * t
    lamb_r = 0.05 * t
    if is_qsh:
        lamb_nu = 0.1 * t  # QSH phase
    else:
        lamb_nu = 0.4 * t  # normal insulating phase

    # Function for generating unit vector in xOy plane
    def _ev(_theta):
        _theta *= (math.pi / 180)
        return np.array([math.cos(_theta), math.sin(_theta), 0.0])

    # Lattice vectors
    vectors = np.array([_ev(-30)*lat, _ev(30)*lat, [0, 0, 1]])
    prim_cell = PrimitiveCell(vectors, unit=1.0)

    # Add orbitals
    prim_cell.add_orbital((1./3, 1./3), energy=lamb_nu, label="A+")
    prim_cell.add_orbital((1./3, 1./3), energy=lamb_nu, label="A-")
    prim_cell.add_orbital((2./3, 2./3), energy=-lamb_nu, label="B+")
    prim_cell.add_orbital((2./3, 2./3), energy=-lamb_nu, label="B-")

    # Pauli matrices
    sigma_x = np.array([[0, 1], [1, 0]])
    sigma_y = np.array([[0, -1j], [1j, 0]])
    sigma_z = np.array([[1, 0], [0, -1]])

    # Add 1st nearest hopping terms, which are ordinary hopping + Rashba SOC
    # dij_table is from sub-lattice A to B.
    dij_table = {(-1, 0, 0): _ev(120), (0, -1, 0): _ev(-120), (0, 0, 0): _ev(0)}
    for rn, dij in dij_table.items():
        hop_mat = t * np.eye(2, dtype=np.complex128)
        hop_mat += 1j * lamb_r * (dij[1] * sigma_x - dij[0] * sigma_y)
        for ii in range(2):
            for jj in range(2):
                prim_cell.add_hopping(rn, ii, jj+2, hop_mat.item(ii, jj))

    # Add 2nd nearest hopping terms, which are inter-site SOC
    # vij_table is for sub-lattice A. For B it is the opposite.
    vij_table = {(-1, 1, 0): -1, (-1, 0, 0): 1, (0, -1, 0): -1}
    for rn, vij in vij_table.items():
        hop_mat = 1j * lamb_so * vij * sigma_z
        for ii in range(2):
            for jj in range(2):
                hij = hop_mat.item(ii, jj)
                prim_cell.add_hopping(rn, ii, jj, hij)
                prim_cell.add_hopping(rn, ii+2, jj+2, -hij)
    return prim_cell
