import numpy as np

from symop.modes.envelopes import GaussianEnvelope
from symop.modes.transfer import GaussianBandpass
from symop.modes.transfer.apply import apply_transfer
import symop.viz as viz

env = GaussianEnvelope(
    omega0=10.0,
    sigma=2.0,
    tau=0.0,
    phi0=0.0,
)

filt = GaussianBandpass(
    w0=10.0,
    sigma_w=0.1,
)

out, eta = apply_transfer(filt, env)

t = np.linspace(-30.0, 30.0, 2000)
w = np.linspace(8.0, 12.0, 2000)

viz.plot(
    env,
    t=t,
    w=w,
    title="Input envelope",
    normalize_spectrum=False,
    freq_relative=True,
    show=False,
)

viz.plot(
    out,
    t=t,
    w=w,
    title=f"Filtered envelope, eta={eta:.6f}",
    normalize_spectrum=False,
    freq_relative=True,
    show=False,
)