#! /usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt

import tbplas as tb


# In this tutorial we show how to calculate the band structure and DOS of
# primitive cell. First, we import graphene primitive cell in diamond shape
# from the material repository.
cell = tb.make_graphene_diamond()

# Then we create a path in the 1st Brillouin zone.
k_points = np.array([
        [0.0, 0.0, 0.0],    # Gamma
        [1./2, 0.0, 0.0],   # M
        [2./3, 1./3, 0.0],  # K
        [0.0, 0.0, 0.0],    # Gamma
]) 
k_label = ["G", "M", "K", "G"]
k_path, k_idx = tb.gen_kpath(k_points, [40, 40, 40])

# Then we evaluate the band structure.
use_diag_solver = False
if use_diag_solver:
    solver = tb.DiagSolver(cell)
    solver.config.k_points = k_path
    solver.config.prefix = "graphene"
    k_len, bands = solver.calc_bands()
else:
    k_len, bands = tb.calc_bands(cell, k_path, prefix="graphene")

# The band structure can be plotted using the Visualizer class.
vis = tb.Visualizer()
vis.plot_bands(k_len, bands, k_idx, k_label)

# Or using matplotlib.
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()

# Now we calculate DOS. We need to create a mesh-grid in 1st Brillouin zone
# first.
k_mesh = tb.gen_kmesh((120, 120, 1))
if use_diag_solver:
    solver = tb.DiagSolver(cell)
    solver.config.k_points = k_mesh
    solver.config.prefix = "graphene"
    solver.config.e_min = -10
    solver.config.e_max = 10
    energies, dos = solver.calc_dos()
else:
    energies, dos = tb.calc_dos(cell, k_mesh, prefix="graphene", e_min=-10, e_max=10)

# Visualize DOS using Visualizer
vis = tb.Visualizer()
vis.plot_dos(energies, dos)

# Or using matplotlib.
plt.plot(energies, dos, color="r", linewidth=1.0)
plt.xlabel("Energy (eV)")
plt.ylabel("DOS (1/eV")
plt.tight_layout()
plt.show()
plt.close()
