#!/usr/bin/env python

# Toy graphene model with radiation

# By Luis E. F. Foa Torres (June 2013)
# This code is built on the basis of the pythtb module by Sinisa Coh and David Vanderbilt.

# This code is provided AS IT IS and for educational purposes, use it at your own risk.

#Refs:
# Floquet chiral edge states in graphene
#P.M. Perez-Piskunow, G. Usaj, C. A. Balseiro, and L. E. F. Foa Torres
#Physical Review B, 89, 121401(R) (2014)
#
#Irradiated graphene as a tunable Floquet topological insulator
#G. Usaj, P. M. Perez-Piskunow, L. E. F. Foa Torres, and C. A. Balseiro
#Physical Review B, 90, 115423 (2014)
#
#Multiterminal Conductance of a Floquet topological Insulator
#L. E. F. Foa Torres, P. M. Perez-Piskunow, C. A. Balseiro, and G. Usaj
#Physical Review Letters, 113, 266801 (2014)
#
#Floquet topological transitions in a one-dimensional topological insulator
#V. Dal Lago, M. Atala and L. E. F. Foa Torres
#Physical Review A 92, 023624 (2015).
#
#Hierarchy of edge states in driven honeycomb lattices
#P. M. Perez-Piskunow, L. E. F. Foa Torres, and G. Usaj
#Physical Review A 91, 043625 (2015).

from pythtb import * # import TB model class
import numpy as np
import pylab as pl
import math as math

from scipy.special import jn, jn_zeros


# set model parameters
delta=0.0
hbarOmega=0.7
t=-1.0
zx=0.15 # parameter for the laser intensity x component
zy=0.15 # parameter for the laser intensity y component
phi=-math.pi/2.0
N_ribbon=50 # ribbon width
N_Floquet_ch=7 # Number of Floquet channels
N_Floquet_min=-3 # lowest Floquet channel
nk=500

# define lattice vectors
lat=[[np.sqrt(3.0)/2.0,-0.5,0.0],[np.sqrt(3.0)/2.0,0.5,0.0],[0.0,0.0,1.]]
# define coordinates of orbitals. These coordinates are in reduced units. [1,0,0] means 1 times the first lattice vector
orb=[[1./3.,1./3.,0.],[2./3.,2./3.,0.]]
original_orb=[[1./3.,1./3.,0.],[2./3.,2./3.,0.]]
# make three dimensional tight-binding graphene model: 2d graphene + Floquet channels, 2 periodic directions in 3d space
my_model=tb_model(3,3,lat,orb)

# set hoppings (one for each connected pair of orbitals)
# (amplitude, i, j, [lattice vector to cell containing j])
zx_tilde=zx/2.0
zy_tilde=zy*np.sqrt(3.)/2.0
for m in range(-3*N_Floquet_ch,3*N_Floquet_ch+1):
    hop1=0.0+0.0j
    hop2=0.0+0.0j
    hop3=0.0+0.0j
    for n in range(-4*N_Floquet_ch,4*N_Floquet_ch+1): #
 	hop1 += (1.0j)**n*jn(n,zx_tilde+zy_tilde*np.cos(phi))*jn(m-n,-zy_tilde*np.sin(phi))
	hop3 += (1.0j)**n*jn(n,zx_tilde-zy_tilde*np.cos(phi))*jn(m-n,zy_tilde*np.sin(phi))

    hop2 = (1.0j)**m * jn(m,-zx)
    my_model.set_hop(t*hop1, 0, 1, [ 0, -1, -m]) #  
    my_model.set_hop(t*hop2, 0, 1, [ 0, 0, -m]) #  
    my_model.set_hop(t*hop3, 0, 1, [ -1, 0, -m]) #  

tbmod=my_model.make_supercell([[1, 0, 0],[-1, 2, 0],[0, 0, 1]])

# now we further restrict the dimensionality 
ribbon_model=tbmod.cut_piece(N_ribbon,1,glue_edgs=False)


fin_model=ribbon_model.cut_piece(N_Floquet_ch,2,glue_edgs=False)

array_onsite=np.zeros(fin_model._norb)
for ii in range(ribbon_model._norb):
	for iii in range(N_Floquet_ch):
		array_onsite[ii+iii*ribbon_model._norb]=(N_Floquet_min+iii)*hbarOmega
fin_model.set_onsite(array_onsite,mode="add")



# generate list of k-points following some high-symmetry line in
# the k-space. Variable kpts here is just an array of k-points
path='full' #[-1.0,1.0]
kpts=k_path(path,nk)
print '---------------------------------------'
print 'report of k-point path'
print '---------------------------------------'
print 'Path runs over',len(kpts),'k-points connecting:'
for k in path:
    print k
print
    
print '---------------------------------------'
print 'starting calculation'
print '---------------------------------------'
print 'Calculating bands...'


evals_file = open('evals.dat', "w")
evals_file.close()
evals_file = open('evals.dat', "a")


for k in kpts:
	(evals,evecs)=fin_model.solve_one(k,eig_vectors=True)


	weight=np.zeros_like(evals)

	for ii in range(fin_model._norb):
		vector=[0. for mm in range(len(orb))]
		vector[:]=evecs[ii,:]
		for site_index in range(ribbon_model._norb):
			weight[ii]+=np.absolute(vector[site_index+(0-N_Floquet_min)*ribbon_model._norb])**2

	# Save data to files 
#This saves k, the eigenvalues and the corresponding weight of the eigenstate on the n=0 Floquet channel
	for ii in range(fin_model._norb):
		data=str(k)+'   '+str(evals[ii])+'   '+str(weight[ii])
		evals_file.write(data+ " \n")
#	np.savetxt('weights.dat.gz', weight, fmt='%1.7e')
	


print 'Done.\n'