import numpy as np

from symop.modes.envelopes import GaussianEnvelope
from symop.modes.transfer import TimeDelay
import symop.viz as viz

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

delay = TimeDelay(tau=2.0)

out, eta = delay.apply_to_gaussian(env)

print(out.tau, eta)

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

viz.plot(
    env,
    t=t,
    w=w,
    title="Before delay",
    freq_relative=True,
    show=False,
)

viz.plot(
    out,
    t=t,
    w=w,
    title="After delay",
    freq_relative=True,
    show=False,
)