From 6194cfcd49177f6af22ef313d6cdf3ccc53cfe7a Mon Sep 17 00:00:00 2001
From: Correl Roush <correl@gmail.com>
Date: Sat, 26 Sep 2020 00:59:57 -0400
Subject: [PATCH] Fix up the spectrum graph

---
 turntable/gui.py | 53 ++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 9 deletions(-)

diff --git a/turntable/gui.py b/turntable/gui.py
index f3791ef..871d31a 100644
--- a/turntable/gui.py
+++ b/turntable/gui.py
@@ -32,20 +32,54 @@ class Plot:
         self.width = width
         self.height = height
         self.bars = bars
-        self.audio = b""
+        self.audio = models.PCM(44100, 2)
 
-    def draw(self) -> None:
-        data = np.fromstring(self.audio, dtype=np.int16)
+    def spectrum(self) -> np.array:
+        data = np.fromstring(self.audio.raw, dtype=np.int16)
+        if len(data) == 0:
+            return data
+        merged = np.mean(np.reshape(data, (-1, self.audio.channels)), axis=1)
+        fft = np.abs(np.fft.fft(merged))
+        fft = fft[: len(fft) // 2]
+        dbfs = 20 * np.log10(fft * 2 / (len(fft) * 2 ** 15))
+        dbfs = np.maximum(-100, dbfs) + 100
+        return dbfs
+
+    def draw_lines(self) -> None:
+        data = self.spectrum()
         if len(data) == 0:
             return
-        fft = abs(np.fft.fft(data).real)
-        fft = fft[: len(fft) // 2]
-        fft = abs(scipy.signal.resample(fft, self.bars))
+
+        print((np.min(data), np.max(data)))
+        nlines = self.width // 4
+        data = np.mean(
+            np.reshape(data[: len(data) // nlines * nlines], (-1, len(data) // nlines)),
+            axis=1,
+        )
+        lines = self.height * (data / 100)
+        for i, line in enumerate(lines):
+            pygame.draw.line(
+                self.screen,
+                (128, 128, 128),
+                (self.x + i * 4, self.height),
+                (self.x + i * 4, self.height - int(line)),
+            )
+
+    def draw_bars(self) -> None:
+        data = self.spectrum()
+        if len(data) == 0:
+            return
+        fft = np.mean(
+            np.reshape(
+                data[: len(data) // self.bars * self.bars], (-1, len(data) // self.bars)
+            ),
+            axis=1,
+        )
 
         light_width = self.width // (2 * self.bars - 1)
         light_height = self.height // 2 // light_width
         light_steps = self.height // light_height // 2
-        lights = fft * light_steps / 2 ** 16
+        lights = fft * light_steps / 100
 
         colors = [
             (0, (0, 255, 0)),
@@ -147,11 +181,12 @@ def main():
                 ...
             try:
                 while sample := pcm_in.get(False):
-                    plot.audio = sample.raw
+                    plot.audio = sample
             except queue.Empty:
                 ...
             screen.fill((0, 0, 0))
-            plot.draw()
+            plot.draw_lines()
+            plot.draw_bars()
             title_text = font.render(title, True, (255, 255, 255))
             title_rect = title_text.get_rect()
             title_rect.left = 25