import numpy as np
from scipy import signal
from numpy.fft import fft, ifft
import matplotlib.pyplot as plt


# Filtering things I have mainly stolen from the internet. I reallt need to learn more so I can do something compitent.
class SignalProcess:
    def __init__(self, path):
        self.data = np.loadtxt(path, delimiter="\t", dtype=np.float64)
        self.times = self.data[:, 0]
        self.counts = self.data[:, 1]
        self.ave_counts = self.data[:, 2]

    def band_pass(self):
        fs = 30.0
        lowcut = 1.5 * 1.0 / 30.0
        highcut = 2.5 * 1.0 / 30.0
        nyqs = 0.5 * fs
        low = lowcut / nyqs
        high = highcut / nyqs
        b, a = signal.butter(2, [low, high], "bandpass", analog=False)
        filtered_counts = signal.filtfilt(b, a, self.ave_counts, axis=0, padlen=150)
        return filtered_counts

    def butter_filter(self):

        b, a = signal.butter(2, 0.02)
        filtered_counts = signal.filtfilt(b, a, self.ave_counts, padlen=150)
        return filtered_counts

    def fft(self):
        sr = 30
        ts = 1.0 / sr
        X = fft(self.counts)
        N = len(X)
        n = np.arange(N)
        T = N / sr
        freq = n / T
        plt.stem(freq, np.abs(X), "b", markerfmt=" ", basefmt="-b")
        plt.xlabel("Freq (Hz)")
        plt.ylabel("FFT Amplitude |X(freq)|")
        plt.xlim(0, 10)
        plt.show()