package extractfreq

import (
	"fmt"
	"log"
	"math"
	"math/cmplx"

	"git.crumpington.com/public/jlaudio/lib/flac"
	"gonum.org/v1/gonum/fourier"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"
)

// Steps:
//
// * Create amplitude spectrum (abs of real FFT).
// * Multiply by A-weight curve
// * Find peaks
// * Sort peaks by amplitude
// * Use strongest peak to estimate fFundamental?
// *
func PlotFFT(path string, nSamples int, fFundamental float64) {
	L, R, err := flac.Load(path)
	if err != nil {
		panic(err)
	}

	if nSamples < len(L) {
		nSamples = len(L)
	}

	data := make([]float64, nSamples)
	for i := range data {
		data[i] = float64(L[i]) + float64(R[i])
	}

	// Make initial fft.
	fft := fourier.NewFFT(nSamples)
	cCoeffs := fft.Coefficients(nil, data)
	amps := make([]float64, len(cCoeffs))
	freqs := make([]float64, len(cCoeffs))
	for i := range freqs {
		freqs[i] = fft.Freq(i) * 48000
		amps[i] = cmplx.Abs(cCoeffs[i])
	}

	// Trim to frequencies of interest.
	deltaF := freqs[1] - freqs[0]
	idx0 := int(20 / deltaF)
	idx1 := int(16000 / deltaF)
	freqs = freqs[idx0:idx1]
	amps = amps[idx0:idx1]

	// Multiply by A-weight function.
	for i := range amps {
		amps[i] *= AWeight(freqs[i])
	}

	// Find peaks.
	dFreq := freqs[1] - freqs[0]
	halfWinSize := int(math.Round((fFundamental / dFreq) / 3))
	peakInds, peakAmps := getPeaks(halfWinSize, amps)
	peakInds, peakAmps = sortPeaks(peakInds, peakAmps)

	p, err := plot.New()
	if err != nil {
		panic(err)
	}

	fMin := float64(20)
	fMax := float64(16000)

	log.Printf("Making points...")
	pts := make(plotter.XYs, 0, len(amps))
	for i := range amps {
		if freqs[i] > fMin && freqs[i] < fMax {
			pts = append(pts, plotter.XY{
				X: freqs[i],
				Y: amps[i],
			})
		}
	}

	pts2 := make(plotter.XYs, 0, len(amps))
	for i := range peakInds {
		if freqs[peakInds[i]] > fMin && freqs[peakInds[i]] < fMax {
			pts2 = append(pts2, plotter.XY{
				X: freqs[peakInds[i]],
				Y: peakAmps[i],
			})
		}
	}

	log.Printf("Plotting...")
	if err = plotutil.AddScatters(p, "fft", pts, "max", pts2); err != nil {
		panic(err)
	}
	p.Add(plotter.NewGrid())

	log.Printf("Saving...")
	if err := p.Save(16*vg.Inch, 8*vg.Inch, "fft.png"); err != nil {
		panic(err)
	}

	for i, idx := range peakInds {
		//amp := peakAmps[i]

		// Make window.
		ampsWin := amps[idx-halfWinSize : idx+halfWinSize]
		freqWin := freqs[idx-halfWinSize : idx+halfWinSize]

		log.Printf("Making points...")
		pts := make(plotter.XYs, 0, 2*halfWinSize)
		for i := range ampsWin {
			pts = append(pts, plotter.XY{
				X: freqWin[i],
				Y: ampsWin[i],
			})
		}

		log.Printf("Plotting...")
		p, err := plot.New()
		if err != nil {
			panic(err)
		}

		if err = plotutil.AddLines(p, "fft", pts); err != nil {
			panic(err)
		}

		//p.Y.Scale = plot.LogScale{}
		//p.X.Scale = plot.LogScale{}

		log.Printf("Saving...")
		filename := fmt.Sprintf("plot-%03d.png", i)
		if err := p.Save(8*vg.Inch, 8*vg.Inch, filename); err != nil {
			panic(err)
		}
	}
}