import matplotlib
matplotlib.use('TkAgg')

import numpy as np
from scipy import signal
import matplotlib.pyplot as plt


class LinearPhaseFilter:
    def __init__(self, numtaps, *args, **kwargs):
        self.numtaps = numtaps
        self.coef = signal.firwin(numtaps, *args, **kwargs)
        self.state = np.zeros(self.numtaps - 1)

    def input(self, block):
        output, self.state = signal.lfilter(self.coef, [1.0], block, zi=self.state)
        return output

    def delay(self):
        assert(self.numtaps % 2 == 1)
        return (self.numtaps - 1)//2


fs = 192000
t = np.arange(0.0, 0.15, step=1./fs)
path = 0.1e-3*np.sin(2.*np.pi*47.*t)
interferogram = np.sin(2*np.pi*path/1550e-9) + np.sin(2*np.pi*path/1270.0e-9)

filt = LinearPhaseFilter(8191, [80., 120.], pass_zero=False, fs=fs)
demod = filt.input(np.diff(interferogram)**2)

d = filt.delay() + 3000
t = t[d:]
path = path[d:]
interferogram = interferogram[d:]
demod = demod[d:]

zero_crossings = np.where(np.diff(np.sign(demod)))[0]
x = []
y = []
for i, zero in enumerate(zero_crossings):
    x.append(t[zero])
    y.append(i*np.pi)
A = np.vstack([x, np.ones(len(x))]).T
m, c = np.linalg.lstsq(A, y, rcond=None)[0]
normspeed = (1. + np.sin(m*t+c))/2.

segments = []
i = 0
try:
    while True:
        while normspeed[i] < 0.25:
            i += 1
        segment = []
        nco = 0.0
        while normspeed[i] >= 0.25:
            nco += normspeed[i]
            if nco >= 1.0:
                nco -= 1.0
                segment.append(interferogram[i])
            i += 1
        segments.append(segment)
except IndexError:
    pass



fig, axs = plt.subplots(3+len(segments))
axs[0].plot(path)
axs[1].plot(interferogram)
axs[2].plot(demod)

for i, segment in enumerate(segments):
    print(i, len(segment))
    axs[3+i].plot(np.abs(np.fft.rfft(segment*signal.blackmanharris(len(segment)))))


plt.show()