import sys
import numpy as np
import matplotlib.pyplot as plt
from spectrum import pburg, Periodogram


#infile = 'input_time_series.csv'
infile = ''
ndata = 300
yscale = 'log'
order  = 8  # ARモデルの次数

argv = sys.argv
nargs = len(argv)
if nargs > 1: infile = argv[1]
if nargs > 2: order = int(argv[2])
if nargs > 3: yscale = argv[3]


def generate_input(nmax: int):
    if nmax < 100:
        print("Warning: nmax < 100. Setting nmax = 100.")
        nmax = 100

    dt = 1.0 / (nmax - 1)
    t = np.linspace(0, 1.0, nmax)

#    x = np.cos(2 * np.pi * t * 1.7)

    w1 = 124.5
    w2 = 62.0
    w3 = 67.0
    w4 = 3.0
    x = (np.sin(2 * np.pi * t * w1) +
         np.sin(2 * np.pi * t * w2) +
         np.sin(2 * np.pi * t * w3) +
         np.sin(2 * np.pi * t * w4))


    return t, x


print()
print(f"{infile=}")
print(f"{order=}")
print(f"{yscale=}")

if infile == "":
    t, x = generate_input(ndata)
else:
    _a = np.loadtxt(infile, delimiter=',', skiprows=1, usecols=[0,1])
    t, x = _a.T

nfft = len(x)

# MEMスペクトル（Burg法）
burg_spec = pburg(x, order=order, NFFT=nfft)
mem_psd   = burg_spec.psd                      # MEM PSD
freqs_mem = burg_spec.frequencies()            # 周波数軸

# FFTベースPSD
fft_spec  = Periodogram(x, NFFT=nfft)
fft_psd   = fft_spec.psd
freqs_fft = fft_spec.frequencies()


# プロット(入力)
plt.figure(figsize=(8, 4))
plt.plot(t, x, label='input', linewidth=2)
plt.xlabel('t')
plt.ylabel('signal')
plt.title('Input signal')
#plt.legend()
#plt.grid(True)
plt.pause(0.1)

# プロット比較
plt.figure(figsize=(8, 4))
plt.plot(freqs_mem, mem_psd, label='MEM (pburg)', linewidth=2)
plt.plot(freqs_fft, fft_psd, '--',      label='FFT PSD', linewidth=1)
#plt.plot(freqs_mem, np.log10(mem_psd), label='MEM (pburg)', linewidth=2)
#plt.plot(freqs_fft, np.log10(fft_psd), '--',      label='FFT PSD', linewidth=1)
if yscale == 'log':
    plt.yscale('log')
plt.xlabel('Frequency (Hz)')
plt.ylabel('PSD (dB)')
plt.title('Spectral Estimate: MEM vs FFT')
plt.legend()
#plt.grid(True)
plt.show()

