import math
import os
import sys

# Constants (approximating BASIC's graphic settings for context, though not used for drawing)
X0 = 75
Y0 = 185
XL = 512
YL = 150
MS = 8

def clear_screen():
    """Clears the console screen."""
    os.system('cls' if os.name == 'nt' else 'clear')

def data_input_routine(ND, X):
    """
    Data Input Routine
    Reads data from a specified file.
    """
    while True:
        nm = input(" *** データファイルメイ   =")
        try:
            with open(nm, 'r') as f:
                for i in range(ND):
                    line = f.readline().strip()
                    if not line:
                        print(f"ファイルの終わりに到達しました。{i+1}行目が読み取れませんでした。")
                        break
                    # Assuming the data file contains three values per line,
                    # similar to the BASIC code's INPUT#1,X(I),BBB,CCC
                    try:
                        parts = line.split(',')
                        if len(parts) >= 1: # Only X(I) is strictly needed
                            X[i] = float(parts[0])
                            print(f"I={i+1}   X={X[i]}")
                        else:
                            print(f"警告: ファイルの行 {i+1} が期待される形式ではありません。")
                            X[i] = 0.0 # Default value
                    except ValueError:
                        print(f"警告: ファイルの行 {i+1} のデータが数値ではありません。スキップします。")
                        X[i] = 0.0 # Default value
                else: # This else block executes if the loop completes without a break
                    break
        except FileNotFoundError:
            print(f"エラー: ファイル '{nm}' が見つかりません。")
        except Exception as e:
            print(f"ファイル読み込み中にエラーが発生しました: {e}")

def ar_model_estimation_routine(ND, MMAX, X, Y, FPE, R, RR, RFPE, ISW):
    """
    AR-model Estimation Routine
    Estimates auto-regressive coefficients using Burg's method.
    """
    # Calculate average and subtract it from data
    data_sum = sum(X[:ND])
    av = data_sum / ND

    for i in range(ND):
        z = X[i] - av
        X[i] = z
        if i < ND -1 : # Y is 0-indexed and stores up to ND-1 elements
            Y[i] = z

    pm_sum = sum(z * z for z in X[:ND])
    pm = pm_sum / ND

    fpemin = (ND + 1) / (ND - 1) * pm
    FPE[0] = fpemin
    
    min_m = 0
    p_mm = pm
    
    # Initialize RFPE with zeros or appropriate default values
    for i in range(MMAX + 1): # RFPE should be sized correctly to hold MMAX values
        RFPE[i] = 0.0 

    for m in range(1, MMAX + 1):
        sumn = 0
        sumd = 0
        for i in range(ND - m):
            sumn += X[i] * Y[i]
            sumd += X[i] * X[i] + Y[i] * Y[i]
        
        if sumd == 0: # Avoid division by zero
            rm = 0
        else:
            rm = -2 * sumn / sumd
        R[m] = rm

        pm = pm * (1 - rm * rm)

        if m > 1:
            for i in range(1, m):
                R[i] = RR[i] + rm * RR[m - i]
        
        for i in range(1, m + 1):
            RR[i] = R[i]
        
        for i in range(ND - m - 1):
            temp_x_i = X[i] # Store X[i] before it's updated
            X[i] = X[i] + rm * Y[i]
            Y[i] = Y[i + 1] + rm * temp_x_i # Use the original X[i] for Y[i] calculation
            
        if ISW == 0:  # If model order is not given, find optimal order
            current_fpe = (ND + m + 1) / (ND - m - 1) * pm
            FPE[m] = current_fpe
            if current_fpe < fpemin:
                fpemin = current_fpe
                min_m = m
                p_mm = pm
                for i in range(1, m + 1):
                    RFPE[i] = R[i]
        
    if ISW == 1:  # If model order is given, use it
        min_m = MMAX
        p_mm = pm
        for i in range(1, MMAX + 1):
            RFPE[i] = R[i]
            
    return min_m, p_mm

def power_spectrum_estimation_routine(MINM, PMM, DT, WNMIN, WNMAX, NS, X, RFPE):
    """
    Power Spectrum Estimation Routine
    Calculates the power spectrum using the estimated AR coefficients.
    """
    ci = 2 * math.pi * DT
    
    # Ensure X is large enough to store NS spectrum points
    # If X was sized 1024, it should be fine as NS is typically smaller.
    
    step_size = (WNMAX - WNMIN) / NS if NS > 0 else 0
    
    i = 0
    f = WNMIN
    while i < NS:
        sum1 = 1
        sum2 = 0
        for j in range(1, MINM + 1):
            sum1 += RFPE[j] * math.cos(ci * f * j)
            sum2 += RFPE[j] * math.sin(ci * f * j)
        
        denominator = (sum1 * sum1 + sum2 * sum2)
        if denominator == 0:
            X[i] = float('inf') # Handle division by zero
        else:
            X[i] = PMM * DT / denominator
        
        f += step_size
        i += 1

def output_routine(NS, X, N_str, WNMIN, WNMAX):
    """
    Output Routine
    Displays results and handles saving data.
    Note: Graphic display elements are commented out as they require specific libraries.
    """
    # X0=75:Y0=185:XL=512:YL=150:MS=8 (Constants already defined globally)

    if NS == 0:
        print("スペクトルデータがありません。")
        return

    xmax_val = 0
    if NS > 0:
        xmax_val = X[0]
        for i in range(1, NS):
            if X[i] > xmax_val:
                xmax_val = X[i]
    
    clear_screen()
    
    # Graphic drawing commands are commented out
    # LINE (X0,Y0+3)-(X0,Y0-YL)
    # LINE (X0-3,Y0)-(X0+XL,Y0)
    # FOR I=1 TO MS
    #   XX=X0+XL/MS*I:LINE@(XX,Y0+1)-(XX,Y0+3)
    # NEXT
    
    print(f"最小周波数: {WNMIN}")
    print(f"最大周波数: {WNMAX}")
    print(f"最大スペクトル値: {xmax_val}")

    # SC=XL/NS:SMAX=XMAX/YL (These are for graphic scaling, not directly used in text output)
    # FOR I=1 TO NS
    #   PSET(X0+SC*I,Y0-X(I)/XMAX)
    # NEXT

    nn = 15 - len(N_str) // 2
    display_title = "*" * nn + N_str + "*" * nn
    print(f"{display_title:^80}") # Center the title
    
    while True:
        an_hardcopy = input(f"{' ' * 20}Hard Copy(Y/N)? ").upper()
        if an_hardcopy == "Y":
            print("ハードコピー機能は実装されていません。") # Placeholder for graphic hardcopy
            break
        elif an_hardcopy == "N":
            break
        else:
            print("YまたはNを入力してください。")

    while True:
        an_save = input(f"{' ' * 20}Data Save(Y/N)? ").upper()
        if an_save == "N":
            break
        elif an_save == "Y":
            nm_save = input("   Data File Name: ")
            try:
                # The original BASIC used "A:" prefix, which is for floppy drives.
                # Here, we'll just use the filename directly in the current directory.
                with open(nm_save, 'w') as f:
                    for i in range(NS):
                        f.write(f"{X[i]}\n")
                print(f"データは '{nm_save}' に保存されました。")
            except Exception as e:
                print(f"データの保存中にエラーが発生しました: {e}")
            break
        else:
            print("YまたはNを入力してください。")


## Main Routine
def main():
    clear_screen()
    print("      ***** MEM *****\n")

    while True:
        try:
            nd = int(input(" *** サンプル テンスウ    ="))
            if nd <= 0:
                print("サンプル点数は正の整数である必要があります。")
            else:
                break
        except ValueError:
            print("無効な入力です。数値を入力してください。")
    print()

    while True:
        try:
            dt = float(input(" *** サンプル カンカク    ="))
            if dt <= 0:
                print("サンプル間隔は正の数である必要があります。")
            else:
                break
        except ValueError:
            print("無効な入力です。数値を入力してください。")
    print()

    while True:
        an = input("モデル ジスウ ヲ アタエマスカ (Y/N)? ").upper()
        if an == "Y":
            isw = 1
            prompt_str = " *** モデル ジスウ      ="
            break
        elif an == "N":
            isw = 0
            prompt_str = " *** サイダイ モデル ジスウ   ="
            break
        else:
            print("YまたはNを入力してください。")
    print()

    while True:
        try:
            mmax = int(input(prompt_str))
            if mmax <= 0:
                print("モデル次数は正の整数である必要があります。")
            else:
                break
        except ValueError:
            print("無効な入力です。数値を入力してください。")
    print()

    # Initialize arrays (Python lists are dynamic, but pre-allocating for clarity)
    # BASIC DIM X(1024) -> Python list of 1024 elements (0 to 1023)
    # Ensure ND, MMAX are within reasonable bounds for array sizes if fixed size arrays were strictly needed.
    X = [0.0] * 1024 
    Y = [0.0] * nd # Y is used up to ND-1
    FPE = [0.0] * (mmax + 1) # FPE(0) to FPE(MMAX)
    R = [0.0] * (mmax + 1)   # R(0) to R(MMAX)
    RR = [0.0] * (mmax + 1)  # RR(0) to RR(MMAX)
    RFPE = [0.0] * (mmax + 1) # RFPE(0) to RFPE(MMAX)

    data_input_routine(nd, X)  # データ入力

    ns = nd
    wnmin = 0
    wnmax = nd
    
    n_str_data = " Data "
    output_routine(ns, X, n_str_data, wnmin, wnmax)  # データのグラフ表示 (テキスト出力のみ)

    min_m, p_mm = ar_model_estimation_routine(nd, mmax, X, Y, FPE, R, RR, RFPE, isw) # Burg法による自己回帰係数と予測誤差の推定

    while True:
        try:
            wnmin = float(input("   サイショウ シュウハスウ ="))
            break
        except ValueError:
            print("無効な入力です。数値を入力してください。")

    while True:
        try:
            wnmax = float(input("   サイダイ シュウハスウ ="))
            if wnmax <= wnmin:
                print("最大周波数は最小周波数より大きい必要があります。")
            else:
                break
        except ValueError:
            print("無効な入力です。数値を入力してください。")

    while True:
        try:
            ns = int(input("   スペクトル テンスウ  ="))
            if ns <= 0:
                print("スペクトル点数は正の整数である必要があります。")
            else:
                break
        except ValueError:
            print("無効な入力です。数値を入力してください。")

    # Recalculate X for spectrum storage if NS is significantly different from 1024
    if ns > 1024:
        X = [0.0] * ns # Resize X if needed for spectrum points
    
    power_spectrum_estimation_routine(min_m, p_mm, dt, wnmin, wnmax, ns, X, RFPE)  # スペクトルの計算

    n_str_mem = " MEM Spectrum "
    output_routine(ns, X, n_str_mem, wnmin, wnmax)  # スペクトル推定結果のグラフ (テキスト出力のみ)

    while True:
        an_again = input(" モウイチド シマスカ (Y/N)? ").upper()
        if an_again == "Y":
            # This would typically loop back to line 1220 in BASIC, which is setting WNMIN/WNMAX/NS
            # For Python, we'll re-run the main part of the program from ar_model_estimation onward,
            # but keep the current data (X, ND, DT, MMAX, ISW, MINM, P_MM, RFPE).
            # To truly replicate, we'd need to re-call ar_model_estimation_routine if that data changes.
            # For now, we assume it just means re-doing the spectrum calculation with new WNMIN/MAX/NS.
            # A more robust structure would involve a function for the main loop.
            print("\nスペクトル計算を再度実行します。\n")
            
            while True:
                try:
                    wnmin = float(input("   サイショウ シュウハスウ ="))
                    break
                except ValueError:
                    print("無効な入力です。数値を入力してください。")

            while True:
                try:
                    wnmax = float(input("   サイダイ シュウハスウ ="))
                    if wnmax <= wnmin:
                        print("最大周波数は最小周波数より大きい必要があります。")
                    else:
                        break
                except ValueError:
                    print("無効な入力です。数値を入力してください。")

            while True:
                try:
                    ns = int(input("   スペクトル テンスウ  ="))
                    if ns <= 0:
                        print("スペクトル点数は正の整数である必要があります。")
                    else:
                        break
                except ValueError:
                    print("無効な入力です。数値を入力してください。")

            if ns > len(X): # Resize X if needed for spectrum points
                X = [0.0] * ns

            power_spectrum_estimation_routine(min_m, p_mm, dt, wnmin, wnmax, ns, X, RFPE)
            output_routine(ns, X, n_str_mem, wnmin, wnmax)
            # After re-running, ask again
            continue # Continue the 'while True' loop for `an_again`

        elif an_again == "N":
            break
        else:
            print("YまたはNを入力してください。")

    clear_screen()
    sys.exit()

if __name__ == "__main__":
    main()