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(10, 20, 1/1000):
  tasks_list.append(i)
for i in np.arange(12, 14.2, 1/10000):
  tasks_list.append(i)
for i in np.arange(13.085,13.095,1/1000000):
  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")
  max_speed_time=0
  last_status = None
  last_ti = ti_list[0]
  for ti in ti_list:
    print(f"calculating time_point={ti}",file=sys.stderr)
    try:
      res = loong.CalcStatusListByTime(mp.mpf(ti), last_status if (last_status and abs(ti-last_ti)<=0.1) else None)
      for node in res:
        if node["v"] > max_speed_found:
          max_speed_found = node["v"]
          max_speed_time = ti
      last_status = res
      last_ti = ti
    except ValueError as e:
      print(f"Error: {e}",file=sys.stdout)
  return max_speed_found, max_speed_time

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=mp.mpf("0.0")
  max_speed_time=0
  for res in res_list:
    if res[0]>max_speed_found:
      max_speed_found=res[0]
      max_speed_time=res[1]
  valid_head_speed = mp.mpf("1.0") * (mp.mpf("2.0")/max_speed_found)
  print(f"max_speed_found={max_speed_found} at {max_speed_time}, valid_head_speed={valid_head_speed}")