#! /usr/bin/env python
"""
Example for constructing twisted bilayer tmdc.
"""

import math
from typing import Tuple

import numpy as np
from numpy.linalg import norm

import tbplas as tb


def calc_twist_angle(i: int) -> float:
    """
    Calculate twisting angle according to ref. [1].

    :param i: twisting angle index
    :return: twisting angle in RADIANs, NOT degrees
    """
    cos_ang = (3 * i ** 2 + 3 * i + 0.5) / (3 * i ** 2 + 3 * i + 1)
    return math.acos(cos_ang)


def calc_twist_angle2(n: int, m: int) -> float:
    """
    Calculate twisting angle according to ref. [2].

    :param n: twisting angle index
    :param m: twisting angle index
    :return: twisting angle in RADIANs, NOT degrees
    """
    cos_ang = (n**2 + 4 * n * m + m**2) / (2 * (n**2 + n * m + m**2))
    return math.acos(cos_ang)


def calc_hetero_lattice(i: int, prim_cell_fixed: tb.PrimitiveCell) -> np.ndarray:
    """
    Calculate Cartesian coordinates of lattice vectors of hetero-structure
    according to ref. [1].

    :param i: parameter controlling the twisting angle
    :param prim_cell_fixed: primitive cell of fixed layer
    :return: (3, 3) float64 array, Cartesian coordinates of hetero-structure
        lattice vectors in NANOMETER
    """
    hetero_lattice = np.array([[i, i + 1, 0],
                               [-(i + 1), 2 * i + 1, 0],
                               [0, 0, 1]])
    hetero_lattice = tb.frac2cart(prim_cell_fixed.lat_vec, hetero_lattice)
    return hetero_lattice


def calc_hetero_lattice2(n: int, m: int,
                         prim_cell_fixed: tb.PrimitiveCell) -> np.ndarray:
    """
    Calculate Cartesian coordinates of lattice vectors of hetero-structure
    according to ref. [2].

    :param n: parameter controlling the twisting angle
    :param m: parameter controlling the twisting angle
    :param prim_cell_fixed: primitive cell of fixed layer
    :return: (3, 3) float64 array, Cartesian coordinates of hetero-structure
        lattice vectors in NANOMETER
    """
    hetero_lattice = np.array([[n, m, 0],
                               [-m, n + m, 0],
                               [0, 0, 1]])
    hetero_lattice = tb.frac2cart(prim_cell_fixed.lat_vec, hetero_lattice)
    return hetero_lattice


def calc_hop(sk: tb.SK, rij: np.ndarray, label_i: str, label_j: str) -> float:
    """
    Calculate hopping parameter according to Slater-Koster relation.

    Reference: Table V of PHYSICAL REVIEW B 92, 205108 (2015).

    :param sk: Slater-Koster table
    :param rij: hopping term generated by find_neighbors in nm
    :param label_i: label of orbital from which hop term is defined
    :param label_j: label of orbital to which hop term is defined
    :return: hopping parameter in eV
    """
    nu_pp_sigma = 2.627
    nu_pp_pi = -0.708
    r_pp_sigma = 0.3128  # unit in NM
    r_pp_pi = 0.2923  # unit in NM
    eta_pp_sigma = 3.859
    eta_pp_pi = 5.724
    r_norm = norm(rij).item()
    v_pp_sigma = nu_pp_sigma * math.exp(-(r_norm / r_pp_sigma)**eta_pp_sigma)
    v_pp_pi = nu_pp_pi * math.exp(-(r_norm / r_pp_pi)**eta_pp_pi)
    return sk.pp(rij, v_pps=v_pp_sigma, v_ppp=v_pp_pi,
                 label_i=label_i, label_j=label_j)


def make_layers(i: int = 1, shift: float = 0.7) -> Tuple[tb.PrimitiveCell, tb.PrimitiveCell]:
    """
    Prepare the layers for assembling the hetero-structure.

    :param i: integer controlling the twisting angle
    :param shift: interlayer distance in nm
    :return: fixed and twisted layers of the hetero-structure
    """
    # Evaluate twisting angle.
    angle = calc_twist_angle(i)

    # To build a twist bilayer tmdc we build the twisted primitive cells for
    # each layer first. The 'fixed' cell which is fixed at z=0 and not rotated
    # can be imported from the material repository directly.
    prim_cell_fixed = tb.make_tmdc()

    # Reset the c-axis since it is only 0.3 nm.
    new_lat_vec = prim_cell_fixed.lat_vec.copy()
    new_lat_vec[2, 2] = 1.5  # new c-axis with length of 1.5 nm.
    prim_cell_fixed.reset_lattice(lat_vec=new_lat_vec, unit=tb.NM)

    # On the contrary, the 'twisted' cell must be rotated counter-clockwise by
    # the twisting angle and shifted towards +z by 0.3349 nanometers, which is
    # done by calling the function 'spiral_prim_cell'.
    prim_cell_twisted = tb.make_tmdc()
    prim_cell_twisted.reset_lattice(lat_vec=new_lat_vec, unit=tb.NM)
    tb.spiral_prim_cell(prim_cell_twisted, angle=angle, shift=shift)

    # Evaluate coordinates of lattice vectors of hetero-structure.
    # The reference papers give the fractional coordinates in basis of 'fixed'
    # primitive cell. However, we want the Cartesian coordinates. This is done
    # by calling the 'calc_hetero_lattice' function.
    hetero_lattice = calc_hetero_lattice(i, prim_cell_fixed)

    # With all the primitive cells ready, we build the 'fixed' and 'twisted'
    # layers by reshaping corresponding cells to the lattice vectors of
    # hetero-structure. This is done by calling 'make_hetero_layer'.
    layer_fixed = tb.make_hetero_layer(prim_cell_fixed, hetero_lattice)
    layer_twisted = tb.make_hetero_layer(prim_cell_twisted, hetero_lattice)

    # Align the layers for better appearance though it is not essential.
    angle = -math.atan(hetero_lattice[0, 1] / hetero_lattice[0, 0])
    tb.spiral_prim_cell(layer_fixed, angle=angle)
    tb.spiral_prim_cell(layer_twisted, angle=angle)
    return layer_fixed, layer_twisted


def make_tbg() -> None:
    """Set up tbg at PrimitiveCell level."""
    # Prepare the layers
    layer_fixed, layer_twisted = make_layers(i=2)

    # Construct interlayer hopping terms
    inter_hop = tb.PCInterHopping(layer_fixed, layer_twisted)
    neighbors = tb.find_neighbors(layer_fixed, layer_twisted,
                                  a_max=1, b_max=1, max_distance=0.75)
    sk = tb.SK()
    for term in neighbors:
        i, j = term.pair
        atom_i, label_i = layer_fixed.orbitals[i].label.split("_")
        atom_j, label_j = layer_twisted.orbitals[j].label.split("_")
        if atom_i == "SA" and atom_j == "SB":
            inter_hop.add_hopping(term.rn, i, j, calc_hop(sk, term.rij, label_i, label_j))

    merged_cell = tb.merge_prim_cell(layer_fixed, layer_twisted, inter_hop)
    merged_cell.plot(with_orbitals=True, hop_as_arrows=False, view="bc")

    # Visualize Moire's pattern
    sample = tb.Sample(tb.SuperCell(merged_cell, dim=(4, 4, 1),
                                    pbc=(True, True, False)))
    sample.plot(with_orbitals=True, hop_as_arrows=False, view="bc")


def main():
    make_tbg()


if __name__ == "__main__":
    main()
