<gh_stars>0
from operator import ge
from KMC.AFP.AFP_lattice import AFP_square_lattice
import numpy as np

def generate_data():
    a   = 0.1
    Lx  = 10
    Ly  = 10
    L   = 2
    W   = 3

    sl  = AFP_square_lattice(a, Lx, Ly)

    angle60 = sl.square_offset(L, W,60 )
    angle120 = sl.square_offset(L, W, 120)
    angle0  = sl.square_offset(L, W, 0)

    np.save("ref_data/Angle60.npy", angle60)
    np.save("ref_data/Angle120.npy", angle120)
    np.save("ref_data/Angle0.npy", angle0)


def check_data():
    a   = 0.1
    Lx  = 10
    Ly  = 10
    L   = 2
    W   = 3

    sl  = AFP_square_lattice(a,Lx, Ly)

    angle60 = sl.square_offset(L, W,60 )
    angle120 = sl.square_offset(L, W, 120)
    angle0  = sl.square_offset(L, W, 0)

    angle60_check = np.load("ref_data/Angle60.npy")
    angle120_check = np.load("ref_data/Angle120.npy")
    angle0_check = np.load("ref_data/Angle0.npy")

    assert(np.array_equal(angle60,angle60_check))
    assert(np.array_equal(angle120, angle120_check))
    assert(np.array_equal(angle0, angle0_check))

if __name__ == "__main__":
    #generate_data()
    check_data()
    print("Passed")
