from loong import *
import json
import numpy as np
import sys
import matplotlib.pyplot as plt


class GoodOrbit(Orbit):

  def __init__(self):
    self.kAlpha = mp.mpf("1.7") / (2 * mp.pi)
    self.kCriticalTheta = 2.86 / ((2 / 3) * self.kAlpha)
    self.r = (1 / 3) * self.kAlpha * mp.sqrt(1 + self.kCriticalTheta**2)
    self.point_A_cartesian = (
        self.kAlpha * self.kCriticalTheta * mp.cos(self.kCriticalTheta),
        self.kAlpha * self.kCriticalTheta * mp.sin(self.kCriticalTheta),
    )
    self.point_B_cartesian = (-self.kAlpha * self.kCriticalTheta * mp.cos(self.kCriticalTheta),
                              -self.kAlpha * self.kCriticalTheta * mp.sin(self.kCriticalTheta))
    self.kPhi = mp.atan(self.kCriticalTheta)
    dx, dy = self.point_A_cartesian[0] - self.point_B_cartesian[0], self.point_A_cartesian[1] - self.point_B_cartesian[1]
    self.angle = mp.atan2(dy, dx)
    self.point_C1_cartesian = (self.point_A_cartesian[0] - 2 * self.r * dx, self.point_A_cartesian[1] - 2 * self.r * dy)
    self.point_C2_cartesian = (self.point_B_cartesian[0] + 1 * self.r * dx, self.point_B_cartesian[1] + 1 * self.r * dy)
    self.radius_of_C1 = 2 * self.r
    self.radius_of_C2 = 1 * self.r
    self.arclength = 6 * self.r * self.kPhi
    self.edge_k = self.kAlpha * mp.sqrt(1 + self.kCriticalTheta * self.kCriticalTheta)
    self.n = -1
    for i in range(3, 20, 2):
      self.a = (self.arclength - 2 * self.edge_k * self.kCriticalTheta) / (2 * (1 - i) * self.kCriticalTheta**i)
      self.b = (i * self.arclength - 2 * self.edge_k * self.kCriticalTheta) / (2 * (i - 1) * self.kCriticalTheta)
      if self.a > 0 and self.b > 0:
        self.n = i
        break
    print(f"arclength={self.arclength}", file=sys.stderr)
    print(f"edge_k={self.edge_k}", file=sys.stderr)
    print(f"a={self.a}", file=sys.stderr)
    print(f"b={self.b}", file=sys.stderr)
    print(f"n={self.n}", file=sys.stderr)
    print(f"now k={self.n*self.a*self.kCriticalTheta**(self.n-1)+self.b}", file=sys.stderr)
    if self.n == -1:
      raise Exception("n must be set")
    self.edge_raw_C = self.kAlpha * 0.5 * (
        self.kCriticalTheta * mp.sqrt(1 + self.kCriticalTheta * self.kCriticalTheta) -
        mp.log(-self.kCriticalTheta + mp.sqrt(1 + self.kCriticalTheta * self.kCriticalTheta)))

  def InitIdx(self):
    return mp.mpf("0.0")

  def InitC(self):
    return mp.mpf("0.0")

  def Idx2C(self, idx):  # this function must be monotonically increasing
    if idx >= 0:
      theta = idx + self.kCriticalTheta
      tmp = mp.sqrt(1 + theta * theta)
      return self.kAlpha * 0.5 * (theta * tmp - mp.log(-theta + tmp)) - self.edge_raw_C
    elif idx >= -2 * self.kCriticalTheta:
      x = idx + self.kCriticalTheta
      y = (self.a * (x**self.n) + self.b * x) - 0.5 * self.arclength
      return y
    else:
      theta = -idx - self.kCriticalTheta
      tmp = mp.sqrt(1 + theta * theta)
      return -self.kAlpha * 0.5 * (theta * tmp - mp.log(-theta + tmp)) + self.edge_raw_C - self.arclength

  def Idx2Cartesian(self, idx):
    return mp.matrix([mp.cos(idx), mp.sin(idx)])

  def C2Idx(self, C):

    def f(idx):
      return self.Idx2C(idx) - C

    return mp.findroot(f, 0, solver='secant')

  def GenerateNextPointIdx(self, cur_point_idx, expected_distance):
    return cur_point_idx + expected_distance


orbit = GoodOrbit()


def f(x):
  return float(orbit.Idx2C(x))


# 获取 kCriticalTheta
kCriticalTheta = float(orbit.kCriticalTheta)
print(f"kCriticalTheta={kCriticalTheta}")

# 定义范围 [-kCriticalTheta-1, kCriticalTheta +1]
start = -2.5 * kCriticalTheta - 1
end = 0.5 * kCriticalTheta + 1

# 生成范围内的点
x_vals = np.linspace(start, end, 1000)

# 计算 f(x) 的值
y_vals = [f(mp.mpf(x)) for x in x_vals]

# 绘制图像
plt.figure(figsize=(10, 6))
plt.plot(x_vals, y_vals, label='Idx2C(x)')
plt.title('Idx2C function plot')
plt.xlabel('Index (x)')
plt.ylabel('C Value')
plt.grid(True)
plt.axhline(0, color='black', linewidth=0.5)
plt.axvline(0, color='black', linewidth=0.5)
plt.legend()
plt.show()
sys.exit()
if __name__ == "__main__":
  orbit = GoodOrbit()
  loong = Loong(orbit, 224, mp.mpf("2.0"), mp.mpf("1e-8"))
  res_list = []
  for ti in range(-100, 101):
    print(f"calculating time_point={ti}")
    res_list.append(loong.CalcStatusListByTime(mp.mpf(ti)))
    # 转换成内置浮点数并保留6位
  float_res_list = [[{
      "idx": round(float(node["idx"]), 6),
      "node": [
          round(float(node["node"][0]), 6),
          round(float(node["node"][1]), 6),
      ],
      "C": round(float(node["C"]), 6),
      "v": round(float(node["v"]), 6),
  } for node in res] for res in res_list]
  with open("A4_res.json", "w") as file:
    json.dump(float_res_list, file, indent=4)