import pandas as pd
import seaborn as sns

from nnfwtbn import Variable, Process, Cut, roc
from nnfwtbn import toydata

df = toydata.get()
df['noise'] = df.fpid + 3 * (df.random - 0.5)**3

p_bkg = Process(r"Background", range=(0, 0))
p_sig = Process(r"Signal", range=(1, 1))

import matplotlib.pyplot as plt
fig, ax = plt.subplots()

v_higgs_m = Variable(r"$m^H$", "higgs_m", "GeV")
roc(df.compute(), p_sig, p_bkg, v_higgs_m, axes=ax, steps=400, atlas=False)
roc(df.compute(), p_sig, p_bkg, Variable("Noise ID", "noise"), axes=ax, enlarge=1.5)
None