import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

#微分方程式を定義、単振り子に周期的な外力が働く非減衰保存系、x=x[0]は角度、y=x[1]は角速度、井上・秦 1999, p. 118.参考
def forced_pend(x: np.ndarray, t:float) -> np.ndarray:
    dx = x[1]
    dy = -np.sin(x[0]) + B * np.sin(Omega*t)
    return(np.array([dx, dy]))

# パラメータ
B = 0.1
Omega = 1

divnum = 100                # 計算時間刻み幅を2piの何等分にするか
dt = 2*np.pi/Omega/divnum   # 計算時間刻み幅
t0 = 0                      # 初期時刻
totalnum = 200000           # 計算繰り返し数
T = dt * totalnum           # 計算総時間
times = np.arange(t0, T, dt)

xin = 0                     # xの初期値の初期値
yin = -2.4                  # yの初期値の初期値

stroboID=np.arange(0, totalnum, divnum) # ストロボ写像で抜き出す分のリスト番号（t = 2pi x 自然数 に相当するリスト番号））を定義

# グラフの用意
fig, ax = plt.subplots()
xlabel = 'x'
ylabel = 'y'
ax.set_xlabel(xlabel,fontsize=14, fontstyle="italic")
ax.set_ylabel(ylabel,fontsize=14, fontstyle="italic")

# プロットカラーのリストの用意
color_list = ["red","blue","orange","green", "pink", "gray", "purple", "teal", "brown", "black", "gold"]    
colornum = 0 

while yin <= 3:  
    # x,yの初期値    
    x0 = [xin, yin]

    # 微分方程式を解き、ストロボ写像の点を抽出
    orbit = odeint(forced_pend, x0, times)
    strobo_orbit = orbit[stroboID]

    # xをθ（[-pi, pi]のあいだの数値）に変換
    theta_int = np.trunc(strobo_orbit[:,0] /np.pi)
    numb=len(theta_int)
    twopi_num = []
    for i in range(numb):
        if theta_int[i] % 2 == 0:
            ev=theta_int[i]/2
            twopi_num.append(ev)
        else:
            if theta_int[i] >= 0:
                op = (theta_int[i]+1)/2
                twopi_num.append(op)
            else:
                on=(theta_int[i]-1)/2
                twopi_num.append(on)
    theta_strobo_pre = strobo_orbit[:,0]/ (2*np.pi) - twopi_num
    theta_strobo_orbit = theta_strobo_pre * 2*np.pi

    # 1つの軌道のストロボ写像の点をプロット
    ax.plot(theta_strobo_orbit, strobo_orbit[:,1], '.', markersize=1, color=color_list[colornum])

    yin += 0.5
    colornum +=1

plt.show()