import numpy as np
import pandas as pd
from modules.ndxplorer.ndxplorer.data_source import compute_values

df = pd.DataFrame({
    'S prompt green (kHz) | 0-2048': [1.0, 2.0, 3.0],
    'Bg': [0.1, 0.2, 0.3],
})

equations = [
    {'Sg(PIE)': "'S prompt green (kHz)'"},
    {'Fg(PIE)': "'Sg(PIE)' - 'Bg'"},
]

compute_values(df, {}, equations)

src = df['S prompt green (kHz) | 0-2048'].values
sg = df['Sg(PIE)'].values
fg = df['Fg(PIE)'].values

assert np.allclose(sg, src)
assert np.allclose(fg, src - df['Bg'].values)
print('OK _tmp_test')
