from scipy import signal
from nmigen import *
from nmigen.back.pysim import *
from nmigen.cli import main
import matplotlib.pyplot as plt
import numpy as np
import wave, struct

class Crossover(Elaboratable):
    def __init__(self, cutoff, sampleRate = 48000):
        self.input = Signal(32)
        self.output = Signal(32)
        self.sampleRate = sampleRate
        self.cutoff = cutoff
        (self.num, self.den) = signal.butter(2, [cutoff], btype='lowpass', analog=False, fs=sampleRate)

    def polynome(self):
        return (self.num, self.den)

    def plot(self):
        w, h = signal.freqz(self.num, self.den, fs=self.sampleRate)
        plt.plot(w, 20 * np.log10(abs(h)))
        plt.xscale('log')
        plt.title('Butterworth filter frequency response')
        plt.xlabel('Frequency [radians / second]')
        plt.ylabel('Amplitude [dB]')
        plt.margins(0, 0.1)
        plt.grid(which='both', axis='both')
        plt.axvline(self.cutoff, color='green') # cutoff frequency
        plt.show()

    def elaborate(self, platform):
        m = Module()

        x_sigs = []
        for i in range(len(self.num)):
            x_sigs.append(Signal(32))

        m.d.comb += x_sigs[0].eq(self.input)

        for i in range(1, len(self.num)):
            m.d.sync += x_sigs[i].eq(x_sigs[i-1])

        y_sigs = []
        for i in range(len(self.den)):
            y_sigs.append(Signal(32))

        m.d.comb += y_sigs[0].eq(self.output)

        for i in range(1, len(self.den)):
            m.d.sync += y_sigs[i].eq(y_sigs[i-1])

        outputSum = 0
        for i in range(len(self.num)):
            outputSum += x_sigs[i]*int(self.num[i]*256)
        for i in range(1, len(self.den)):
            outputSum -= y_sigs[i]*int(self.den[i]*256)

        buf = Signal(32)
        m.d.comb += buf.eq(outputSum)

        m.d.comb += self.output.eq(outputSum >> 8)

        return m

if __name__ == "__main__":
    crossover = Crossover(1000, 48000)
    #crossover.plot()
    #print(crossover.polynome())

    with Simulator(crossover) as sim:
        def process():
            audioIn = wave.open("input.wav", "r")
            audioOut = wave.open("output.wav", "w")
            audioOut.setnchannels(1)
            audioOut.setsampwidth(2)
            audioOut.setframerate(48000)
            print("Frames ", audioIn.getnframes())
            for i in range(audioIn.getnframes()):
                print("Sample ", i)
                sample = struct.unpack("<h", audioIn.readframes(1))
                yield crossover.input.eq(sample[0]+32768)
                yield Tick()
                #print((yield crossover.output)-32768)
                audioOut.writeframes(struct.pack("<h", (yield crossover.output)-32768))
            audioOut.close()

        sim.add_clock(1/48000)
        sim.add_process(process)
        sim.run()