from matplotlib import pyplot as plt
from matplotlib import patches as pat
import numpy as np
import reader

# load data
data_reader = reader.Reader()
num = 14
data = data_reader.read_file(f".\data\T00{num}ALL.csv")
str_end = 'Near end' if num % 2 == 0 else 'Far end'
str_desc = '230 Ohm'

# FFT
np_data1 = np.array(data['CH1'])
np_data2 = np.array(data['CH2'])
np_time = np.array(data['TIME'])
X1 = np.fft.rfft(np_data1)
X2 = np.fft.rfft(np_data2)
N = np.fft.rfftfreq(len(np_data1), (np_time[1]-np_time[0]))

phi1 = np.angle(X1, True)
phi2 = np.angle(X2, True)
diff_phi = phi1[np.argmax(abs(X1))]-phi2[np.argmax(abs(X2))]

fig, ax = plt.subplots()
plt.plot(data['TIME'], data['CH1'], label='CH1')
plt.plot(data['TIME'], data['CH2'], label='CH2')

handles, labels = ax.get_legend_handles_labels()
patch = pat.Patch(
    color='C0', 
    label='pk-pk = ${:.2f}\ V$'.format(max(data['CH1'])-min(data['CH1'])))
handles.append(patch)
patch = pat.Patch(
    color='C1', 
    label='pk-pk = ${:.2f}\ V$'.format(max(data['CH2'])-min(data['CH2'])))
handles.append(patch)
patch = pat.Patch(
    color='C0', 
    label='F = ${:.1f}\ MHz$'.format(N[np.argmax(abs(X1))]/1e6))
handles.append(patch)
patch = pat.Patch(
    color='C1', 
    label='F = ${:.1f}\ MHz$'.format(N[np.argmax(abs(X2))]/1e6))
handles.append(patch)
patch = pat.Patch(
    color='grey', 
    label='$\Delta \phi$ = ${:.1f} \degree$'.format(diff_phi))
handles.append(patch)

plt.title('Voltage vs Time, ' + str_desc + ', ' +  str_end)
plt.legend(loc='upper right', handles=handles)
plt.grid()
plt.xlabel('Time [s]')
plt.ylabel('Voltage [V]')
plt.xlim(data['TIME'][0], data['TIME'][-1])
name = f'Graph - {str_desc} - {str_end}.png'
plt.savefig(f'./doc/crosstalk/img/{name}')
print('Saved:', name)
# plt.show()