#!/usr/bin/env python3
import sys
import numpy as np
import pandas as pd # pandasを追加
import matplotlib.pyplot as plt
from spectrum import pburg, Periodogram
from statsmodels.tsa.stattools import pacf
from statsmodels.tsa.ar_model import AutoReg
import warnings
warnings.filterwarnings("ignore")

# ------------------------
# Parameters
# ------------------------
infile = 'input_time_series.csv'
preset_order = 8  # default AR model order if no override
argv = sys.argv
if len(argv) > 1:
    infile = argv[1]
if len(argv) > 2:
    preset_order = int(argv[2])

# ------------------------
# Load data
# ------------------------
if infile == "":
    print("No input file specified. Exiting.")
    sys.exit(1) # 入力ファイルがない場合は終了するように修正
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:
            # skiprows=1はヘッダーがある場合を想定。ない場合は削除してください。
            _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:
            # skiprows=1はヘッダーがある場合を想定。ない場合は削除してください。
            _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)

#x = x - np.mean(x) # 信号の平均値を0にする場合はコメントを解除
nfft = len(x)

# ------------------------
# PACF and confidence intervals
# ------------------------
max_lag = min(40, len(x)//2)
pacf_vals, confint = pacf(x, nlags=max_lag, alpha=0.05)
print(f"PACF with 95% CI for lags 1..{max_lag}:")
significant_lags = []
for lag in range(1, max_lag+1):
    lower, upper = confint[lag]
    val = pacf_vals[lag]
    sig = val < lower or val > upper
    marker = '*' if sig else ' '
    print(f"lag={lag:2d}: pacf={val: .4f} CI=({lower: .4f},{upper: .4f}) {marker}")
    if sig and not significant_lags:
        significant_lags.append(lag)
# Determine PACF-based order
order_pacf = significant_lags[0] if significant_lags else None
if order_pacf:
    print(f"Selected order by PACF cutoff: {order_pacf}\n")
else:
    print("No PACF spike outside CI; cannot select order by PACF.\n")

# ------------------------
# AIC / BIC selection
# ------------------------
best_aic, best_bic = np.inf, np.inf
best_p_aic, best_p_bic = None, None
max_ar = min(40, len(x)-1)
for p in range(1, max_ar+1):
    try:
        model = AutoReg(x, lags=p, old_names=False).fit()
        if model.aic < best_aic:
            best_aic, best_p_aic = model.aic, p
        if model.bic < best_bic:
            best_bic, best_p_bic = model.bic, p
    except Exception:
        pass
print(f"Selected order by AIC: {best_p_aic}, AIC={best_aic:.2f}")
print(f"Selected order by BIC: {best_p_bic}, BIC={best_bic:.2f}\n")

# ------------------------
# Final order selection
# ------------------------
# Priority: PACF -> BIC -> AIC -> preset
if order_pacf:
    order_selected = order_pacf
elif best_p_bic:
    order_selected = best_p_bic
else:
    order_selected = best_p_aic if best_p_aic else preset_order
print(f"Using AR model order = {order_selected}\n")

# ------------------------
# MEM spectrum (Burg method)
# ------------------------
burg_spec = pburg(x, order=order_selected, NFFT=nfft)
mem_psd   = burg_spec.psd
freqs_mem = burg_spec.frequencies()

# ------------------------
# FFT-based PSD
# ------------------------
fft_spec  = Periodogram(x, NFFT=nfft)
fft_psd   = fft_spec.psd
freqs_fft = fft_spec.frequencies()

# ------------------------
# Save results to Excel file
# ------------------------
if infile != "":
    # 入力ファイル名から出力ファイル名を生成
    filebody = infile.rsplit('.', 1)[0]
    output_filename = f"{filebody}-mem_order={order_selected}.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}")

# ------------------------
# Plotting with subplots
# ------------------------
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) # 2行1列のサブプロットを作成

# 1つ目のサブプロット：入力信号
axes[0].plot(t, x, label='Input signal', linewidth=2)
axes[0].set_xlabel('Time')
axes[0].set_ylabel('Amplitude')
axes[0].set_title('Input Time Series')
axes[0].grid(True)
axes[0].legend() # 凡例を追加

# 2つ目のサブプロット：スペクトル比較
axes[1].plot(freqs_mem, mem_psd, label=f'MEM (order={order_selected})', linewidth=2)
axes[1].plot(freqs_fft, fft_psd, '--', label='FFT PSD', linewidth=1)
axes[1].set_xlabel('Frequency (Hz)')
axes[1].set_ylabel('PSD')
axes[1].set_title(f'Spectral Estimate: MEM (order={order_selected}) vs FFT')
axes[1].legend()
axes[1].grid(True)
if 'yscale' in locals() and yscale == 'log': # yscaleパラメータが定義されていて'log'の場合
    axes[1].set_yscale('log')

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


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