import mpmath as mp
import numba as nb

mp.dps = 50  # 设置精度为50位小数

class Orbit:
  def __init__(self):
    self.kAlpha = mp.mpf('1.7') / (2 * mp.pi)
  def InitIdx(self) -> mp.mpf:
    raise NotImplementedError
  def InitC(self) -> mp.mpf:
    raise NotImplementedError
  def Idx2C(self, idx) -> mp.mpf:
    raise NotImplementedError
  def Idx2Cartesian(self, idx) -> [mp.mpf, mp.mpf]:
    raise NotImplementedError
  def C2Idx(self, C:[mp.mpf, mp.mpf]) -> mp.mpf:
    raise NotImplementedError
  def GenerateNextPointIdx(self, cur_point_idx:mp.mpf, expected_distance:mp.mpf)->mp.mpf:
    raise NotImplementedError

class Loong:
  def __init__(self, orbit:Orbit, total_points:int, speed:mp.mpf, delta_idx:mp.mpf):
    self.orbit = orbit
    self.total_points = total_points
    self.speed = speed
    self.delta_idx = delta_idx
    self.kSegLength1 = mp.mpf('2.86')
    self.kSegLength2 = mp.mpf('1.65')
  def CalcStatusListByIdx(self, cur_idx:mp.mpf):
    first_node_idx=cur_idx
    first_node_C=self.orbit.Idx2C(first_node_idx)
    first_node_dot = self.orbit.Idx2Cartesian(first_node_idx)
    virtual_first_node_idx = first_node_idx + self.delta_idx
    virtual_first_node_C = self.orbit.Idx2C(virtual_first_node_idx)
    delta_T = (virtual_first_node_C - first_node_C) / self.speed
    node_list = [{"idx": first_node_idx, "node": first_node_dot, "C": first_node_C, "v": mp.mpf('1.0')}]
    for i in range(1, self.total_points):
      expected_distance = self.kSegLength1 if i == 1 else self.kSegLength2
      cur_node_idx = self.orbit.GenerateNextPointIdx(node_list[-1]["idx"], expected_distance)
      cur_node_dot = self.orbit.Idx2Cartesian(cur_node_idx)
      cur_node_C = self.orbit.Idx2C(cur_node_idx)

      virtual_cur_node_idx = self.orbit.GenerateNextPointIdx(virtual_first_node_idx, expected_distance)
      virtual_cur_node_C = self.orbit.Idx2C(virtual_cur_node_idx)
      v = (virtual_cur_node_C - cur_node_C) / delta_T

      node_list.append({"idx": cur_node_idx, "node": cur_node_dot, "C": cur_node_C, "v": v})
      virtual_first_node_idx = virtual_cur_node_idx
    return node_list
  def CalcStatusListByTime(self, time_point:mp.mpf):
    first_node_C = self.orbit.InitC() - time_point * self.speed
    first_node_idx = self.orbit.C2Idx(first_node_C)
    return self.CalcStatusListByIdx(first_node_idx)