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

print("🚀 INITIATING FFT TRANSFORM (TIME -> FREQUENCY DOMAIN)...")

# 1. LOAD DATA
possible_paths = ['vG.0.1/real_tokamak_data_v2.csv', 'real_tokamak_data_v2.csv']
df = None
for path in possible_paths:
    if os.path.exists(path):
        print(f"   ✅ Found data at: {path}")
        df = pd.read_csv(path)
        break
if df is None: exit()

df.replace([np.inf, -np.inf], np.nan, inplace=True)
df.fillna(0, inplace=True)
y = df['label'].values

# 2. PERFORM FFT ON CRITICAL CHANNELS
# We focus on 'n1' (Instability) and 'ip' (Current)
# These are the signals that "wobble" before a crash.
target_signals = ['n1', 'ip']
fft_features = []
feature_names = []

for p in target_signals:
    # Extract the 100 time steps
    cols = [c for c in df.columns if c.startswith(p + '_')]
    cols.sort(key=lambda x: int(x.split('_')[1]))
    signal_data = df[cols].values # (N, 100)
    
    # Apply FFT
    # We take the absolute magnitude of the FFT
    # We ignore the first component (DC offset / Mean value)
    fft_vals = np.abs(np.fft.rfft(signal_data, axis=1))[:, 1:] 
    
    # We select 4 key frequencies per signal (Low, Mid-Low, Mid-High, High)
    # This gives us 8 total features (4 for n1, 4 for ip)
    # rfft on 100 points gives ~51 frequencies. We bin them.
    
    n_freqs = fft_vals.shape[1]
    indices = np.linspace(0, n_freqs-1, 4, dtype=int)
    
    extracted = fft_vals[:, indices]
    fft_features.append(extracted)
    
    for i in indices:
        feature_names.append(f"{p}_freq_{i}")

# Stack Features
X_fft = np.hstack(fft_features)

print(f"   FFT Shape: {X_fft.shape} (8 Features representing Plasma Rhythm)")
print(f"   Features: {feature_names}")

# 3. SAVE
np.savez('golden_qgan_data.npz', # Overwriting for the QSVC script
         X=X_fft, 
         y=y, 
         feature_names=feature_names)

print("\n💾 Saved Frequency Data to 'golden_qgan_data.npz'.")

# 4. VISUALIZE (Sanity Check)
# Let's see if Healthy vs Disruptive look different in Frequency
idx_0 = np.where(y==0)[0]
idx_1 = np.where(y==1)[0]

avg_0 = np.mean(X_fft[idx_0], axis=0)
avg_1 = np.mean(X_fft[idx_1], axis=0)

plt.figure(figsize=(10, 6))
plt.plot(avg_0, label='Healthy Spectrum', marker='o')
plt.plot(avg_1, label='Disruptive Spectrum', marker='x')
plt.title("The 'Sound' of a Plasma: Frequency Domain Analysis")
plt.xlabel("Frequency Bins (n1 then ip)")
plt.ylabel("Magnitude")
plt.legend()
plt.grid(True)
plt.savefig('fft_spectrum_debug.png')
print("📸 Saved spectrum plot to 'fft_spectrum_debug.png'")