#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
exafs_mem_fft.py

EXAFS-oriented FFT/MEM transform for k-space data.

Input data:
    Column 0 : k [A^-1]
    Column 1 : signal, usually chi(k), k^2 chi(k), etc.

Main features:
    - EXAFS R-axis correction: chi(k) ~ sin(2 k R + phase)
      Therefore R = pi * frequency, not 2*pi * frequency.
    - mode = fft, mem, both
    - FFT output: real, imag, amplitude, power
    - Zero padding by --npadding, default 1024
    - k-weight by --korder
    - EXAFS-style Hanning window by --hrange and --hwidth

Examples:
    # Input is already k^2 chi(k), so do not apply extra k^2
    python exafs_mem_fft.py a-GeO2_glass_0Pa.xlsx --mode fft --korder 0 --hrange 2 12 --hwidth 0.5

    # Input is raw chi(k), apply k^2 weighting
    python exafs_mem_fft.py chi.xlsx --mode fft --korder 2 --hrange 2 12 --hwidth 0.5

    # MEM, fixed AR order
    python exafs_mem_fft.py a-GeO2_glass_0Pa.xlsx --mode mem --order 20 --korder 0 --hrange 2 12 --hwidth 0.5

    # Both FFT and MEM
    python exafs_mem_fft.py a-GeO2_glass_0Pa.xlsx --mode both --order 20 --korder 0 --hrange 2 12 --hwidth 0.5
"""

import argparse
import os
import sys
import warnings

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from spectrum import pburg
from statsmodels.tsa.ar_model import AutoReg
from statsmodels.tsa.stattools import pacf

warnings.filterwarnings("ignore")


# ============================================================
# Input
# ============================================================

def load_xy(infile, sheet=0, xcol=0, ycol=1):
    """
    Load two-column data.

    Parameters
    ----------
    infile : str
        Input file.
    sheet : int or str
        Excel sheet index or name.
    xcol : int
        k-axis column, 0-based.
    ycol : int
        signal column, 0-based.

    Returns
    -------
    k : ndarray
        k-axis [A^-1]
    signal : ndarray
        signal column
    """
    ext = os.path.splitext(infile)[1].lower()

    if ext == ".xlsx":
        df = pd.read_excel(infile, sheet_name=sheet)
        k = df.iloc[:, xcol].to_numpy(dtype=float)
        signal = df.iloc[:, ycol].to_numpy(dtype=float)

    elif ext == ".csv":
        df = pd.read_csv(infile)
        k = df.iloc[:, xcol].to_numpy(dtype=float)
        signal = df.iloc[:, ycol].to_numpy(dtype=float)

    elif ext in [".txt", ".dat"]:
        arr = np.loadtxt(infile)
        if arr.ndim != 2:
            raise ValueError("Text input must be a 2D table.")
        if arr.shape[1] <= max(xcol, ycol):
            raise ValueError(
                f"Text input must have at least {max(xcol, ycol) + 1} columns."
            )
        k = arr[:, xcol].astype(float)
        signal = arr[:, ycol].astype(float)

    else:
        raise ValueError("Unsupported file format. Use .xlsx, .csv, .txt, or .dat.")

    mask = np.isfinite(k) & np.isfinite(signal)
    k = k[mask]
    signal = signal[mask]

    if len(k) < 4:
        raise ValueError("Not enough valid data points.")

    if np.any(np.diff(k) < 0):
        print("WARNING: k-axis is not monotonic. Sorting by k.")
        idx = np.argsort(k)
        k = k[idx]
        signal = signal[idx]

    return k, signal


def estimate_dk(k, rtol=1e-3):
    """
    Estimate dk from k-axis.
    """
    dk_arr = np.diff(k)
    dk = np.mean(dk_arr)

    if dk == 0:
        raise ValueError("dk is zero.")

    rel_std = np.std(dk_arr) / abs(dk)

    if rel_std > rtol:
        print("WARNING: k-axis is not perfectly uniform.")
        print(f"  mean dk      = {dk:.10g}")
        print(f"  std(diff)    = {np.std(dk_arr):.10g}")
        print(f"  relative std = {rel_std:.3e}")
        print("  FFT/MEM assume uniformly sampled k data.")
        print("  Consider interpolation onto a uniform k grid if needed.")

    return dk, rel_std


# ============================================================
# Window and preprocessing
# ============================================================

def make_hanning_range_window(k, hrange=None, hwidth=0.5):
    """
    EXAFS-style Hanning window on k-axis.

    Parameters
    ----------
    k : ndarray
        k-axis [A^-1]
    hrange : None or list/tuple
        If None, window = 1 over all data.
        If given, hrange = [k_start, k_end].
    hwidth : float
        Width of cosine taper at both ends [A^-1].

    Window definition
    -----------------
    window = 0 outside [k1, k2]

    left edge:
        k1 <= k < k1 + hwidth
        w = 0.5 * (1 - cos(pi * (k-k1)/hwidth))

    flat region:
        k1 + hwidth <= k <= k2 - hwidth
        w = 1

    right edge:
        k2 - hwidth < k <= k2
        w = 0.5 * (1 + cos(pi * (k-(k2-hwidth))/hwidth))
    """
    k = np.asarray(k, dtype=float)

    if hrange is None:
        return np.ones_like(k, dtype=float)

    k1 = float(hrange[0])
    k2 = float(hrange[1])

    if k2 < k1:
        k1, k2 = k2, k1

    if k2 <= k1:
        raise ValueError("--hrange requires two different values: KSTART KEND")

    if hwidth < 0:
        raise ValueError("--hwidth must be >= 0")

    if 2.0 * hwidth > (k2 - k1):
        new_hwidth = 0.5 * (k2 - k1)
        print("WARNING: 2*hwidth is larger than hrange width.")
        print(f"         hwidth changed from {hwidth:.10g} to {new_hwidth:.10g}")
        hwidth = new_hwidth

    win = np.zeros_like(k, dtype=float)

    if hwidth == 0:
        win[(k >= k1) & (k <= k2)] = 1.0
        return win

    left = (k >= k1) & (k < k1 + hwidth)
    flat = (k >= k1 + hwidth) & (k <= k2 - hwidth)
    right = (k > k2 - hwidth) & (k <= k2)

    win[left] = 0.5 * (
        1.0 - np.cos(np.pi * (k[left] - k1) / hwidth)
    )

    win[flat] = 1.0

    win[right] = 0.5 * (
        1.0 + np.cos(np.pi * (k[right] - (k2 - hwidth)) / hwidth)
    )

    return win


def preprocess_signal(k, signal_raw, korder=0.0, hrange=None, hwidth=0.5,
                      remove_mean=False):
    """
    Apply k-weight, optional mean subtraction, and Hanning window.

    Important:
        If input is already k^2 chi(k), use --korder 0.
        If input is raw chi(k), use --korder 2 for k^2 chi(k).
    """
    k = np.asarray(k, dtype=float)
    signal_raw = np.asarray(signal_raw, dtype=float)

    kweight = np.power(k, korder)
    signal_kweighted = signal_raw * kweight

    if remove_mean:
        signal_centered = signal_kweighted - np.mean(signal_kweighted)
    else:
        signal_centered = signal_kweighted.copy()

    window = make_hanning_range_window(k, hrange=hrange, hwidth=hwidth)
    signal_used = signal_centered * window

    return signal_used, signal_kweighted, signal_centered, window


# ============================================================
# AR order selection
# ============================================================

def select_ar_order(signal, preset_order=8, max_order=40, order_method="bic"):
    """
    Select AR model order by AIC/BIC or fixed preset.

    PACF is printed only as diagnostic information.
    """
    n = len(signal)

    max_lag = min(max_order, n // 2)

    try:
        pacf_vals, confint = pacf(signal, 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 = (lower > 0.0) or (upper < 0.0)
            marker = "*" if sig else " "

            print(
                f"lag={lag:2d}: pacf={val: .4f} "
                f"CI=({lower: .4f},{upper: .4f}) {marker}"
            )

            if sig:
                significant_lags.append(lag)

        if significant_lags:
            print(
                "PACF significant lags, diagnostic only: "
                f"{significant_lags[:10]}"
            )
        else:
            print("No PACF spike outside zero-confidence interval.")

        print("PACF is not used as the primary order selector.\n")

    except Exception as e:
        print(f"PACF calculation failed: {e}")
        print("Continue with AIC/BIC order selection.\n")

    best_aic = np.inf
    best_bic = np.inf
    best_p_aic = None
    best_p_bic = None

    max_ar = min(max_order, n - 1)

    for p in range(1, max_ar + 1):
        try:
            model = AutoReg(signal, lags=p, old_names=False).fit()

            if model.aic < best_aic:
                best_aic = model.aic
                best_p_aic = p

            if model.bic < best_bic:
                best_bic = model.bic
                best_p_bic = p

        except Exception:
            pass

    print(f"Selected order by AIC: {best_p_aic}, AIC={best_aic:.6g}")
    print(f"Selected order by BIC: {best_p_bic}, BIC={best_bic:.6g}\n")

    if order_method == "bic":
        if best_p_bic is not None:
            order_selected = best_p_bic
        elif best_p_aic is not None:
            order_selected = best_p_aic
        else:
            order_selected = preset_order

    elif order_method == "aic":
        if best_p_aic is not None:
            order_selected = best_p_aic
        elif best_p_bic is not None:
            order_selected = best_p_bic
        else:
            order_selected = preset_order

    elif order_method == "preset":
        order_selected = preset_order

    else:
        raise ValueError("order_method must be one of: bic, aic, preset")

    print(f"Order selection method = {order_method}")
    print(f"Using AR model order = {order_selected}\n")

    return order_selected, {
        "order_method": order_method,
        "best_p_aic": best_p_aic,
        "best_aic": best_aic,
        "best_p_bic": best_p_bic,
        "best_bic": best_bic,
    }


# ============================================================
# FFT
# ============================================================

def calc_fft_exafs(signal_used, dk, nfft):
    """
    Calculate EXAFS-style FFT.

    EXAFS phase:
        chi(k) ~ sin(2 k R + phi)

    numpy rfftfreq:
        freq [cycles / A^-1]

    Since
        2 k R = 2 pi f k

    the R-axis is:
        R = pi * f

    FT is multiplied by dk to mimic the integral scale.
    """
    ft_complex = np.fft.rfft(signal_used, n=nfft) * dk
    freq = np.fft.rfftfreq(nfft, d=dk)

    R = np.pi * freq

    ft_real = ft_complex.real
    ft_imag = ft_complex.imag
    ft_amp = np.abs(ft_complex)
    ft_power = ft_amp ** 2

    return R, ft_complex, ft_real, ft_imag, ft_amp, ft_power


# ============================================================
# MEM
# ============================================================

def calc_mem_exafs(signal_used, dk, order, nfft):
    """
    Calculate MEM spectrum using spectrum.pburg.

    pburg.frequencies() returns normalized frequency in cycles/sample.
    Convert to physical cycles per A^-1 by dividing by dk.
    For EXAFS R-axis:
        R = pi * f_phys = pi * f_sample / dk
    """
    burg_spec = pburg(signal_used, order=order, NFFT=nfft)

    mem_psd = np.asarray(burg_spec.psd, dtype=float)
    freqs_sample = np.asarray(burg_spec.frequencies(), dtype=float)

    R = np.pi * freqs_sample / dk

    return R, freqs_sample, mem_psd


# ============================================================
# Utilities
# ============================================================

def find_main_peak(R, y, rmin=1e-12):
    """
    Find main peak excluding R=0.
    """
    R = np.asarray(R)
    y = np.asarray(y)

    mask = np.isfinite(R) & np.isfinite(y) & (R > rmin)

    if not np.any(mask):
        return np.nan, np.nan

    R2 = R[mask]
    y2 = y[mask]

    idx = np.argmax(y2)

    return R2[idx], y2[idx]


def make_output_filename(infile, mode, order, korder):
    stem, _ = os.path.splitext(infile)

    if mode == "mem":
        return f"{stem}_exafs_mem_order={order}_korder={korder:g}.xlsx"
    elif mode == "fft":
        return f"{stem}_exafs_fft_korder={korder:g}.xlsx"
    else:
        return f"{stem}_exafs_fft_mem_order={order}_korder={korder:g}.xlsx"


def pad_array(a, n):
    a = np.asarray(a)
    out = np.full(n, np.nan, dtype=float)
    out[:len(a)] = a
    return out


# ============================================================
# Save Excel
# ============================================================

def save_results_excel(
    outfile,
    args,
    k,
    signal_raw,
    signal_kweighted,
    signal_centered,
    window,
    signal_used,
    dk,
    rel_std_dk,
    order_selected,
    order_info,
    fft_result,
    mem_result,
):
    """
    Save input, FFT/MEM results, and summary to Excel.
    """
    df_input = pd.DataFrame({
        "k_A^-1": k,
        "signal_raw": signal_raw,
        f"signal_k^{args.korder:g}_weighted": signal_kweighted,
        "signal_centered": signal_centered,
        "hanning_window": window,
        "signal_used": signal_used,
    })

    summary_items = [
        "input_file",
        "sheet",
        "xcol",
        "ycol",
        "mode",
        "n_data",
        "npadding",
        "nfft_used",
        "dk [A^-1]",
        "relative_std_dk",
        "korder",
        "remove_mean",
        "hrange",
        "hwidth [A^-1]",
        "AR_order_selected",
        "order_method",
        "best_p_aic",
        "best_aic",
        "best_p_bic",
        "best_bic",
        "R_axis_definition",
        "FFT_definition",
    ]

    summary_values = [
        args.infile,
        args.sheet,
        args.xcol,
        args.ycol,
        args.mode,
        len(k),
        args.npadding,
        max(args.npadding, len(k)),
        dk,
        rel_std_dk,
        args.korder,
        args.remove_mean,
        "" if args.hrange is None else str(args.hrange),
        args.hwidth,
        order_selected if order_selected is not None else "",
        order_info.get("order_method", "") if order_info else "",
        order_info.get("best_p_aic", "") if order_info else "",
        order_info.get("best_aic", "") if order_info else "",
        order_info.get("best_p_bic", "") if order_info else "",
        order_info.get("best_bic", "") if order_info else "",
        "EXAFS: R = pi * frequency, because chi(k) ~ sin(2 k R + phase)",
        "FT = rFFT(signal_used, nfft) * dk",
    ]

    with pd.ExcelWriter(outfile, engine="openpyxl") as writer:
        df_input.to_excel(writer, sheet_name="Input", index=False)

        if fft_result is not None:
            (
                R_fft,
                ft_complex,
                ft_real,
                ft_imag,
                ft_amp,
                ft_power,
            ) = fft_result

            peak_R_fft, peak_y_fft = find_main_peak(R_fft, ft_amp)

            df_fft = pd.DataFrame({
                "R_A": R_fft,
                "FT_real": ft_real,
                "FT_imag": ft_imag,
                "FT_amplitude": ft_amp,
                "FT_power": ft_power,
            })
            df_fft.to_excel(writer, sheet_name="FFT", index=False)

            summary_items += [
                "main_peak_R_FFT_amplitude [A]",
                "main_peak_FFT_amplitude",
            ]
            summary_values += [
                peak_R_fft,
                peak_y_fft,
            ]

        if mem_result is not None:
            R_mem, freqs_mem_sample, mem_psd = mem_result

            peak_R_mem, peak_y_mem = find_main_peak(R_mem, mem_psd)

            df_mem = pd.DataFrame({
                "Frequency_MEM_cycles_per_sample": freqs_mem_sample,
                "R_A": R_mem,
                "MEM_PSD": mem_psd,
            })
            df_mem.to_excel(writer, sheet_name="MEM", index=False)

            summary_items += [
                "main_peak_R_MEM [A]",
                "main_peak_MEM_PSD",
            ]
            summary_values += [
                peak_R_mem,
                peak_y_mem,
            ]

        df_summary = pd.DataFrame({
            "item": summary_items,
            "value": summary_values,
        })

        df_summary.to_excel(writer, sheet_name="Summary", index=False)


# ============================================================
# Plot
# ============================================================

def plot_results(
    args,
    k,
    signal_raw,
    signal_used,
    window,
    fft_result,
    mem_result,
    order_selected,
):
    """
    Plot input/window and transform result.
    """
    fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 8))

    ax0 = axes[0]
    ax0.plot(k, signal_raw, label="raw signal", linewidth=1.2)
    ax0.plot(k, signal_used, label="used signal", linewidth=1.2)

    # Plot window on secondary axis
    ax0b = ax0.twinx()
    ax0b.plot(k, window, linestyle="--", linewidth=1.0, label="Hanning window")
    ax0b.set_ylabel("window")

    ax0.set_xlabel("k [A$^{-1}$]")
    ax0.set_ylabel("signal")
    ax0.set_title("Input and processed signal")
    ax0.grid(True)

    lines0, labels0 = ax0.get_legend_handles_labels()
    lines1, labels1 = ax0b.get_legend_handles_labels()
    ax0.legend(lines0 + lines1, labels0 + labels1, loc="best")

    ax1 = axes[1]

    if fft_result is not None:
        (
            R_fft,
            ft_complex,
            ft_real,
            ft_imag,
            ft_amp,
            ft_power,
        ) = fft_result

        if args.fft_plot == "realimag":
            ax1.plot(R_fft, ft_real, label="FFT real", linewidth=1.2)
            ax1.plot(R_fft, ft_imag, label="FFT imag", linewidth=1.2)
            y_for_peak = ft_amp
        elif args.fft_plot == "power":
            ax1.plot(R_fft, ft_power, label="FFT power", linewidth=1.2)
            y_for_peak = ft_power
        else:
            ax1.plot(R_fft, ft_amp, label="FFT amplitude", linewidth=1.2)
            y_for_peak = ft_amp

        peak_R_fft, peak_y_fft = find_main_peak(R_fft, y_for_peak)
        if np.isfinite(peak_R_fft):
            ax1.axvline(peak_R_fft, linestyle="--", linewidth=0.8)
            ax1.text(
                peak_R_fft,
                np.nanmax(y_for_peak[np.isfinite(y_for_peak)]) * 0.8,
                f"FFT R={peak_R_fft:.3g} A",
                rotation=90,
                va="top",
                ha="right",
            )

    if mem_result is not None:
        R_mem, freqs_mem_sample, mem_psd = mem_result
        ax1.plot(
            R_mem,
            mem_psd,
            label=f"MEM PSD, order={order_selected}",
            linewidth=1.8,
        )

        peak_R_mem, peak_y_mem = find_main_peak(R_mem, mem_psd)
        if np.isfinite(peak_R_mem):
            ax1.axvline(peak_R_mem, linestyle=":", linewidth=0.8)
            ax1.text(
                peak_R_mem,
                np.nanmax(mem_psd[np.isfinite(mem_psd)]) * 0.6,
                f"MEM R={peak_R_mem:.3g} A",
                rotation=90,
                va="top",
                ha="left",
            )

    ax1.set_xlabel("R [A]")
    ax1.set_ylabel("FT / PSD")
    ax1.set_title("EXAFS transform")
    ax1.grid(True)
    ax1.legend()

    if args.yscale == "log":
        ax1.set_yscale("log")

    if args.xlim is not None:
        ax1.set_xlim(args.xlim[0], args.xlim[1])

    plt.tight_layout()

    if args.savefig:
        plt.savefig(args.savefig, dpi=300)
        print(f"Figure saved to {args.savefig}")

    if not args.no_show:
        plt.show()


# ============================================================
# Main
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description=(
            "EXAFS-oriented FFT/MEM transform. "
            "Input column 0 is k [A^-1], column 1 is chi(k)-like signal."
        )
    )

    parser.add_argument(
        "infile",
        help="Input file: .xlsx, .csv, .txt, or .dat",
    )

    parser.add_argument(
        "preset_order",
        nargs="?",
        type=int,
        default=8,
        help="Fallback/fixed AR model order for MEM. Default: 8.",
    )

    parser.add_argument(
        "--mode",
        choices=["fft", "mem", "both"],
        default="fft",
        help="Transform mode. Default: fft.",
    )

    parser.add_argument(
        "--sheet",
        default=0,
        help="Excel sheet name or index. Default: 0.",
    )

    parser.add_argument(
        "--xcol",
        type=int,
        default=0,
        help="k-axis column index, 0-based. Default: 0.",
    )

    parser.add_argument(
        "--ycol",
        type=int,
        default=1,
        help="signal column index, 0-based. Default: 1.",
    )

    parser.add_argument(
        "--npadding",
        type=int,
        default=1024,
        help=(
            "FFT/MEM NFFT length. If smaller than data length, data length is used. "
            "Default: 1024."
        ),
    )

    parser.add_argument(
        "--korder",
        type=float,
        default=0.0,
        help=(
            "Apply k^korder to the input signal before transform. "
            "If input is already k^2 chi(k), use --korder 0. Default: 0."
        ),
    )

    parser.add_argument(
        "--hrange",
        type=float,
        nargs=2,
        default=None,
        metavar=("KSTART", "KEND"),
        help="Hanning window k range [A^-1], e.g. --hrange 2 12.",
    )

    parser.add_argument(
        "--hwidth",
        type=float,
        default=0.5,
        help="Hanning taper width at both ends of --hrange [A^-1]. Default: 0.5.",
    )

    parser.add_argument(
        "--remove-mean",
        action="store_true",
        help="Subtract mean after k-weighting and before windowing.",
    )

    parser.add_argument(
        "--max-order",
        type=int,
        default=40,
        help="Maximum AR order tested by AIC/BIC. Default: 40.",
    )

    parser.add_argument(
        "--order-method",
        choices=["bic", "aic", "preset"],
        default="preset",
        help=(
            "AR order selection method for MEM. "
            "Default: preset, i.e. use positional preset_order."
        ),
    )

    parser.add_argument(
        "--yscale",
        choices=["linear", "log"],
        default="linear",
        help="Y scale for transform plot. Default: linear.",
    )

    parser.add_argument(
        "--fft-plot",
        choices=["amplitude", "power", "realimag"],
        default="amplitude",
        help="What to plot for FFT mode. Default: amplitude.",
    )

    parser.add_argument(
        "--xlim",
        type=float,
        nargs=2,
        default=None,
        help="R-axis plot range, e.g. --xlim 0 6.",
    )

    parser.add_argument(
        "--outfile",
        default="",
        help="Output Excel file. Default: auto-generated.",
    )

    parser.add_argument(
        "--savefig",
        default="",
        help="Save figure filename, e.g. result.png.",
    )

    parser.add_argument(
        "--no-show",
        action="store_true",
        help="Do not show matplotlib window.",
    )

    args = parser.parse_args()

    # Excel sheet can be integer-like string
    try:
        args.sheet = int(args.sheet)
    except Exception:
        pass

    infile = args.infile

    try:
        k, signal_raw = load_xy(
            infile,
            sheet=args.sheet,
            xcol=args.xcol,
            ycol=args.ycol,
        )
    except Exception as e:
        print(f"Error loading input file: {e}")
        sys.exit(1)

    try:
        dk, rel_std_dk = estimate_dk(k)
    except Exception as e:
        print(f"Error estimating dk: {e}")
        sys.exit(1)

    try:
        signal_used, signal_kweighted, signal_centered, window = preprocess_signal(
            k,
            signal_raw,
            korder=args.korder,
            hrange=args.hrange,
            hwidth=args.hwidth,
            remove_mean=args.remove_mean,
        )
    except Exception as e:
        print(f"Error in preprocessing: {e}")
        sys.exit(1)

    n_data = len(k)
    nfft = max(args.npadding, n_data)

    print()
    print("Input data summary")
    print(f"  file          = {infile}")
    print(f"  mode          = {args.mode}")
    print(f"  n_data        = {n_data}")
    print(f"  npadding      = {args.npadding}")
    print(f"  nfft used     = {nfft}")
    print(f"  k_min         = {np.min(k):.10g} [A^-1]")
    print(f"  k_max         = {np.max(k):.10g} [A^-1]")
    print(f"  dk            = {dk:.10g} [A^-1]")
    print(f"  rel std dk    = {rel_std_dk:.3e}")
    print(f"  korder        = {args.korder:g}")
    print(f"  remove mean   = {args.remove_mean}")
    print(f"  hrange        = {args.hrange}")
    print(f"  hwidth        = {args.hwidth}")
    print(f"  mean raw      = {np.mean(signal_raw):.10g}")
    print(f"  mean used     = {np.mean(signal_used):.10g}")
    print()

    fft_result = None
    mem_result = None
    order_selected = None
    order_info = {}

    if args.mode in ["fft", "both"]:
        fft_result = calc_fft_exafs(signal_used, dk, nfft)

        R_fft, ft_complex, ft_real, ft_imag, ft_amp, ft_power = fft_result
        peak_R_fft, peak_y_fft = find_main_peak(R_fft, ft_amp)

        print("FFT result")
        print(f"  main peak by amplitude: R = {peak_R_fft:.10g} [A], amp = {peak_y_fft:.10g}")
        print()

    if args.mode in ["mem", "both"]:
        try:
            order_selected, order_info = select_ar_order(
                signal_used,
                preset_order=args.preset_order,
                max_order=args.max_order,
                order_method=args.order_method,
            )
        except Exception as e:
            print(f"Error selecting AR order: {e}")
            sys.exit(1)

        try:
            mem_result = calc_mem_exafs(signal_used, dk, order_selected, nfft)
        except Exception as e:
            print(f"Error calculating MEM spectrum: {e}")
            sys.exit(1)

        R_mem, freqs_mem_sample, mem_psd = mem_result
        peak_R_mem, peak_y_mem = find_main_peak(R_mem, mem_psd)

        print("MEM result")
        print(f"  order      = {order_selected}")
        print(f"  main peak  = R = {peak_R_mem:.10g} [A], PSD = {peak_y_mem:.10g}")
        print()

    if args.outfile:
        outfile = args.outfile
    else:
        outfile = make_output_filename(
            infile,
            args.mode,
            order_selected if order_selected is not None else args.preset_order,
            args.korder,
        )

    try:
        save_results_excel(
            outfile,
            args,
            k,
            signal_raw,
            signal_kweighted,
            signal_centered,
            window,
            signal_used,
            dk,
            rel_std_dk,
            order_selected,
            order_info,
            fft_result,
            mem_result,
        )
        print(f"Results saved to {outfile}")
    except Exception as e:
        print(f"Error saving Excel file: {e}")

    plot_results(
        args,
        k,
        signal_raw,
        signal_used,
        window,
        fft_result,
        mem_result,
        order_selected,
    )

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


if __name__ == "__main__":
    main()