#!/usr/bin/env python
# coding: utf-8

# (graphene-nn)=
# # Graphene neighbor-shell hoppings
# 
# This graphene model calculation illustrates a case where one can use `TBModel.set_nn_hops` to set the n'th nearest neighbor hoppings.

# In[ ]:


from pythtb import TBModel, Lattice
import numpy as np


# In[ ]:


lat_vecs = [[1, 0], [1 / 2, np.sqrt(3) / 2]]
orb_vecs = [[1 / 3, 1 / 3], [2 / 3, 2 / 3]]

lat = Lattice(lat_vecs, orb_vecs, periodic_dirs=[0, 1])


# ## First nearest neighbor hopping
# 
# For reference, we first create a graphene tight-binding model with only nearest neighbor hoppings.

# In[ ]:


nn_model = TBModel(lat)

shell_hops = {1: -np.exp(-1)}
nn_model.set_shell_hops(shell_hops=shell_hops, mode="set")


# ## N'th neighbor hopping
#  
# We will now create a graphene tight-binding model and set the hopping amplitudes for the first five nearest neighbors. To do this we will use the `set_shell_hops` method of the `TBModel` class. This method takes a dictionary where the keys are the shell numbers (1 for nearest neighbor, 2 for next nearest neighbor, etc.) and the values are the hopping amplitudes.

# In[ ]:


model = TBModel(lat)

N = 5
shell_hops = {i + 1: -np.exp(-i) for i in range(N)}
model.set_shell_hops(shell_hops=shell_hops, mode="set")

print(model)


# The `visualize` method will allow us to see the connectivity of the model after setting the hoppings.

# In[ ]:


model.visualize()


# Lastly, we will plot the band structure of the model to see the effects of including up to the 5th nearest neighbor hoppings. We expect to see modifications in the band structure as we include more distant neighbor interactions. 

# In[ ]:


path = [[0, 0], [2 / 3, 1 / 3], [1 / 2, 1 / 2], [0, 0]]
label = (r"$\Gamma $", r"$K$", r"$M$", r"$\Gamma $")
nk = 100

fig, ax = nn_model.plot_bands(path, label, nk, bands_label="NN", fig=None, ax=None)
model.plot_bands(
    path, label, nk, bands_label=f"{N}'th NN", fig=fig, ax=ax, ls="--", lc="r"
)
ax.set_title("Graphene Tight-Binding Model")