"""
Functions for making models at PrimitiveCell level.
"""
from typing import Tuple

import numpy as np

import tbplas as tb


__all__ = ["make_graphene_pc", "make_graphene", "make_overlap", "make_haldane"]


def make_graphene_pc(t: float = -2.7,
                     a: float = 0.142,
                     with_onsite: bool = False) -> tb.PrimitiveCell:
    """
    Make graphene primitive cell with non-zero onsite energies.

    :param t: hopping integral in eV
    :param a: C-C bond length in NM
    :param with_onsite: whether to use non-zero onsite energies
    :return: graphene primitive cell
    """
    lat = a * np.sqrt(3)
    vectors = tb.gen_lattice_vectors(a=lat, b=lat, c=1.0, gamma=60)
    cell = tb.PrimitiveCell(vectors, unit=tb.NM)
    if with_onsite:
        cell.add_orbital((0.0, 0.0), energy=-0.1, label="C_pz")
        cell.add_orbital((1 / 3., 1 / 3.), energy=0.1, label="C_pz")
    else:
        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(t: float = -2.7,
                  a: float = 0.142,
                  with_onsite: bool = False,
                  dim: Tuple[int, int, int] = (1024, 1024, 1)) -> tb.PrimitiveCell:
    """
    Make a large graphene model by extending the primitive cell.

    :param t: hopping integral in eV
    :param a: C-C bond length in NM
    :param with_onsite: whether to use non-zero onsite energies
    :param dim: dimension of the model
    :return: graphene model with given dimension
    """
    prim_cell = make_graphene_pc(t, a, with_onsite)
    prim_cell = tb.extend_prim_cell(prim_cell, dim=dim)
    return prim_cell


def make_overlap(model: tb.PrimitiveCell) -> tb.PrimitiveCell:
    """
    Make an overlap for given primitive cell.

    :param model: primitive cell from which the overlap will be generated
    :return: overlap with the same numbers of orbitals and hopping terms as the model
    """
    # As a demo, overlap has only diagonal (on-site) terms which equals 1.0.
    overlap = tb.PrimitiveCell(model.lat_vec, model.origin, 1.0)
    for i in range(model.num_orb):
        orbital = model.orbitals[i]
        overlap.add_orbital(orbital.position, 1.0)
    # Must have a fake hopping term to avoid empty hopping arrays,
    # which causes c++ extension fail to work.
    overlap.add_hopping((0, 0, 0), 0, 1, 0.0)
    return overlap


def make_haldane() -> tb.PrimitiveCell:
    """
    Make Haldane model for testing Berry utilities.

    Reference: examples/haldane.py of pythtb
    """
    lat = 1.0
    delta = 0.2
    t = -1.0
    t2 = 0.15 * np.exp(1j * np.pi / 2.)
    t2c = t2.conjugate()

    vectors = tb.gen_lattice_vectors(a=lat, b=lat, c=1.0, gamma=60)
    cell = tb.PrimitiveCell(vectors, unit=tb.NM)
    cell.add_orbital((1 / 3., 1 / 3.), energy=-delta)
    cell.add_orbital((2 / 3., 2 / 3.), energy=delta)
    cell.add_hopping((0, 0), 0, 1, t)
    cell.add_hopping((1, 0), 1, 0, t)
    cell.add_hopping((0, 1), 1, 0, t)
    cell.add_hopping((1, 0), 0, 0, t2)
    cell.add_hopping((1, -1), 1, 1, t2)
    cell.add_hopping((0, 1), 1, 1, t2)
    cell.add_hopping((1, 0), 1, 1, t2c)
    cell.add_hopping((1, -1), 0, 0, t2c)
    cell.add_hopping((0, 1), 0, 0, t2c)
    return cell
