import numpy as np
import matplotlib.pyplot as plt

N = 512            # サンプル数（デジタル化する際の点の数）
dt = 0.01          # サンプリング間隔（0.01秒ごとに点を打つ）
t = np.arange(0, N * dt, dt) # 時間軸の配列（0から5.12秒まで）

y = np.sin(2 * np.pi * 3 * t) + np.sin(2 * np.pi * 5 * t)

# コンピュータで無駄なく計算し、かつ元通りに復元できるようにするために
# サンプル数と全く同じ個数・同じ間隔のスピードだけを調べる
# 円の重心の位置っぽいもの (複素数) の配列を返す
def my_dft(x):
    N_data = len(x)
    X = []
    # k は探りたい周波数のインデックス（円に巻くスピード）
    for k in range(N_data):
        X_k = 0j # 複素数のゼロで初期化
        
        # n は各データポイント（ビーズ）
        for n in range(N_data):
            # 角度θ = -2π * k * n / N
            # cmath.exp(1j * θ) がオイラーの公式による「円への巻き付け」
            theta = -2 * np.pi * k * n / N_data
            X_k += x[n] * np.exp(1j * theta)
            
        X.append(X_k) # 総和（重心ベクトル）を保存
    return X

def my_idft(x):
    N_data = len(x)
    X = []
    for k in range(N_data):
        x_k = 0j

        for n in range(N_data):
            theta = 2 * np.pi * k * n / N_data
            x_k += x[n] * np.exp(1j * theta)

        X.append(x_k / N_data)
    return X

F = my_dft(y)
F_i = my_idft(F)

freqs = []
amps = []

for k in range(N):
    # ナイキスト周波数（N/2）を境に、後半はマイナスの周波数として扱う
    if k < N / 2:
        freq = k / (N * dt)
    else:
        freq = (k - N) / (N * dt)
    freqs.append(freq) # これは plot のために必要なだけ
    
    # 複素数の絶対値（ベクトルの長さ）を求め、N/2で割って振幅を正規化
    amps.append(abs(F[k]) / (N / 2))

plt.figure(figsize=(10, 6))

plt.subplot(3, 1, 1)
plt.plot(t, y)
plt.xlim(0, 2)
plt.title("Original Wave: sin(2 * pi * 3 * t) + sin(2 * pi * 5 * t)")
plt.xlabel("Time [s]")
plt.ylabel("Amplitude")
plt.grid(True)

plt.subplot(3, 1, 2)

plt.plot(freqs, amps, marker='o', linestyle='-')
plt.xlim(-8, 8)
plt.title("DFT Result")
plt.xlabel("Frequency [Hz]")
plt.ylabel("Amplitude")
plt.grid(True)

plt.subplot(3, 1, 3)
plt.plot(t, [val.real for val in F_i]) # 実部だけ取り出して plot
plt.xlim(0, 2)
plt.title("IDFT Result")
plt.xlabel("Time [s]")
plt.ylabel("Amplitude")
plt.grid(True)

plt.tight_layout()
plt.show()