#! /usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


cell = tb.wan2pc("graphene", hop_eng_cutoff=0.0)
k_points = np.array([
    [0.0, 0.0, 0.0],
    [1./3, 1./3, 0.0],
    [1./2, 0.0, 0.0],
    [0.0, 0.0, 0.0],
])
k_label = ["G", "K", "M", "G"]
k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40])
k_len, bands = tb.calc_bands(cell, k_path)

num_bands = bands.shape[1]
for i in range(num_bands):
    plt.plot(k_len, bands[:, i], color="r", linewidth=1.0)
for idx in k_idx:
    plt.axvline(k_len[idx], color="k", linewidth=1.0)
plt.xlim((0, np.amax(k_len)))
plt.xticks(k_len[k_idx], k_label)
plt.xlabel("k / (1/nm)")
plt.ylabel("Energy (eV)")
plt.tight_layout()
plt.show()
