import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def solve_kepler(M, e, tol=1e-10):
    """ケプラー方程式をニュートン法で解く"""
    E = M  # 初期値
    for _ in range(100):
        dE = (E - e * np.sin(E) - M) / (1 - e * np.cos(E))
        E -= dE
        if np.all(np.abs(dE) < tol):
            break
    return E

def orbital_elements_to_position(a, e, i, omega_big, omega_small, M, mu=3.986e14):
    """軌道要素からECI座標系での位置を計算"""
    # ケプラー方程式を解く
    E = solve_kepler(M, e)

    # 真近点角
    nu = 2 * np.arctan2(np.sqrt(1+e) * np.sin(E/2),
                         np.sqrt(1-e) * np.cos(E/2))

    # 軌道面内の距離
    r = a * (1 - e * np.cos(E))

    # 軌道面内の座標
    x_orb = r * np.cos(nu)
    y_orb = r * np.sin(nu)

    # ECI座標系への変換（回転行列）
    cos_O = np.cos(omega_big)
    sin_O = np.sin(omega_big)
    cos_w = np.cos(omega_small)
    sin_w = np.sin(omega_small)
    cos_i = np.cos(i)
    sin_i = np.sin(i)

    x = (cos_O*cos_w - sin_O*sin_w*cos_i)*x_orb + \
        (-cos_O*sin_w - sin_O*cos_w*cos_i)*y_orb
    y = (sin_O*cos_w + cos_O*sin_w*cos_i)*x_orb + \
        (-sin_O*sin_w + cos_O*cos_w*cos_i)*y_orb
    z = (sin_w*sin_i)*x_orb + (cos_w*sin_i)*y_orb

    return x, y, z

# 軌道パラメータの例（ISSに近い軌道）
Re = 6378.137e3  # 地球半径 [m]
a = Re + 408e3   # 軌道長半径（高度408km）
e = 0.0007       # 離心率（ほぼ円軌道）
i = np.radians(51.6)     # 軌道傾斜角
omega_big = np.radians(30)  # 昇交点赤経
omega_small = np.radians(60)  # 近地点引数

# 1周回分の軌道を計算
M_array = np.linspace(0, 2*np.pi, 1000)
x_orb, y_orb, z_orb = orbital_elements_to_position(
    a, e, i, omega_big, omega_small, M_array)

# 可視化
fig = plt.figure(figsize=(14, 6))

# 3D軌道
ax1 = fig.add_subplot(121, projection='3d')

# 地球（球）
u = np.linspace(0, 2 * np.pi, 50)
v = np.linspace(0, np.pi, 50)
x_earth = Re * np.outer(np.cos(u), np.sin(v))
y_earth = Re * np.outer(np.sin(u), np.sin(v))
z_earth = Re * np.outer(np.ones_like(u), np.cos(v))
ax1.plot_surface(x_earth, y_earth, z_earth, alpha=0.3, color='blue')

# 軌道
ax1.plot(x_orb, y_orb, z_orb, 'r-', linewidth=2, label='Orbit')
ax1.scatter(x_orb[0], y_orb[0], z_orb[0], c='green', s=100, label='Start')

ax1.set_xlabel('X [m]')
ax1.set_ylabel('Y [m]')
ax1.set_zlabel('Z [m]')
ax1.set_title('Satellite Orbit in ECI Frame')
ax1.legend()

# 離心率と軌道形状の関係
ax2 = fig.add_subplot(122)
nu = np.linspace(0, 2*np.pi, 500)

for e_val in [0.0, 0.2, 0.5, 0.7, 0.9]:
    r = a * (1 - e_val**2) / (1 + e_val * np.cos(nu))
    x_plot = r * np.cos(nu) / Re
    y_plot = r * np.sin(nu) / Re
    ax2.plot(x_plot, y_plot, label=f'e = {e_val}')

ax2.scatter(0, 0, c='blue', s=100, zorder=5, label='Earth')
ax2.set_xlabel('x [Earth Radii]')
ax2.set_ylabel('y [Earth Radii]')
ax2.set_title('Orbit Shape vs Eccentricity')
ax2.legend()
ax2.set_aspect('equal')
ax2.grid(True)

plt.tight_layout()
plt.show()

# 軌道パラメータの表示
mu = 3.986e14
T = 2 * np.pi * np.sqrt(a**3 / mu)
print(f"=== 軌道パラメータ ===")
print(f"軌道長半径 a = {a/1e3:.1f} km")
print(f"離心率 e = {e}")
print(f"軌道傾斜角 i = {np.degrees(i):.1f} deg")
print(f"軌道周期 T = {T/60:.1f} min")
print(f"近地点高度 = {a*(1-e)/1e3 - Re/1e3:.1f} km")
print(f"遠地点高度 = {a*(1+e)/1e3 - Re/1e3:.1f} km")