import sys
import numpy as np
import pandas as pd # pandasを追加
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:
    # 拡張子によって読み込み方法を分岐
    if infile.endswith('.xlsx'):
        try:
            df = pd.read_excel(infile)
            t = df.iloc[:, 0].values
            x = df.iloc[:, 1].values
        except Exception as e:
            print(f"Error reading Excel file: {e}")
            sys.exit(1)
    elif infile.endswith('.csv'):
        try:
            _a = np.loadtxt(infile, delimiter=',', skiprows=1, usecols=[0,1])
            t, x = _a.T
        except Exception as e:
            print(f"Error reading CSV file: {e}")
            sys.exit(1)
    elif infile.endswith('.txt'):
        try:
            _a = np.loadtxt(infile, skiprows=1, usecols=[0,1])
            t, x = _a.T
        except Exception as e:
            print(f"Error reading TXT file: {e}")
            sys.exit(1)
    else:
        print("Unsupported file format. Please use .xlsx, .csv, or .txt.")
        sys.exit(1)

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()


# 結果をExcelファイルに出力
if infile != "":
    # 入力ファイル名から出力ファイル名を生成
    filebody = infile.rsplit('.', 1)[0]
    output_filename = f"{filebody}-mem_order={order}.xlsx"

    df_output = pd.DataFrame({
        'Frequency_MEM': freqs_mem,
        'PSD_MEM': mem_psd,
        'Frequency_FFT': freqs_fft,
        'PSD_FFT': fft_psd
    })
    try:
        df_output.to_excel(output_filename, index=False)
        print(f"\nMEM and FFT results saved to {output_filename}")
    except Exception as e:
        print(f"Error saving results to Excel file: {e}")


# グラフ表示: subplotsを使って1つのウィンドウに縦に並べる
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) # 2行1列のサブプロットを作成

# 1つ目のサブプロット（入力信号）
axes[0].plot(t, x, label='input', linewidth=2)
axes[0].set_xlabel('t')
axes[0].set_ylabel('signal')
axes[0].set_title('Input signal')
#axes[0].legend()
axes[0].grid(True) # グリッドを追加しました

# 2つ目のサブプロット（スペクトル比較）
axes[1].plot(freqs_mem, mem_psd, label='MEM (pburg)', linewidth=2)
axes[1].plot(freqs_fft, fft_psd, '--', label='FFT PSD', linewidth=1)
if yscale == 'log':
    axes[1].set_yscale('log')
axes[1].set_xlabel('Frequency (Hz)')
axes[1].set_ylabel('PSD (dB)')
axes[1].set_title('Spectral Estimate: MEM (order={order}) vs FFT')
axes[1].legend()
axes[1].grid(True) # グリッドを追加しました

plt.tight_layout() # サブプロット間のスペースを自動調整
plt.pause(0.1)


input("\nPress ENTER to terminate>>\n")