"""
FFTを用いた周期関数の補間を実行するスクリプト。

概要:
    このスクリプトは、入力ファイルから読み込んだ、または内部で生成したデータに対して、
    高速フーリエ変換（FFT）を使用して周期関数の補間を行います。
    補間されたデータと元のデータ、必要であれば厳密な関数をプロットして比較します。

詳細説明:
    スクリプトは以下のいずれかの方法で入力データを取得します。
    1. コマンドライン引数またはデフォルトで指定されたExcelファイルからデータを読み込みます。
       do_mirror フラグがTrueの場合、読み込んだデータはミラーリングされて拡張されます。
    2. Excelファイルが指定されていない、または "generate" が指定された場合、
       periodic_function で定義されたサンプルデータが生成されます。

    取得したデータはFFTによって周波数領域に変換され、
    指定された補間係数に基づいてパディングされた後、逆FFTによって時間領域に戻されます。
    これにより、元のデータ点よりも密な補間データが生成されます。
    最終的に、元のデータ、補間されたデータ、および生成された場合は厳密な関数のプロットが表示されます。

関連リンク:
    interpolate_fft_usage
"""

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

#===================
# デフォルトパラメータ
#===================
INFILE_DEF = "interpolate_fft_test.xlsx"
DO_MIRROR_DEF = False
KRANGE_DEF = [-0.5, 0.5]
N_SAMPLES_DEF = 10
INTERP_FACTOR = 10
FIGSIZE_DEF = (4, 6)

def periodic_function(k):
    """
    概要:
        特定の周期関数を計算します。
    詳細説明:
        入力値 k に基づいて、数式 -cos(2πk) * (1 + 5k^2) を評価し、その結果を返します。
        この関数は、周期的なデータ生成のためのサンプルとして使用されます。
    引数:
        :param k: 関数を評価する入力値。numpy.ndarray または float で指定します。
        :type k: numpy.ndarray or float
    戻り値:
        :returns: 関数の計算結果。入力 k と同じ型（numpy.ndarray または float）で返されます。
        :rtype: numpy.ndarray or float
    """
    return -np.cos(2.0 * np.pi * k) * (1.0 + 5.0 * k**2)

def read_data(infile, do_mirror=False):
    """
    概要:
        指定されたExcelファイルからx-yデータを読み込みます。
    詳細説明:
        与えられた Excel ファイル (infile) から、最初の2列のデータをxとyとして読み込みます。
        読み込まれたデータは NaN 値を除去されます。
        do_mirror が True の場合、読み込んだデータはy軸に対してミラーリングされ、元のデータの前に結合されます。
        これにより、データセットが周期性を維持したまま拡張されます。
    引数:
        :param infile: 読み込むExcelファイルのパス。
        :type infile: str
        :param do_mirror: データをミラーリングして拡張するかどうかを示すフラグ。Trueの場合、データがミラーリングされます。
        :type do_mirror: bool
    戻り値:
        :returns: xデータとyデータのリストのタプル。エラーが発生した場合は空のリストのタプル。
        :rtype: tuple of list
    例外:
        :raises Exception: Excelファイルの読み込み中にエラーが発生した場合にコンソールにエラーメッセージを出力します。
    """
    try:
        df = pd.read_excel(infile)
        x = df.iloc[:,0].values.tolist()
        y = df.iloc[:,1].values.tolist()
        # NaNを除外
        x = [i for i in x if str(i) != 'nan']
        y = [i for i in y if str(i) != 'nan']
        
        if do_mirror:
            _x = [-x[i] for i in range(len(x) - 1, 0, -1)]
            _x.extend(x)
            _y = [y[i] for i in range(len(y) - 1, 0, -1)]
            _y.extend(y)
            return _x, _y
        else:
            return x, y
    except Exception as e:
        print(f"Error reading {infile}: {e}")
        return [], []

def main():
    """
    概要:
        FFTを用いた周期関数の補間処理全体を制御し、結果を可視化します。
    詳細説明:
        この関数は、まずコマンドライン引数またはデフォルト設定から補間パラメータを初期化します。
        次に、指定されたExcelファイルからデータを読み込むか、
        periodic_function を使用してサンプルデータを生成します。
        取得したデータに対して高速フーリエ変換 (FFT) を適用し、周波数領域でゼロパディングを行います。
        その後、逆FFTを実行して補間された時間領域のデータを生成します。
        最後に、元のデータ、補間されたデータ、および厳密な関数（存在する場合）をmatplotlibを用いてプロットし、
        結果を可視化します。
    """
    # パラメータの初期化
    infile = INFILE_DEF
    do_mirror = DO_MIRROR_DEF
    krange = KRANGE_DEF.copy()
    n_samples = N_SAMPLES_DEF
    
    # コマンドライン引数の処理
    argv = sys.argv
    narg = len(argv)
    if narg > 1: infile = argv[1] 
    if narg > 2: do_mirror = bool(int(argv[2]))

    print(f"\nInput mode: {infile}")

    xe, ye = None, None
    
    # データの取得
    if infile is not None and infile != "" and infile != "generate":
        print(f"Reading from Excel: [{infile}]")
        x_raw, y_raw = read_data(infile, do_mirror)
        if not x_raw:
            return
            
        krange[0], krange[1] = min(x_raw), max(x_raw)
        n = len(x_raw)
        # 周期性を考慮し最後の1点を除去してFFT
        x = np.array(x_raw[:n-1])
        y = np.array(y_raw[:n-1])
        n_samples = len(x)
    else:
        print("Generating sample data")
        x = np.linspace(krange[0], krange[1], n_samples, endpoint=False)
        y = periodic_function(x)
        # 比較用の厳密解
        n_exact = n_samples * INTERP_FACTOR
        xe = np.linspace(krange[0], krange[1], n_exact, endpoint=False)
        ye = periodic_function(xe)

    # FFT補間の実行
    n_interp = len(x) * INTERP_FACTOR
    y_fft = np.fft.fft(y)

    # 周波数領域でのパディング（ゼロパディング）
    y_fft_padded = np.zeros(n_interp, dtype=complex)
    # 低周波成分を両端に配置
    half = len(x) // 2
    y_fft_padded[:half] = y_fft[:half]
    y_fft_padded[-half:] = y_fft[-half:]

    # 逆FFTによる時間領域への復元
    x_interp = np.linspace(krange[0], krange[1], n_interp, endpoint=False)
    # 強度を維持するため INTERP_FACTOR を乗じる
    y_interp = np.fft.ifft(y_fft_padded) * INTERP_FACTOR

    # 可視化
    plt.figure(figsize=FIGSIZE_DEF)
    plt.plot(x, y, 'o', label='input data', markersize=6)
    plt.plot(x_interp, y_interp.real, '-', label='interpolated', marker='x', markersize=3, alpha=0.7)
    if xe is not None:
        plt.plot(xe, ye, '--', label='exact', alpha=0.5)
        
    plt.xlabel('x')
    plt.ylabel('y')
    plt.legend()
    plt.title('FFT Interpolation of Periodic Function')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()