import numpy as np
import matplotlib.pyplot as plt
from dtk.process import find_timeshift, truncate_data

t = np.linspace(0.0, 4.0, num=401)
sig1 = np.sin(2.0*t) + np.random.normal(0.0, 0.1, size=len(t))
sig2 = np.sin(2.0*t + 0.3) + np.random.normal(0.0, 0.1, size=len(t))

tau = find_timeshift(sig1, sig2, 100, guess=0.2)

sigtr1, sigtr2 = truncate_data(tau, sig1, sig2, 100)

fig, ax = plt.subplots(layout='constrained')
ax.plot(t[:len(sigtr1)], sigtr1, t[:len(sigtr2)], sigtr2)
ax.legend(['Signal 1', 'Signal 2'])
ax.set_title('Shift: {:1.3f} s'.format(tau))
ax.set_xlabel('Time [s]')