#!/usr/bin/env python3
import sys
import numpy as np
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
# ------------------------
_a = np.loadtxt(infile, delimiter=',', skiprows=1, usecols=[0,1])
t, x = _a.T
x = x - np.mean(x)
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()

# ------------------------
# Plot input signal
# ------------------------
plt.figure(figsize=(8,4))
plt.plot(t, x, label='Input signal', linewidth=2)
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.title('Input Time Series')
plt.grid(True)
plt.show()

# ------------------------
# Plot spectral comparison
# ------------------------
plt.figure(figsize=(8,4))
plt.plot(freqs_mem, mem_psd, label=f'MEM (order={order_selected})', linewidth=2)
plt.plot(freqs_fft, fft_psd, '--', label='FFT PSD', linewidth=1)
plt.xlabel('Frequency (Hz)')
plt.ylabel('PSD')
plt.title('Spectral Estimate: MEM vs FFT')
plt.legend()
plt.grid(True)
plt.show()
