import json
from obspy.taup import TauPyModel
import numpy as np
from tqdm import tqdm

model = TauPyModel(model='iasp91')

computed_depths = 70
computed_angles = 180
max_angle = 90 # deg
max_depth = 700 # km

depth_list = np.linspace(0,max_depth,computed_depths)
angle_list = np.linspace(0,max_angle,computed_angles)

p_table = []
for depth in tqdm(range(len(depth_list))):
    arrivals = []
    for angle in angle_list:
        arrival = model.get_travel_times(source_depth_in_km=depth_list[depth],
                                    distance_in_degree=angle,phase_list=["p","P"])[0].time
        arrivals.append(arrival)
    p_table.append(arrivals)

s_table = []
for depth in tqdm(range(len(depth_list))):
    arrivals = []
    for angle in angle_list:
        arrival = model.get_travel_times(source_depth_in_km=depth_list[depth],
                                    distance_in_degree=angle,phase_list=["s","S"])[0].time
        arrivals.append(arrival)
    s_table.append(arrivals)

data = {"angles": list(angle_list), "depths":list(depth_list), "p_table":p_table,"s_table":s_table}

print(data)

f = open("computed_tables.json", "w")
f.write(json.dumps(data))
f.close()