#!/usr/bin/env python

#
# This code is a complete example of the example given
# on the Wiki page.
#

from __future__ import print_function, division

import sisl
import numpy as np

# Lattice constant
alat = 1.42
length = ( 1.5 ** 2 + 3. / 4. ) ** .5 * alat

# Create carbon atom
C = sisl.Atom(6, R=alat+0.1)

# Create super cell (this is ill-skewed)
sc = sisl.SuperCell([length,length,10,90,90,60], nsc=[3,3,1])
# Rotate super cell
sc = sc.rotate(-30, v=[0,0,1])

# Create graphene unit cell
gr = sisl.Geometry([[0,0,0],[alat,0,0]], atom=C, sc=sc)

# Write to xyz file (repeated a couple of times for visual effect)
gr.tile(3, 0).tile(3, 1).write('graphene_3x3.xyz')

# Create tight-binding object
tb = sisl.Hamiltonian(gr)

# Denote interactions ranges
#     on-site , nearest neighbour
dR = (0.1     , alat + 0.1)
for ia in gr:
    # Get interacting atoms
    idx = gr.close(ia, dR=dR)
    tb[ia,idx[0]] =  0.
    tb[ia,idx[1]] = -2.7

# Now we can calculate the band-structure
nk = 200
kpts = np.linspace(0, 1./3, nk)
eigs = np.empty([nk,2],np.float64)
for ik, k in enumerate(kpts):
    eigs[ik,:] = tb.eigh(k=[k,-k,0])

import matplotlib.pyplot as plt
plt.figure(figsize=(3, 3))
plt.locator_params(nbins=3)
plt.plot(kpts,eigs[:,0])
plt.plot(kpts,eigs[:,1])
plt.savefig('python-graphene-nn-G-K.png', bbox_inches='tight')