import sys
import time
import numpy as np
import mwa_hyperbeam
if len(sys.argv) > 1 and sys.argv[1] == "rts":
beam = mwa_hyperbeam.AnalyticBeam(rts_behaviour=True)
else:
beam = mwa_hyperbeam.AnalyticBeam()
n = 1000000
az = np.linspace(0, 0.9 * np.pi, n)
za = np.linspace(0.1, 0.9 * np.pi / 2, n)
freq = [167000000]
delays = np.zeros((2, 16), dtype=np.uint).flatten()
amps = np.ones((2, 16)).flatten()
norm_to_zenith = True
latitude_rad = -0.4660608448386394
start_time = time.time()
jones = beam.calc_jones_gpu(az, za, freq, delays, amps, latitude_rad, norm_to_zenith)
duration = time.time() - start_time
print("Time to calculate {} directions: {:.3}s".format(n, duration))
print("First Jones matrix:")
print(jones[0, 0, 0])
print(f"dtype of the array: {jones.dtype}")