using Plots

a = 0.5
x0 = (i*sqrt(3)*a for i in -3:3)
x1 = (i*sqrt(3)*a + sqrt(3)*a/2 for i in -4:3)
y0 = ( 2j   *1.5*a for j in -2:2)
y1 = ((2j+1)*1.5*a for j in -2:1)

c = "red"
lw = 0.5
ls = :dash

plot(size=(500,500), legend=false)

for y in y0
    for x in x0
        plot!([x, x            ], [y, y+a  ], color=c, lw=lw, ls=ls)
        plot!([x, x-sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw, ls=ls)
        plot!([x, x+sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw, ls=ls)
    end
end

for y in y1
    for x in x1
        plot!([x, x            ], [y, y+a  ], color=c, lw=lw, ls=ls)
        plot!([x, x-sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw, ls=ls)
        plot!([x, x+sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw, ls=ls)
    end
end

plot!(xlim=(-3,3), ylim=(-2.5,3))
plot!(aspect_ratio=:equal)

using Plots

a = 0.5
x0 = (i*sqrt(3)*a for i in -2:2)
x1 = (i*sqrt(3)*a + sqrt(3)*a/2 for i in -3:1)
y0 = ( 2j   *1.5*a for j in -1:1)
y1 = ((2j+1)*1.5*a for j in -1:0.5)

c = "red"
b = "blue"
lw = 0.5
ls = :dash

plot(size=(500,500), legend=false)

for y in y0
    for x in x0
        plot!([x, x            ], [y, y+a  ], color=c, lw=lw, ls=ls)
        plot!([x],[y],marker=:circle,color=c)
        plot!([x-sqrt(3)*a/2], [y-a/2],marker=:circle,color=b)
        plot!([x, x-sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw)
        plot!([x, x+sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw,ls=ls)
    end
end

for y in y1
    for x in x1
        plot!([x, x            ], [y, y+a  ], color=c, lw=lw, ls=ls)
        plot!([x, x-sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw)
        plot!([x],[y],marker=:circle,color=c)
        plot!([x-sqrt(3)*a/2], [y-a/2],marker=:circle,color=b)
        plot!([x, x+sqrt(3)*a/2], [y, y-a/2], color=c, lw=lw, ls=ls)
    end
end


plot!(aspect_ratio=:equal)


function calc_HGraphene(Nx,Ny,μ)
    N = Nx*Ny*2
    mat_Htb = zeros(Float64,N,N)   
    mat_Htb += (-μ)*Matrix(I,N,N) #eye(N)
    t = 1.0
    for ix =1:Nx
        for iy=1:Ny            
            for dx in -1:1
                for dy in -1:1
                    jx = ix + dx
                    
                    jx += ifelse(jx > Nx,-Nx,0)
                    jx += ifelse(jx < 1,Nx,0)
                    
                    jy = iy + dy
                    jy += ifelse(jy > Ny,-Ny,0)
                    jy += ifelse(jy < 1,Ny,0)
                    
 
                    for a=1:2
                        b = ifelse(a ==1,2,1)
                        ii = ((iy-1)*Nx+ix-1)*2+a
                        jj = ((jy-1)*Nx+jx-1)*2+b

                        if dx == 0 && dy == 0
                            mat_Htb[ii,jj] = t
                        elseif dx == +1 && dy==0 && a ==1
                            mat_Htb[ii,jj] = t
                        elseif dx == 0 && dy == 1 && a ==1
                            mat_Htb[ii,jj] = t
                        elseif dx == -1 && dy ==0 && a ==2
                            mat_Htb[ii,jj] = t
                        elseif dx ==0 && dy == -1 && a ==2
                            mat_Htb[ii,jj] = t
                        end
                    end
                    

                end
            end
        end
    end
                    
        
    return mat_Htb
end

using LinearAlgebra
Nx = 40
Ny = 40
μ = 0
mat_H = calc_HGraphene(Nx,Ny,μ)
#println(mat_H)
energy,mat_v = eigen(mat_H)
histogram(energy,bins=100)

function calc_HGraphenekx(kx,Ny,μ)
    N = Ny*2
    mat_Htb = zeros(Complex{Float64},N,N)   
    mat_Htb += (-μ)*Matrix(I,N,N)#eye(N)
    t = 1.0
    for iy=1:Ny            
        for dy in -1:1              
            jy = iy + dy
            jy += ifelse(jy > Ny,-Ny,0)
            jy += ifelse(jy < 1,Ny,0)
                    
            for a=1:2
                b = ifelse(a ==1,2,1)
                ii = (iy-1)*2+a
                jj = (jy-1)*2+b

                if dy == 0 && a == 1
                    mat_Htb[ii,jj] = t+t*exp(im*kx)                        
                elseif dy == 1 && a ==1
                    mat_Htb[ii,jj] = t
                elseif dy ==0 && a ==2
                    mat_Htb[ii,jj] = t+t*exp(-im*kx)
                elseif dy == -1 && a ==2
                    mat_Htb[ii,jj] = t
                end
            end
                    

        end
    end
                    
        
    return mat_Htb
end

μ=0.0
Ny = 50
nkx = 100
vkx = range(-π,π,length = nkx)
ep = zeros(Float64,nkx,Ny*2)
cnt = 0
for kx in vkx    
    cnt += 1
    mat_H = calc_HGraphenekx(kx,Ny,μ)
    energy,mat_v = eigen(mat_H)
    for i=1:Ny*2
        #println(energy[i])
        ep[cnt,i] = energy[i]
    end
end
plot(vkx,ep)
plot!(legend=false)

function calc_HGraphenekxky(kx,ky,μ)
    N = 2
    mat_Htb = zeros(Complex{Float64},N,N)   
    mat_Htb += (-μ)*Matrix(I,N,N)#eye(N)
    t = 1.0
    c = sqrt(3)/2
#    for iy=1:Ny            
#        for dy in -1:1              
#            jy = iy + dy
#            jy += ifelse(jy > Ny,-Ny,0)
#            jy += ifelse(jy < 1,Ny,0)
                    
            for a=1:2
                b = ifelse(a ==1,2,1)
                ii = a
                jj = b

                if a == 1
                    mat_Htb[ii,jj] = t+t*exp(im*kx)+t*exp(im*ky+c)                        
#                elseif dy == 1 && a ==1
#                    mat_Htb[ii,jj] = t
                elseif a ==2
                    mat_Htb[ii,jj] = t+t*exp(-im*kx)+t*exp(-im*ky*c)
#                elseif dy == -1 && a ==2
#                    mat_Htb[ii,jj] = t
                end
            end
                    

#        end
#    end
                    
        
    return mat_Htb
end

 calc_HGraphenekxky(1.0,2.2,μ)


f(kx,ky) = sqrt(3 + 2*cos(kx)+2*cos(ky*sqrt(3)/2) + 2*cos(kx-ky*sqrt(3)/2))

x = -2π:2π/100:2π
y = -2π:2π/100:2π
z = [f(i,j) for i in x, j in y]'

plot(x,y,z)
plot!(aspect_ratio=:equal)

f(kx,ky) = sqrt(3 + 2*cos(kx+ky/2)+2*cos(ky) + 2*cos(kx+ky/2-ky))

x = -2π:2π/100:2π
y = -2π:2π/100:2π
z = [f(i,j) for i in x, j in y]'

plot(x,y,z)
plot!(aspect_ratio=:equal)

function calc_HGraphenekx_w(kx,Ny,μ)
    N = Ny*2
    mat_Htb = zeros(ComplexF64,N,N)   
    mat_Htb += (-μ)*Matrix(I,N,N)#eye(N)
    t = 1.0
    for iy=1:Ny            
        for dy in -1:1              
            jy = iy + dy
            #jy += ifelse(jy > Ny,-Ny,0)
            #jy += ifelse(jy < 1,Ny,0)
                    
            for a=1:2
                b = ifelse(a ==1,2,1)
                ii = (iy-1)*2+a
                jj = (jy-1)*2+b
                if 1 <= jy <= Ny
                    if dy == 0 && a == 1
                        mat_Htb[ii,jj] = t+t*exp(im*kx)                        
                    elseif dy == 1 && a ==1
                        mat_Htb[ii,jj] = t
                    elseif dy ==0 && a ==2
                        mat_Htb[ii,jj] = t+t*exp(-im*kx)
                    elseif dy == -1 && a ==2
                        mat_Htb[ii,jj] = t
                    end
                end
            end
                    

        end
    end
                    
        
    return mat_Htb
end

μ=0.0
Ny = 50
nkx = 100
vkx = range(-π,π,length=nkx)
ep = zeros(Float64,nkx,Ny*2)
cnt = 0
for kx in vkx    
    cnt += 1
    mat_H = calc_HGraphenekx_w(kx,Ny,μ)
    energy,mat_v = eigen(mat_H)
    for i=1:Ny*2
        #println(energy[i])
        ep[cnt,i] = energy[i]
    end
end
plot(vkx,ep)
plot!(legend=false)

kx = π-0.1
Ny = 50
mat_H = calc_HGraphenekx_w(kx,Ny,μ)
energy,mat_v = eigen(mat_H)
yv = []
ψ1 = []
ψ2 = []
for i in 1:Ny
    push!(yv,i)
    push!(ψ1,mat_v[2*i,Ny])
    push!(ψ2,mat_v[2*i-1,Ny+1])
end
plot(yv,[real.(ψ1),imag.(ψ1),real.(ψ2),imag.(ψ2)],label=["1:Real","1:Imag","2:Real","2:Imag"])