package tracker

import (
	"math"
	"math/cmplx"

	"github.com/viterin/vek/vek32"
	"github.com/vsariola/sointu"
)

type (
	SpecAnalyzer struct {
		settings SpecAnSettings
		broker   *Broker
		chunker  chunker
		temp     specTemp
	}

	SpecAnSettings struct {
		ChnMode    SpecChnMode
		Smooth     SpecSmoothing
		Resolution int
	}

	SpecChnMode   int
	SpecSmoothing int
	Spectrum      [2][]float32

	specTemp struct {
		power      [2][]float32
		window     []float32    // window weighting function
		normFactor float32      // normalization factor, to account for the windowing
		bitPerm    []int        // bit-reversal permutation table
		tmpC       []complex128 // temporary buffer for FFT
		tmp1, tmp2 []float32    // temporary buffers for processing
	}
)

const (
	SpecResolutionMin = 7
	SpecResolutionMax = 16
)

const (
	SpecChnModeOff      SpecChnMode = iota // no spectrum analysis is done to save CPU resources
	SpecChnModeCombine                     // calculate a single combined spectrum for both channels
	SpecChnModeSeparate                    // calculate separate spectrums for left and right channels
	SpecChnModeLeft                        // calculate spectrum only for the left channel
	SpecChnModeRight                       // calculate spectrum only for the right channel
	NumSpecChnModes
)

const (
	SpecSmoothingSlow SpecSmoothing = iota
	SpecSmoothingMedium
	SpecSmoothingFast
	NumSpecSmoothing
)

var spectrumSmoothingMap map[SpecSmoothing]float32 = map[SpecSmoothing]float32{
	SpecSmoothingSlow:   0.05,
	SpecSmoothingMedium: 0.2,
	SpecSmoothingFast:   1.0,
}

func NewSpecAnalyzer(broker *Broker) *SpecAnalyzer {
	ret := &SpecAnalyzer{broker: broker}
	ret.init(SpecAnSettings{
		ChnMode:    SpecChnModeCombine,
		Smooth:     SpecSmoothingMedium,
		Resolution: 10,
	})
	return ret
}

func (s *SpecAnalyzer) Run() {
	for {
		select {
		case <-s.broker.CloseSpecAn:
			close(s.broker.FinishedSpecAn)
			return
		case msg := <-s.broker.ToSpecAn:
			s.handleMsg(msg)
		}
	}
}

func (s *SpecAnalyzer) handleMsg(msg MsgToSpecAn) {
	if msg.HasSettings {
		s.init(msg.SpecSettings)
	}
	switch m := msg.Data.(type) {
	case *sointu.AudioBuffer:
		if s.settings.ChnMode != SpecChnModeOff {
			buf := *m
			l := len(s.temp.window)
			// 50% overlap with the windows
			s.chunker.Process(buf, l, l>>1, func(chunk sointu.AudioBuffer) {
				TrySend(s.broker.ToModel, MsgToModel{Data: s.update(chunk)})
			})
		}
		s.broker.PutAudioBuffer(m)
	default:
		// unknown message type; ignore
	}
}

func (a *SpecAnalyzer) init(s SpecAnSettings) {
	s.Resolution = min(max(s.Resolution, SpecResolutionMin), SpecResolutionMax)
	a.settings = s
	n := 1 << s.Resolution
	a.temp = specTemp{
		power:   [2][]float32{make([]float32, n/2), make([]float32, n/2)},
		window:  make([]float32, n),
		bitPerm: make([]int, n),
		tmpC:    make([]complex128, n),
		tmp1:    make([]float32, n),
		tmp2:    make([]float32, n),
	}
	for i := range n {
		// Hanning window
		w := float32(0.5 * (1 - math.Cos(2*math.Pi*float64(i)/float64(n-1))))
		a.temp.window[i] = w
		a.temp.normFactor += w
		// initialize the bit-reversal permutation table
		a.temp.bitPerm[i] = i
	}
	// compute the bit-reversal permutation
	for i, j := 1, 0; i < n; i++ {
		bit := n >> 1
		for ; j&bit != 0; bit >>= 1 {
			j ^= bit
		}
		j ^= bit

		if i < j {
			a.temp.bitPerm[i], a.temp.bitPerm[j] = a.temp.bitPerm[j], a.temp.bitPerm[i]
		}
	}
}

func (s *SpecAnalyzer) update(buf sointu.AudioBuffer) *Spectrum {
	ret := s.broker.GetSpectrum()
	switch s.settings.ChnMode {
	case SpecChnModeLeft:
		s.process(buf, 0)
		ret[0] = append(ret[0], s.temp.power[0]...)
	case SpecChnModeRight:
		s.process(buf, 1)
		ret[1] = append(ret[1], s.temp.power[1]...)
	case SpecChnModeSeparate:
		s.process(buf, 0)
		s.process(buf, 1)
		ret[0] = append(ret[0], s.temp.power[0]...)
		ret[1] = append(ret[1], s.temp.power[1]...)
	case SpecChnModeCombine:
		s.process(buf, 0)
		s.process(buf, 1)
		ret[0] = append(ret[0], s.temp.power[0]...)
		vek32.Add_Inplace(ret[0], s.temp.power[1])
		vek32.MulNumber_Inplace(ret[0], 0.5)
	}
	// convert to decibels
	for c := range 2 {
		vek32.MaximumNumber_Inplace(ret[c], 1e-8)
		vek32.MinimumNumber_Inplace(ret[c], 1e8)
		vek32.Log10_Inplace(ret[c])
		vek32.MulNumber_Inplace(ret[c], 10)
	}
	return ret
}

func (sd *SpecAnalyzer) process(buf sointu.AudioBuffer, channel int) {
	for i := range buf { // de-interleave
		sd.temp.tmp1[i] = buf[i][channel]
	}
	vek32.Mul_Inplace(sd.temp.tmp1, sd.temp.window)                // apply windowing
	vek32.Gather_Into(sd.temp.tmp2, sd.temp.tmp1, sd.temp.bitPerm) // bit-reversal permutation
	// convert into complex numbers
	c := sd.temp.tmpC
	for i := range c {
		c[i] = complex(float64(sd.temp.tmp2[i]), 0)
	}
	// FFT
	n := len(c)
	for len := 2; len <= n; len <<= 1 {
		ang := 2 * math.Pi / float64(len)
		wlen := complex(math.Cos(ang), math.Sin(ang))
		for i := 0; i < n; i += len {
			w := complex(1, 0)
			for j := 0; j < len/2; j++ {
				u := c[i+j]
				v := c[i+j+len/2] * w
				c[i+j] = u + v
				c[i+j+len/2] = u - v
				w *= wlen
			}
		}
	}
	// take absolute values of the first half, including nyquist frequency but excluding DC
	m := n / 2
	t1 := sd.temp.tmp1[:m]
	t2 := sd.temp.tmp2[:m]
	for i := 0; i < m; i++ {
		t1[i] = float32(cmplx.Abs(c[1+i])) // do not include DC
	}
	// square the amplitudes to get power
	vek32.Mul_Into(t2, t1, t1)
	vek32.DivNumber_Inplace(t2, sd.temp.normFactor*sd.temp.normFactor) // normalize for windowing
	// Since we are using a real-valued FFT, we need to double the values except for Nyquist (and DC, but we don't have that here)
	vek32.MulNumber_Inplace(t2[:m-1], 2)
	// calculate difference to current spectrum and add back, multiplied by smoothing factor
	vek32.Sub_Inplace(t2, sd.temp.power[channel])
	alpha := spectrumSmoothingMap[sd.settings.Smooth]
	vek32.MulNumber_Inplace(t2, alpha)
	vek32.Add_Inplace(sd.temp.power[channel], t2)
}