# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt

# --- 1. Parameters from the Arduino/K230 Test Code ---
FFT_POINTS = 64
SAMPLE_RATE = 44100
PI = np.pi

# --- 2. Replicate the Test Signal ---
def build_test_signal(fft_num):
    """Generates the short-integer test signal used in the K230 C++ code."""
    rx = np.zeros(fft_num, dtype=np.float32)
    for i in range(fft_num):
        # Cosine amplitudes: 10, 20, 30, 0.2, 1000
        tempf1 = 10.0 * np.cos(1.0 * 2.0 * PI * i / fft_num)
        tempf2 = 20.0 * np.cos(2.0 * 2.0 * PI * i / fft_num)
        tempf3 = 30.0 * np.cos(3.0 * 2.0 * PI * i / fft_num)
        tempf4 = 0.2 * np.cos(4.0 * 2.0 * PI * i / fft_num)
        tempf5 = 1000.0 * np.cos(5.0 * 2.0 * PI * i / fft_num)
        rx[i] = tempf1 + tempf2 + tempf3 + tempf4 + tempf5
    # The K230 input is static_cast<short>, which limits the range.
    return rx.astype(np.float32) 

time_domain_signal = build_test_signal(FFT_POINTS)

# --- 3. Perform FFT and Extract K230 Equivalent Data ---
fft_result = np.fft.fft(time_domain_signal)

# Amplitudes (Y-axis): K230 likely uses the raw complex magnitude (abs)
# The C++ code uses f.calculate_amplitudes(), which returns a large int32_t magnitude.
amplitudes_full = np.abs(fft_result)
amplitudes_half = amplitudes_full[:FFT_POINTS // 2] 

# Frequencies (X-axis): Frequencies in Hz, as calculated by K230::calculate_frequencies
freqs_hz = np.fft.fftfreq(FFT_POINTS, d=1/SAMPLE_RATE)
plot_freqs_hz = freqs_hz[:FFT_POINTS // 2] 

# --- 4. Apply C++ BMP Plotting Logic for X-axis Scaling ---

# C++ Logic: Find the highest frequency *present* in the plot data (Bin 5).
# max_freq_data = *std::max_element(freqs.begin(), freqs.end());
max_freq_data = np.max(plot_freqs_hz)
if max_freq_data == 0:
    max_freq_data = 1.0 # Avoid division by zero

# C++ Logic: Scale all frequencies to the range [0, 1]
# x_normalized = freqs[i] / max_freq_data
x_normalized = plot_freqs_hz / max_freq_data


# C++ Logic: Find the maximum amplitude for Y-axis scaling.
max_amplitude = np.max(amplitudes_half)
if max_amplitude == 0:
    max_amplitude = 1.0

# C++ Logic: Scale all amplitudes to the range [0, 1]
y_normalized = amplitudes_half / max_amplitude


# --- 5. Plot the Result (Normalized to match C++ Pixel Drawing) ---
plt.figure(figsize=(10, 6.25)) # Slightly adjusted for standard screen viewing

# Plot the normalized data. X-axis now goes from 0.0 to 1.0.
# The resulting plot SHAPE is identical to the K230 BMP output.
plt.plot(x_normalized, y_normalized, 
         label=f'K230 Emulated Spectrum (Normalized)', 
         color='blue', linewidth=1)

# Highlight the dominant peak (Bin 5)
peak_index = 5
f_expected_hz = plot_freqs_hz[peak_index]
x_peak_normalized = x_normalized[peak_index]
y_peak_normalized = y_normalized[peak_index]

plt.scatter(x_peak_normalized, y_peak_normalized, color='red', marker='o', s=50, 
            label=f'Dominant Peak: {f_expected_hz:.2f} Hz (Normalized X={x_peak_normalized:.2f})')


# Add Chart Title and Labels
plt.title(f'FFT Spectrum Comparison: K230 BMP Visual Style (N={FFT_POINTS})', fontsize=14)
plt.xlabel('Normalized Frequency (0.0 = DC, 1.0 = Max Observed Freq)', fontsize=12)
plt.ylabel('Normalized Amplitude (0.0 = Min, 1.0 = Max Peak)', fontsize=12)

# Set the X-axis limit explicitly to 0.0 to 1.0
plt.xlim(0, 1.05) 
plt.ylim(0, 1.05) 

# Optional: Add helper ticks to show the location of the first few bins
bin_freqs = [plot_freqs_hz[i] for i in range(1, 6)]
bin_ticks = [freq / max_freq_data for freq in bin_freqs]

plt.xticks(bin_ticks, [f'{i+1} Bin\n({f:.1f}Hz)' for i, f in enumerate(bin_freqs)], rotation=45, ha='right')


# Add grid lines and legend
plt.grid(True, which='both', linestyle='--', alpha=0.7)
plt.legend(fontsize=10)
plt.tight_layout()
plt.show()