package audio

import (
	"image"
	"image/color"
	"math"
	"math/cmplx"
)

// Spectrogram считает STFT сигнала: режет samples на перекрывающиеся
// окна размером windowSize (округляется вверх до степени двойки) с
// шагом hopSize, на каждое окно накладывает окно Ханна и считает БПФ.
// Возвращает матрицу [номер_окна][номер_частотного_бина] = амплитуда
// в дБ. Используются только первые windowSize/2 бинов — для
// вещественного сигнала спектр симметричен, вторая половина не несёт
// новой информации.
func Spectrogram(samples []float64, windowSize, hopSize int) [][]float64 {
	windowSize = nextPowerOfTwo(windowSize)
	if hopSize < 1 {
		hopSize = windowSize / 4
	}
	window := hannWindow(windowSize)

	var frames [][]float64
	for start := 0; start+windowSize <= len(samples); start += hopSize {
		buf := make([]complex128, windowSize)
		for i := 0; i < windowSize; i++ {
			buf[i] = complex(samples[start+i]*window[i], 0)
		}
		fft(buf)

		bins := make([]float64, windowSize/2)
		for i := 0; i < windowSize/2; i++ {
			mag := cmplx.Abs(buf[i]) / float64(windowSize)
			// +эпсилон — чтобы log10(0) не улетел в -Inf на полной тишине
			bins[i] = 20 * math.Log10(mag+1e-12)
		}
		frames = append(frames, bins)
	}
	return frames
}

// RenderSpectrogram красит матрицу STFT в изображение: время по
// горизонтали (номер окна), частота по вертикали (низкие частоты
// внизу, высокие — вверху, как в любом спектрограф-вьюере вроде
// Sonic Visualiser). Амплитуда нормализуется по фактическому диапазону
// дБ в самой записи (а не по абсолютной шкале) — так тихие записи не
// превращаются в сплошной чёрный экран.
func RenderSpectrogram(frames [][]float64) *image.RGBA {
	if len(frames) == 0 || len(frames[0]) == 0 {
		return image.NewRGBA(image.Rect(0, 0, 1, 1))
	}
	width := len(frames)
	height := len(frames[0])

	minDB, maxDB := frames[0][0], frames[0][0]
	for _, f := range frames {
		for _, v := range f {
			if v < minDB {
				minDB = v
			}
			if v > maxDB {
				maxDB = v
			}
		}
	}
	if maxDB-minDB < 1 {
		maxDB = minDB + 1 // защита от деления на почти ноль на абсолютно однородном сигнале
	}

	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for x, f := range frames {
		for bin, v := range f {
			t := (v - minDB) / (maxDB - minDB)
			if t < 0 {
				t = 0
			}
			if t > 1 {
				t = 1
			}
			y := height - 1 - bin // низкие частоты внизу изображения
			img.Set(x, y, heatColor(t))
		}
	}
	return img
}

// heatColor — точки градиента: чёрный (тихо) → синий → циан → жёлтый →
// белый (громко). Обычная "тепловая карта" для спектрограмм — резкий
// контраст между тихим фоном и громким сигналом, именно это и нужно,
// чтобы разглядеть спрятанный в спектре текст/QR-код/паттерн.
func heatColor(t float64) color.RGBA {
	stops := []color.RGBA{
		{R: 0, G: 0, B: 0, A: 255},
		{R: 0, G: 0, B: 180, A: 255},
		{R: 0, G: 200, B: 200, A: 255},
		{R: 255, G: 230, B: 0, A: 255},
		{R: 255, G: 255, B: 255, A: 255},
	}
	if t <= 0 {
		return stops[0]
	}
	if t >= 1 {
		return stops[len(stops)-1]
	}
	seg := t * float64(len(stops)-1)
	i := int(seg)
	if i >= len(stops)-1 {
		return stops[len(stops)-1]
	}
	frac := seg - float64(i)
	a, b := stops[i], stops[i+1]
	return color.RGBA{
		R: lerp8(a.R, b.R, frac),
		G: lerp8(a.G, b.G, frac),
		B: lerp8(a.B, b.B, frac),
		A: 255,
	}
}

func lerp8(a, b uint8, t float64) uint8 {
	return uint8(float64(a) + (float64(b)-float64(a))*t)
}