#! /usr/bin/env python

import math

import numpy as np

import tbplas as tb


def make_empty_cell(algo: int = 0) -> tb.PrimitiveCell:
    """
    Generate lattice vectors of monolayer graphene.
    :param algo: algorithm to generate lattice vectors
    :return: empty primitive cell
    """
    a = 2.46  # cell length a in Angstrom
    b = a  # cell length b in Angstrom
    c = 10.0  # cell length c in Angstrom
    gamma = 60.0  # angle between a1 and a2 in Degrees
    a_half = a * 0.5
    sqrt3 = math.sqrt(3)
    if algo == 0:
        vectors = tb.gen_lattice_vectors(a=a, b=b, c=c, gamma=gamma)
    elif algo == 1:
        vectors = np.array([
            [a, 0, 0,],
            [a_half, sqrt3*a_half, 0],
            [0, 0, c]])
    else:
        vectors = np.array([
            [sqrt3*a_half, -a_half, 0],
            [sqrt3*a_half, a_half, 0],
            [0, 0, c]])
    cell = tb.PrimitiveCell(vectors, unit=tb.ANG)
    return cell


def add_orbitals(cell: tb.PrimitiveCell, algo: int = 0) -> None:
    """
    Add orbitals to the model.
    :param algo: algorithm to add orbitals
    :return: None. The incoming cell is modified.
    """
    if algo == 0:
        cell.add_orbital((1./3, 1./3),  energy=0.0, label="pz")
        cell.add_orbital((2./3, 2./3),  energy=0.0, label="pz")
    elif algo == 1:
        cell.add_orbital((1./3, 1./3, 0.5), energy=0.0, label="pz")
        cell.add_orbital((2./3, 2./3, 0.5), energy=0.0, label="pz")
    else:
        # NOTE: Tt is assumed that a_1 should be parallel to x-axis
        cell.add_orbital_cart((1.23, 0.71014083), unit=tb.ANG, energy=0.0, label="pz")
        cell.add_orbital_cart((2.46, 1.42028166), unit=tb.ANG, energy=0.0, label="pz")


def add_hopping_terms(cell: tb.PrimitiveCell, algo: int = 0) -> None:
    """
    Add hopping terms the model.
    :param algo: algorithm to add hopping terms
    :return: None. The incoming cell is modified.
    """
    if algo == 0:
        cell.add_hopping(rn=(0, 0), orb_i=0, orb_j=1, energy=-2.7)
        cell.add_hopping(rn=(1, 0), orb_i=1, orb_j=0, energy=-2.7)
        cell.add_hopping(rn=(0, 1), orb_i=1, orb_j=0, energy=-2.7)
    else:
        cell.add_hopping(rn=(0, 0, 0), orb_i=0, orb_j=1, energy=-2.7)
        cell.add_hopping(rn=(1, 0, 0), orb_i=1, orb_j=0, energy=-2.7)
        cell.add_hopping(rn=(0, 1, 0), orb_i=1, orb_j=0, energy=-2.7)


def main() -> None:
    cell = make_empty_cell(algo=0)
    add_orbitals(cell, algo=0)
    add_hopping_terms(cell, algo=0)

    # Plot the model
    cell.plot()

    # Print model details
    cell.print()

    # Plot the analytical Hamiltonian
    print("\nAnalytical Hamiltonian:\n")
    cell.print_hk()

    # Print Python and C++ code of analytical Hamiltonian
    print("\nAnalytical Hamiltonian for Python:\n")
    cell.print_py()
    print("\nAnalytical Hamiltonian for C++:\n")
    cell.print_cxx()


if __name__ == "__main__":
    main()
