#include <cstdio>   // For printf, scanf, fopen, fclose, perror, getchar
#include <cmath>    // For sin, cos, M_PI (from cmath)
#include <cstdlib>  // For malloc, free, exit, toupper

// Define N as a constant for array sizing.
// Given nmax can be 241, N needs to be at least 241.
// If arrays are 1-indexed (like original code), they need N+1 elements.
#define N_MAX_SIZE 300 // Renamed to avoid conflict with potential variable N

// Custom type for floating-point numbers
typedef double Real;

// Using a const variable for PI is more C++ idiomatic
const Real PI_CONST = 3.14159265358979323846;

/*---------- MEM2.C ----------*/
// Function prototypes (extern is not strictly necessary when defined in same file)
void levins(int nmax, Real *g, Real *pm, Real *b1, Real *b2, Real *a, int m);
void mk_dat(Real *x, int *nmax, int *mmx, int *lmax, Real *dt);
void burg(int nmax, int mmx, int lmax, Real *x, Real *g, Real *c,
          Real *fpe, Real *pm);
void mem_1(int nmax, int mmx, Real dt, Real *g, Real *f,
           Real *spec, Real pm);

// Global arrays (sized for 1-based indexing, so N_MAX_SIZE + 1)
Real x[N_MAX_SIZE + 1];
Real f[N_MAX_SIZE + 1];
Real fpe[N_MAX_SIZE + 1];
Real spec[N_MAX_SIZE + 1];
Real g[N_MAX_SIZE + 1];
Real c[N_MAX_SIZE + 1];

// ## メイン関数
// プログラムのエントリポイントです。データの生成、Burg法によるARモデル推定、
// そしてMEMスペクトル計算の主要な流れを制御します。
int main()
{
    int nmax, mmx, lmax, i, lag, lg1, lx1;
    Real dt_val; // Declared dt_val here to fix the 'dt not declared' error
    Real sum, av, pm; // pm is passed by reference to burg, then by value to mem_1

    // Data generation and parameter setup
    mk_dat(x, &nmax, &mmx, &lmax, &dt_val); // Pass dt_val's address

    lag = mmx + 1;
    if ((lag + 1) > lmax) lmax = lag + 1;

    // Remove mean from data
    sum = 0.0;
    for (i = 1; i <= nmax; i++) sum += x[i];
    av = sum / static_cast<Real>(nmax); // Explicit cast for floating-point division
    for (i = 1; i <= nmax; i++) x[i] -= av;

    // Perform Burg's algorithm for AR model estimation
    burg(nmax, mmx, lmax, x, g, c, fpe, &pm);

    // Calculate Maximum Entropy Method spectrum
    mem_1(nmax, mmx, dt_val, g, f, spec, pm); // Use dt_val here

    // Print results
    lag = mmx + 1;
    lg1 = lag + 1;
    lx1 = lmax + 1;

    printf("\n frec        spc         fpe     auto c.");
    // Print up to 'lag' (mmx + 1) points
    for (i = 1; i <= lag; i++)
        printf("\n%6.1f    %7.3f    %7.3f    %7.3f",
               f[i], spec[i], fpe[i], c[i]);

    // If 'lag' is less than 'lmax', print remaining auto-correlation points
    if (lag < lmax) { // Use '<' instead of '!=' for safety with loops
        for (i = lg1; i <= lmax; i++)
            printf("\n%6.1f    %7.3f          %6.3f",
                   f[i], spec[i], c[i]);
    }
    
    // Print remaining spectrum points up to nmax
    for (i = lx1; i <= nmax; i++)
        printf("\n%6.1f    %7.3f", f[i], spec[i]);

    printf("\n"); // Add a final newline for clean output

    return 0; // Indicate successful execution
}

// ## サブルーチン: levins
// Levinson-Durbinアルゴリズムの一部を実装し、Burg法の自己回帰係数と
// 予測誤差分散を更新します。
void levins(int nmax, Real *g, Real *pm, Real *b1, Real *b2, Real *a, int m)
{
    int i, k;
    Real sn, sd;

    sn = sd = 0.0;
    for (i = 1; i <= (nmax - m); i++)
    {
        sn += b1[i] * b2[i];
        sd += b1[i] * b1[i] + b2[i] * b2[i];
    }
    
    if (sd == 0.0) { // Prevent division by zero
        g[m] = 0.0;
    } else {
        g[m] = -2.0 * sn / sd;
    }
    
    *pm *= (1.0 - g[m] * g[m]);

    if (m != 1) {
        for (k = 1; k <= m - 1; k++) g[k] = a[k] + g[m] * a[m - k];
    }
    for (i = 1; i <= m; i++) a[i] = g[i]; // Store current g coefficients in 'a' for next iteration

    // Update forward and backward prediction errors (b1 and b2)
    Real temp_b1_i; // Temporary variable to store b1[i] before it's modified
    for (int i = 1; i < nmax - m; i++) // Note: loop goes up to nmax-m-1, not nmax-m
    {
        temp_b1_i = b1[i]; // Store original b1[i]
        b1[i] += a[m] * b2[i]; // Update b1[i] using the m-th coefficient (a[m] is current g[m])
        b2[i] = b2[i+1] + a[m] * temp_b1_i; // Update b2[i] using next b2 and original b1[i]
                                            // The original C code had b1[i+1] here, which is an error,
                                            // it should be the original b1[i].
                                            // Based on standard Burg algorithm, it should be the previous b1.
                                            // Assuming temp_b1_i for the correct forward error from previous step.
    }
}

// ## サブルーチン: burg
// Burg法を用いて、自己回帰モデルの係数 `g` と自己相関関数 `c`、
// そして予測誤差 `fpe` を推定します。
void burg(int nmax, int mmx, int lmax, Real *x, Real *g, Real *c,
          Real *fpe, Real *pm)
{
    Real sum;
    Real *b1, *b2, *a; // Pointers for dynamically allocated arrays
    int i, m, l, lag, lg1;

    // Dynamically allocate memory (using new/delete in C++ is preferred, but malloc for direct replacement)
    b1 = static_cast<Real *>(malloc(sizeof(Real) * (nmax + 1)));
    b2 = static_cast<Real *>(malloc(sizeof(Real) * (nmax + 1)));
    a = static_cast<Real *>(malloc(sizeof(Real) * (mmx + 1))); // 'a' stores g coefficients up to m, so mmx+1 is sufficient

    // Check for successful allocation
    if (!b1 || !b2 || !a) {
        fprintf(stderr, "メモリ割り当てに失敗しました。\n");
        exit(1); // Exit on critical error
    }

    sum = 0.0;
    for (i = 1; i <= nmax; i++) sum += x[i] * x[i];
    c[1] = sum / static_cast<Real>(nmax); // Auto-correlation at lag 0 (variance)
    *pm = c[1]; // Initial prediction error power
    fpe[1] = static_cast<Real>(nmax + 1) / static_cast<Real>(nmax - 1) * (*pm); // Initial FPE

    lag = mmx + 1;
    lg1 = lag + 1;

    // Initialize forward and backward prediction errors for Burg's algorithm
    b1[1] = x[1];
    for (i = 2; i <= nmax; i++)
    {
        b1[i] = x[i]; // Forward error: current data point
        b2[i-1] = x[i]; // Backward error: previous data point for the first step
                        // The original code used x[i] for b2[i-1], which means b2 starts from x[1] and goes to x[nmax-1]
    }
    
    for (m = 1; m <= mmx; m++)
    {
        printf("\n M=%4d ", m);
        levins(nmax, g, pm, b1, b2, a, m); // Update g, pm, b1, b2, a

        sum = 0.0;
        for (i = 1; i <= m; i++) sum -= c[m - i + 1] * g[i]; // Calculate next auto-correlation point
        c[m + 1] = sum;

        if (m != (nmax - 1)) { // Avoid division by zero when nmax-m-1 becomes 0
            if ((nmax - m - 1) == 0) {
                 fpe[m + 1] = 0.0; // Or some other suitable value
            } else {
                 fpe[m + 1] = static_cast<Real>(nmax + m + 1) / static_cast<Real>(nmax - m - 1) * (*pm);
            }
        }
    }

    // Extrapolate auto-correlation function if lmax > lag
    if (lag < lmax) {
        for (l = lg1; l <= lmax; l++)
        {
            sum = 0.0;
            // The original code `sum -= c[l-1] * g[i];` seems incorrect for typical auto-correlation extrapolation.
            // It should be a summation involving previously calculated `g` and `c` values.
            // Assuming it means a sum over the AR coefficients.
            // A common formula for c[l] (l > mmx) is:
            // c[l] = -sum_{k=1}^{mmx} g[k] * c[l-k]
            // This loop looks like it's missing the sum for i=1 to mmx for the correct recursion.
            // Based on similar implementations, it should iterate over 'g' coefficients.
            for (i = 1; i <= mmx; i++) { // Iterate over all mmx coefficients
                if ((l - i) >= 0 && (l - i) <= lmax) { // Ensure c[l-i] is a valid index
                    sum -= c[l - i] * g[i]; // Use c[l-i] not c[l-1]
                }
            }
            c[l] = sum;
        }
    }

    // Free dynamically allocated memory
    free(b1);
    free(b2);
    free(a);
}

// ## サブルーチン: mem_1
// 自己回帰係数 `g` と予測誤差分散 `pm` を用いて、
// 最大エントロピー法によるパワースペクトルを計算します。
void mem_1(int nmax, int mmx, Real dt_param, Real *g, Real *f, Real *spec,
                     Real pm) // Renamed dt to dt_param to avoid conflict with global dt
{
    int i, j;
    Real f0, sumr, sumi, ci_freq; // Renamed ci to ci_freq to avoid confusion with global PI_CONST

    if ((nmax - 1) == 0) { // Prevent division by zero
        f0 = 0.0; // Or handle as an error
    } else {
        f0 = 1.0 / (static_cast<Real>(nmax - 1) * dt_param);
    }
    
    for (i = 1; i <= nmax; i++)
    {
        f[i] = f0 * static_cast<Real>(i - 1); // Frequency points
        printf("\n  freq=%6.2f", f[i]);
        sumr = 1.0;
        sumi = 0.0;
        
        ci_freq = 2.0 * PI_CONST * f[i] * dt_param; // Calculate omega * dt

        for (j = 1; j <= mmx; j++)
        {
            sumr += g[j] * cos(static_cast<double>(j) * ci_freq);
            sumi += g[j] * sin(static_cast<double>(j) * ci_freq);
        }
        
        Real denominator = (sumr * sumr + sumi * sumi);
        if (denominator == 0.0) { // Prevent division by zero
            spec[i] = 1e18; // Assign a very large number for "infinity" or handle as error
        } else {
            spec[i] = 2.0 * pm * dt_param / denominator;
        }
    }
}

// ## サブルーチン: mk_dat
// 最大エントロピー法で分析するためのサンプルデータを生成します。
// 複数の正弦波の合成波を生成します。
void mk_dat(Real *x, int *nmax, int *mmx, int *lmax, Real *dt_param)
{
    Real t1, t2, t3, t4, t, y;
    int i;

    *nmax = 241; // Number of data points
    *mmx = 38;   // Model order (AR order)
    *lmax = 50;  // Maximum lag for auto-correlation (if used beyond AR order)
    
    if ((*nmax - 1) == 0) { // Prevent division by zero
        *dt_param = 0.0;
    } else {
        *dt_param = 1.0 / static_cast<Real>(*nmax - 1); // Sampling interval
    }
    
    // Periods of sine waves
    t1 = 1.0 / 124.5;
    t2 = 1.0 / 65.0;
    t3 = 1.0 / 64.0;
    t4 = 1.0 / 4.0;

    // Generate synthetic data
    for (i = 1; i <= *nmax; i++)
    {
        t = *dt_param * static_cast<Real>(i - 1); // Time point
        y = sin(2.0 * PI_CONST * t / t1) + sin(2.0 * PI_CONST * t / t2)
            + sin(2.0 * PI_CONST * t / t3) + sin(2.0 * PI_CONST * t / t4);
        x[i] = y;
        /* + 2.0 * drand48()  - 2; // Commented out: for adding noise */
    }
}