#!/usr/bin/env python

## compute the bands for bilayer graphene (no g4 or delta)
## third regime

from math import *
import cmath
import numpy
from scipy import optimize

g1=0.1
g3=0.33*g1

def Omega(x,y):
    return(1+2*cmath.exp(-3j/2*x)*cos(sqrt(3)/2*y))

# Hamiltonian
def H(x,y):
    return(numpy.array(\
                       [[0,g1,0,Omega(x,y).conjugate()],\
                        [g1,0,Omega(x,y),0],\
                        [0,Omega(x,y).conjugate(),0,g3*Omega(x,y)*cmath.exp(3j*x)],\
                        [Omega(x,y),0,g3*(Omega(x,y).conjugate())*cmath.exp(-3j*x),0]]\
    ))

# eigenvalues
def eigsH(x,y):
    return(numpy.linalg.eigvals(H(x,y)))

# resolution
nrpoints=101
# xrange
xmin,xmax=2*pi/3-0.004,2*pi/3+0.004
# yrange
ymin,ymax=2*pi/3/sqrt(3)-0.004,2*pi/3/sqrt(3)+0.004

# sample points
x=numpy.linspace(xmin,xmax,nrpoints)
y=numpy.linspace(ymin,ymax,nrpoints)
x,y=numpy.meshgrid(x,y)


## We want a data point at each Fermi point, which must first be computed

# A function that vanishes only at the Fermi points
def fermpt(v):
    eigs=numpy.sort(numpy.real(eigsH(v[0],v[1])))
    return([eigs[1]-eigs[2],0])
# compute Fermi points
## approximate starting value
pf_app=numpy.zeros((4,2))
pf_app[0]=[2*pi/3,2*pi/3/sqrt(3)]
pf_app[1]=[2*pi/3,2/sqrt(3)*acos(1/2-g1*g3/2)]
pf_app[2]=[acos(g1*g3/2-1/2),acos(g1*g3/2-1/2)/sqrt(3)]
pf_app[3]=[4*pi/3-acos(g1*g3/2-1/2),acos(g1*g3/2-1/2)/sqrt(3)]
pf=pf_app
for i in range(0,4):
    pf[i]=optimize.root(fermpt,pf_app[i],tol=1.e-10).x

# Reset the data point closest to each Fermi point to the fermi point
for k in range(0,4):
    mini,minj=0,0
    minval=100
    for i in range(0,nrpoints):
        for j in range(0,nrpoints):
            if(sqrt((x[i,j]-pf[k,0])**2+(y[i,j]-pf[k,1])**2)<minval):
                mini,minj=i,j
                minval=sqrt((x[i,j]-pf[k,0])**2+(y[i,j]-pf[k,1])**2)
    [x[mini,minj],y[mini,minj]]=pf[k]


# data points
z=numpy.zeros((4,nrpoints,nrpoints))
for i in range(0,nrpoints):
    for j in range(0,nrpoints):
        eigs=numpy.sort(numpy.real(eigsH(x[i,j],y[i,j])))
        for k in range(0,4):
            z[k,i,j]=(eigs[k])

# output
for i in range(0,nrpoints):
    for j in range(0,nrpoints):
        for k in [1,2]:
                print("%10f %10f %10f %d   " % (x[i,j],y[i,j],z[k,i,j],k),end='')
        print()
    print()