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 FFT).
// * Multiply by A-weight curve
// * Take log10 of result
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])
	}

	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
		if freqs[i] < 20 {
			amps[i] = 0
		} else {
			amps[i] = math.Log10(cmplx.Abs(cCoeffs[i]) * AWeight(freqs[i]))
		}
	}

	// Normalize coefficients.
	//for i := range amps {
	//amps[i] /= ampMax
	//}

	// Find highest N peaks and sort by size.

	/*
		// Use a moving window to search for local peaks.
		hws := []*HarmonicWindow{}

		// Bypass low-frequency samples.
		i := 0
		for i < len(freqs) && freqs[i] < 24 {
			i++
		}

		iCenter := i
		for iC{

		}
	*/
	// Find frequency window around midi note.

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

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

	log.Printf("Plotting...")
	if err = plotutil.AddLinePoints(p, "fft", pts); err != nil {
		panic(err)
	}

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

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

	harmonics := []float64{0.5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
	hws := []*HarmonicWindow{}
	for _, harmonic := range harmonics {

		hw := NewHarmonicWindow(
			fFundamental,
			harmonic,
			amps,
			freqs)

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

		log.Printf("Making points...")
		pts := make(plotter.XYs, 0, len(hw.Amps))
		for i := range hw.Amps {
			pts = append(pts, plotter.XY{
				X: hw.Freqs[i],
				Y: hw.Amps[i],
			})
		}

		log.Printf("Plotting...")
		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-%03.2f.png", harmonic)
		if err := p.Save(8*vg.Inch, 8*vg.Inch, filename); err != nil {
			panic(err)
		}

		hws = append(hws, hw)
	}

	pts = make(plotter.XYs, len(hws))
	for i, hw := range hws[:11] {
		pts[i].X = hw.Harmonic
		pts[i].Y = hw.FreqPeakInterp
	}

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

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

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