from simulator import *

kTiBeg = -100
kTiEnd = 100
kSampleNum = 1000
kStep = (kTiEnd - kTiBeg) / kSampleNum
kParallelNum=24
tasks_list = [i for i in np.arange(kTiBeg, kTiEnd, kStep)]
for i in np.arange(0, 10, 0.1):
  tasks_list.append(i)
for i in np.arange(5, 10, 5.0/1000):
  tasks_list.append(i)
task_list_per_process=[tasks_list[i::kParallelNum] for i in range(kParallelNum)]

print(f"len(task_list_per_thread)={len(task_list_per_process)}",file=sys.stderr)
def ProcessEntryPoint(arg):
  ti_list, process_id = arg
  orbit = GoodOrbit()
  loong = Loong(orbit, 224, mp.mpf("1.0"), mp.mpf("1e-8"))
  max_speed_found=mp.mpf("0.0")
  for ti in ti_list:
    print(f"calculating time_point={ti}",file=sys.stderr)
    res = loong.CalcStatusListByTime(mp.mpf(ti))
    for node in res:
      max_speed_found=max(max_speed_found,node["v"])
  return max_speed_found

if __name__ == "__main__":
  manager = multiprocessing.Manager()
  task_args_list = [(task_list_per_process[i], i) for i in range(kParallelNum)]
  with multiprocessing.Pool(processes=kParallelNum) as pool:
    res_list=pool.map(ProcessEntryPoint, task_args_list)
  max_speed_found=max(res_list)
  valid_head_speed = mp.mpf("1.0") * (mp.mpf("2.0")/max_speed_found)
  print(f"max_speed_found={max_speed_found}, valid_head_speed={valid_head_speed}")