import numpy as np
from matplotlib import pyplot as plt
from matplotlib import patches
import japanize_matplotlib
import warnings
import sys
warnings.simplefilter('ignore')
# 3体問題-1
#           2025.03.25
#
#  global 変数
crk0=0.292893218813452
crk1=2-crk0
r1=np.zeros(4)  # ルンゲクッタ計算用配列
r2=np.zeros(4)
r3=np.zeros(4)
r4=np.zeros(4)
q1=np.zeros(4)
q2=np.zeros(4)
q3=np.zeros(4)
q4=np.zeros(4)
yw=np.zeros(4)
y0=np.zeros(4)
y1=np.zeros(4)
y2=np.zeros(4)
y3=np.zeros(4)
y4=np.zeros(8)

a=np.zeros(8)
prm=np.zeros(16)

x1=0
x2=0
mu=0  
#-----------------------------------------------------
def newton(x,a):    # 5次方程式をニュートン法で解く
    x3 = x  # 初期値
    for i in range(20):
        yy1 = a[5]*x3**5 + a[4]*x3**4 + a[3]*x3**3 + a[2]*x3**2 + a[1]*x3 + a[0]
        yy2 = 5.*a[5]*x3**4 + 4.*a[4]*x3**3 + 3.*a[3]*x3**2 + 2.*a[2]*x3 + a[1]
        if abs(yy2) < 1.e-4:
            return 0
        dx = yy1/yy2
        if abs(dx) < 1.e-4:    # 差分が一定値以下なら
            return x3
        x3 = x3 -dx
    return x3
#----------------------------------------------------------
def ufunc(k,yy) -> np.float64: # 計算処理
    x=yy[0]
    y=yy[1]
    dx=yy[2]
    dy=yy[3]

    if k == 0:
        return dx

    elif k == 1:
        return dy

    elif k == 2:
        rr1=np.sqrt((x-x1)*(x-x1)+y*y)
        rr2=np.sqrt((x-x2)*(x-x2)+y*y)
        fx=-mu*(x-x1)/rr1**3-(1-mu)*(x-x2)/rr2**3
        ddx=dy*2.+x+fx
        return ddx

    elif k == 3:
        rr1=np.sqrt((x-x1)*(x-x1)+y*y)
        rr2=np.sqrt((x-x2)*(x-x2)+y*y)
        fy=-mu*y/rr1**3-(1.-mu)*y/rr2**3
        ddy=-dx*2+y+fy
        return ddy
#----------------------------------------------------------
def drkgn(y0,y4,dt,nn):  # ルンゲクッタ・ジル法

    for k in range(nn):
        q4[k]=0
        yw[k]=y4[k]=y0[k]

    for i in range(2):

        for k in range(nn):
            # y1, q1
            r1[k]=ufunc(k,yw)*dt
            ww=r1[k]*0.5-q4[k]
            y1[k]=yw[k]+ww
            q1[k]=q4[k]+ww*3-r1[k]*0.5

        for k in range(nn):
            # y2, q2
            r2[k]=ufunc(k,y1)*dt
            ww=(r2[k]-q1[k])*crk0
            y2[k]=y1[k]+ww
            q2[k]=q1[k]+ww*3.-r2[k]*crk0

        for k in range(nn):
            # y3, q3
            r3[k]=ufunc(k,y2)*dt
            ww=(r3[k]-q2[k])*crk1
            y3[k]=y2[k]+ww
            q3[k]=q2[k]+ww*3.-r3[k]*crk1

        for k in range(nn):
            # y4, q4
            r4[k]=ufunc(k,y3)*dt
            ww=(r4[k]-q3[k]*2.)/6.
            yw[k]=y3[k]+ww
            y4[k+i*nn]=yw[k]
            q4[k]=q3[k]+ww*3.-r4[k]*0.5

    return
#----------------------------------------------------------
def parm(yyy,xxx):  # パラメータ計算
    x=yyy[0]
    y=yyy[1]
    dx=yyy[2]
    dy=yyy[3]

    rr1=np.sqrt((x-x1)*(x-x1)+y*y)
    rr2=np.sqrt((x-x2)*(x-x2)+y*y)
    u=(x*x+y*y)*.5+mu/rr1+(1.-mu)/rr2
    c=u*2.-(dx*dx+dy*dy)
    h=(x-x1)*dy-y*dx+rr1*rr1
    hkp=np.sqrt(mu*rr1)

    xxx[0]=rr1
    xxx[1]=rr2
    xxx[2]=u
    xxx[3]=c
    xxx[4]=h
    xxx[5]=hkp

    return

#----------------------------------------------------------------
aa=np.zeros(8)  # ラグランジュ点計算用係数配列
dt=0.0025       # 時間刻み幅
kk=0
# --- 質量比入力
q=float(input('* 質量比( M2/M1) : '))
if q < 1.e-2 or q > 0.95:
    sys.exit()  # 入力エラーで終了

# --- 質量比入力
n=int(input('* 繰り返し回数 : '))
ix=np.zeros(n)
iy=np.zeros(n)

# ---- 重心=原点とする
mu=1./(1.+q)  # M1の質量割合 = 原点からM2までの距離
x1=mu-1.
x2=mu       # M2の質量割合 = 原点からM1までの距離

# ---- ラグランジュ点の計算 (M2からの距離/軌道長)
aa[0]=-q
aa[1]=2*q
aa[2]=-q
aa[3]=3+q
aa[4]=-3-2*q
aa[5]=1+q

x3=0.5        # L1 の初期値
xl0 = newton(x3,aa)
xl1 = x2-xl0  # L1のX座標

print('\n ---> 質量比：','{:.3f}'.format(q),'   (M2/M1)    ---- 繰返し回数 :',n)
print(' ---> M1 の座標：','{:.3f}'.format(x1),' ----  M2 の座標：','{:.3f}'.format(x2))
print(' ---> L1 の座標：','{:.3f}'.format(xl1),'  particle の出発点    計算刻み幅:','{:.3e}'.format(dt))
#
# ラグランジュ点の別な計算
z1=x1+0.0001
z2=x2-0.0001
zy=z1-mu/((z1-x1)*(z1-x1))+(1.-mu)/((z1-x2)*(z1-x2))
while abs(z1-z2)>1.e-5:
    zm=(z1+z2)/2.
    zq=zy*(zm-mu/((zm-x1)*(zm-x1))+(1.-mu)/((zm-x2)*(zm-x2)))
    if zq >= 0.:
       z1=zm
    else:
       z2=zm
xl2=zm
print(' ---> L1の別解：','{:.3f}'.format(xl2))
# ---- 初期値設定 & パラメータ計算
y0[0]=xl1    # x
y0[1]=0         # y
y0[2]=-0.2      # dx=-0.1
y0[3]=0         # dy = 0
t0=0.           # 時間初期値


# ----- 初期値
parm(y0,prm)
rr1=prm[0]
rr2=prm[1]
u=prm[2]
c0=prm[3]
h=prm[4]
hkp=prm[5]
print('* init. x:','{:.3f}'.format(y0[0]),' y:','{:.3f}'.format(y0[1]),' r1:','{:.3f}'.format(rr1),' r2:','{:.3f}'.format(rr2),' h:','{:.3f}'.format(h),' u:','{:.3e}'.format(u),' c0:','{:.3e}'.format(c0),' hkp:','{:.3e}'.format(hkp))

for j in range(n):

    drkgn(y0,y4,dt,4)    # ルンゲクッタ・ジル法

    for k in range(4):
        y0[k]=y4[k+4]

    parm(y0,prm)

    ix[j]=x=y0[0]
    iy[j]=y=y0[1]
    dx=y0[2]
    dy=y0[3]
    rr1=prm[0]
    rr2=prm[1]
    u=prm[2]
    c=prm[3]
    h=prm[4]     # 座標刻み幅
    hkp=prm[5]

#    if j%10==0:
#        print(' loop :',j,'  x:','{:.2f}'.format(x),' y:','{:.2f}'.format(y),' dx :','{:.2f}'.format(dx),' dy:','{:.2f}'.format(dy),' r1:','{:.2f}'.format(rr1),' r2:','{:.2f}'.format(rr2),' c :','{:.2e}'.format(c),' h:','{:.2e}'.format(h),' hkp:','{:.2e}'.format(hkp))

    if h>=hkp:
        kk=kk+1
#    if abs(c-c0)> 1.e-4:
      
print('\n* loop : ',j+1)
title="q="+str(q)

fig=plt.figure(figsize=(12,12))
fig, ax = plt.subplots()
ax.text(0.2,0.45,title,fontsize="xx-large")
ax.set_xlim(-1.0,1.0)
ax.set_ylim(-0.8,0.8)
c0 = patches.Circle( (xl1,0), 0.005, facecolor="green", edgecolor="red")
ax.add_patch(c0)
ax.text(xl1-0.05,0.05,"L1")
c1 = patches.Circle( (x1,0), 0.01, facecolor="pink", edgecolor="red")
ax.add_patch(c1)
c2 = patches.Circle( (x2,0), 0.05, facecolor="yellow", edgecolor="black")
ax.add_patch(c2)
plt.plot(ix,iy,color="blue",lw=0.7)
#for i in range(n-1):
#    plt.plot([ix[i],ix[i+1]], [iy[i],iy[i+1]], color='blue',lw=1)
plt.show()
fig.savefig('three-body.png')