#! /usr/bin/env python

import numpy as np

import tbplas as tb


def make_overlap(prim_cell: tb.PrimitiveCell,
                 on_site: float = 1.0,
                 hop: float = 0.0) -> tb.PrimitiveCell:
    """
    Make an overlap for given primitive cell.

    :param prim_cell: primitive cell from which the overlap will be generated
    :param on_site: on-site term of the overlap
    :param hop: hopping terms of the overlap
    :return: overlap with the same numbers of orbitals and hopping terms as the model
    """
    overlap = tb.PrimitiveCell(prim_cell.lat_vec, prim_cell.origin, 1.0)
    for i in range(prim_cell.num_orb):
        orbital = prim_cell.orbitals[i]
        overlap.add_orbital(orbital.position, on_site)

    overlap.add_hopping((0, 0), 0, 1, hop)
    overlap.add_hopping((1, 0), 1, 0, hop)
    overlap.add_hopping((0, 1), 1, 0, hop)
    return overlap


def main():
    prim_cell = tb.make_graphene_diamond()
    overlap = make_overlap(prim_cell, on_site=1.0, hop=0.0)
    #overlap = make_overlap(prim_cell, on_site=0.8, hop=0.1)
    solver = tb.DiagSolver(prim_cell, overlap)

    k_points = np.array([
        [0.0, 0.0, 0.0],
        [2./3, 1./3, 0.0],
        [0.5, 0.0, 0.0],
        [0.0, 0.0, 0.0]
    ])
    k_path, k_idx = tb.gen_kpath(k_points, (1000, 1000, 1000))
    solver.config.prefix = "graphene"
    solver.config.k_points = k_path

    timer = tb.Timer()
    timer.tic("bands")
    k_len, bands = solver.calc_bands()
    timer.toc("bands")

    if solver.is_master:
        timer.report_total_time()
        vis = tb.Visualizer()
        vis.plot_bands(k_len, bands, k_idx, ["G", "K", "M", "G"])


if __name__ == "__main__":
    main()
