import z2pack
import numpy as np
import math as c

def hamiltonian(k):
    kx,ky,kz=k
    #t=10,mu=10,r=5,A0=A1=2
    return np.array([
        [-2*10*(c.cos(kx)+c.cos(ky))-10,-2*5*c.sin(ky)-2*1j*5*c.sin(kx),2+2*2*(c.cos(kx)+c.cos(ky)),0],
        [-2*5*c.sin(ky)+2*1j*5*c.sin(kx),-2*10*(c.cos(kx)+c.cos(ky))-10,0,2+2*2*(c.cos(kx)+c.cos(ky))],
        [2+2*2*(c.cos(kx)+c.cos(ky)),0,2*10*(c.cos(kx)+c.cos(ky))+10,2*5*c.sin(ky)-2*1j*5*c.sin(kx)],
        [0,2+2*2*(c.cos(kx)+c.cos(ky)),2*5*c.sin(ky)+2*1j*5*c.sin(kx),2*10*(c.cos(kx)+c.cos(ky))+10]
    ])
system = z2pack.hm.System(hamiltonian,bands=4)

result=z2pack.surface.run(
    system=system,
    surface=lambda t1,t2:[t1,t2,0],
    num_lines=100,
    min_neighbour_dist=0.0001
)

print(z2pack.invariant.chern(result))
print(z2pack.invariant.z2(result))