stfft.go

package main

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"log"
	"math"
	"sync"

	"github.com/gordonklaus/portaudio"
	"github.com/mjibson/go-dsp/spectral"
	"github.com/rakyll/portmidi"
	"golang.org/x/exp/shiny/driver"
	"golang.org/x/exp/shiny/screen"
	"golang.org/x/mobile/event/lifecycle"
	"golang.org/x/mobile/event/paint"
)

const sampleRate = 44100
const fftSize = 1024
const windowSize = 256
const bufferSize = fftSize
const squareWidth = 100
const smoothing = .2
const minDB = 60
const maxDB = 200

var (
	valsMu sync.Mutex
	vals   []float64

	buckets = []struct {
		Min, Max int
	}{
		{
			Min: freq2Bin(0), Max: freq2Bin(150),
		},
		{
			Min: freq2Bin(2500), Max: freq2Bin(3000),
		},
		{
			Min: freq2Bin(9000), Max: freq2Bin(11000),
		},
	}
)

func main() {
	if err := portmidi.Initialize(); err != nil {
		log.Fatal(err)
	}
	if err := portaudio.Initialize(); err != nil {
		log.Fatal(err)
	}

	go analyze()

	driver.Main(func(s screen.Screen) {
		w, err := s.NewWindow(&screen.NewWindowOptions{
			Width:  squareWidth * len(buckets),
			Height: squareWidth,
			Title:  "Viz",
		})
		check(err)

		var paintVals []float64

		for {
			switch e := w.NextEvent().(type) {
			case lifecycle.Event:
				if e.To == lifecycle.StageDead {
					return
				}
			case paint.Event:
				valsMu.Lock()
				copy(paintVals, vals)
				valsMu.Unlock()
				for i := range vals {
					v := uint8((1 - vals[i]) * 255)
					w.Fill(image.Rectangle{
						Min: image.Point{squareWidth * i, 0},
						Max: image.Point{squareWidth + squareWidth*i, squareWidth},
					}, color.RGBA{v, v, v, 1}, draw.Src)
				}
				w.Publish()

				w.Send(e)
			case error:
				log.Print(e)
			}
		}
	})
}

func analyze() {
	inputDevice, _ := portaudio.DefaultInputDevice()
	log.Print("in from ", inputDevice.Name)

	midiID := portmidi.DefaultOutputDeviceID()
	midiOut, err := portmidi.NewOutputStream(midiID, 128, 0)
	log.Print("out to ", portmidi.Info(midiID).Name)
	check(err)

	fmt.Println("actual buckets:")
	for i, b := range buckets {
		fmt.Println(i, bin2Freq(b.Min), bin2Freq(b.Max))
	}

	thresholds := make([]float64, len(buckets))
	tmpVals := make([]float64, len(buckets)) // temp buffer

	valsMu.Lock()
	vals = make([]float64, len(buckets))
	valsMu.Unlock()

	p := portaudio.LowLatencyParameters(inputDevice, nil)
	p.Input.Channels = 1
	p.Output.Channels = 0
	p.SampleRate = sampleRate
	p.FramesPerBuffer = windowSize

	var once sync.Once

	in := make([]int32, windowSize)
	inF := make([]float64, bufferSize)

	stream, err := portaudio.OpenStream(p, in)
	//stream, err := portaudio.OpenDefaultStream(1, 0, sampleRate, fftSize, in)
	check(err)
	check(stream.Start())
	for {
		//start := time.Now()
		check(stream.Read())
		//log.Print("Read ", time.Now().Sub(start))
		inF = inF[len(in):]
		for i := range in {
			inF = append(inF, float64(in[i]))
		}

		pxx, freqs := spectral.Pwelch(inF, float64(sampleRate), &spectral.PwelchOptions{
			//Noverlap: windowSize,
			NFFT: fftSize,
		})

		once.Do(func() {
			log.Print("actual actual freqs")
			for _, bucket := range buckets {
				log.Print(freqs[bucket.Min], freqs[bucket.Max])
			}
		})

		valsMu.Lock()
		for i, bucket := range buckets {
			min := bucket.Min
			max := bucket.Max
			tmpVals[i] = 0
			for n := min; n < max; n++ {
				tmpVals[i] += lerp(minDB, maxDB, math.Log10(pxx[n])*10) // db min/max
			}
			tmpVals[i] /= float64(max - min)
			if tmpVals[i] > thresholds[i] {
				thresholds[i] = tmpVals[i]

				// on beat
				vals[i] = tmpVals[i]
			} else {
				//vals[i] = vals[i] * .5 // decay 2048
				//vals[i] = vals[i] * .7 // decay 1024
				//vals[i] = vals[i] * .85 // decay 512
				vals[i] = vals[i] * .92 // decay 256
			}

			bendVal := int64(vals[i] * (1 << 14))
			bendHi := (bendVal >> 7) & 0x7f
			bendLo := bendVal & 0x7f

			midiOut.WriteShort(0xE0|int64(i), bendLo, bendHi) // Bend lo hi (ch n)
		}
		valsMu.Unlock()

		for i := range thresholds {
			//thresholds[i] *= .995 // decay 2048
			//thresholds[i] *= .997 // decay 1024
			//thresholds[i] *= .9985 // decay 512
			thresholds[i] *= .9993 // decay 256
		}
		//log.Print(time.Now().Sub(start))
	}
}

func check(err error) {
	if err != nil {
		panic(err)
	}
}

func lerp(min, max float64, val float64) float64 {
	if val < min {
		return 0
	}
	if val > max {
		return 1
	}
	return (val - min) / (max - min)
}

func freq2Bin(freq float64) int {
	return int(freq / (sampleRate / fftSize))
}

func bin2Freq(bin int) float64 {
	return float64(bin) * (sampleRate / fftSize)
}