valhallir-deconvolver/pkg/plot/plot.go

package plot

import (
	"bytes"
	"fmt"
	"image/color"
	"image/png"
	"math"
	"math/cmplx"
	"os"
	"path/filepath"
	"strings"

	"github.com/mjibson/go-dsp/fft"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/font"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
	"gonum.org/v1/plot/vg/draw"
	"gonum.org/v1/plot/vg/vgimg"
)

// PlotIR plots the frequency response (magnitude in dB vs. frequency in Hz) of the IR to a PNG file
// logoData is the embedded logo image data (can be nil if not available)
func PlotIR(ir []float64, sampleRate int, irFileName string, logoData []byte) error {
	if len(ir) == 0 {
		return nil
	}
	// Use only the first 8192 samples of the IR for plotting
	windowLen := 8192
	if len(ir) < windowLen {
		windowLen = len(ir)
	}
	irWin := ir[:windowLen]
	X := fft.FFTReal(irWin)
	// Plot from 20 Hz up to 20kHz, include every bin
	var plotPts plotter.XYs
	var minDb float64 = 1e9
	var maxDb float64 = -1e9
	var minDbFreq float64
	freqBins := windowLen / 2
	for i := 1; i < freqBins; i++ {
		freq := float64(i) * float64(sampleRate) / float64(windowLen)
		if freq < 20.0 {
			continue
		}
		if freq > 20000.0 {
			break
		}
		mag := cmplx.Abs(X[i])
		if mag < 1e-12 {
			mag = 1e-12
		}
		db := 20 * math.Log10(mag)
		plotPts = append(plotPts, plotter.XY{X: freq, Y: db})
		if db < minDb {
			minDb = db
			minDbFreq = freq
		}
		if db > maxDb {
			maxDb = db
		}
	}
	fmt.Printf("[PlotIR] minDb in plotted range: %.2f dB at %.2f Hz\n", minDb, minDbFreq)
	p := plot.New()
	p.Title.Text = "IR Frequency Response"
	p.X.Label.Text = "Frequency (Hz)"
	p.Y.Label.Text = "Magnitude (dB)"
	p.X.Scale = plot.LogScale{}
	p.X.Tick.Marker = plot.TickerFunc(func(min, max float64) []plot.Tick {
		ticks := []float64{20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000}
		labels := []string{"20", "50", "100", "200", "500", "1k", "2k", "5k", "10k", "20k"}
		var result []plot.Tick
		for i, v := range ticks {
			if v >= min && v <= max {
				result = append(result, plot.Tick{Value: v, Label: labels[i]})
			}
		}
		return result
	})
	line, err := plotter.NewLine(plotPts)
	if err != nil {
		return err
	}
	// Set line color to blue
	line.Color = color.RGBA{R: 30, G: 100, B: 220, A: 255}
	p.Add(line)
	// Find minimum dB value between 20 Hz and 50 Hz for y-axis anchor
	minDb2050 := 1e9
	for i := 1; i < freqBins; i++ {
		freq := float64(i) * float64(sampleRate) / float64(windowLen)
		if freq < 20.0 {
			continue
		}
		if freq > 50.0 {
			break
		}
		mag := cmplx.Abs(X[i])
		if mag < 1e-12 {
			mag = 1e-12
		}
		db := 20 * math.Log10(mag)
		if db < minDb2050 {
			minDb2050 = db
		}
	}
	p.Y.Min = minDb2050
	p.Y.Max = math.Ceil(maxDb)
	p.X.Min = 20.0
	p.X.Max = 20000.0

	// --- Time-aligned waveform plot ---
	p2 := plot.New()
	p2.Title.Text = "IR Waveform"
	p2.X.Label.Text = "Time (ms)"
	p2.Y.Label.Text = "Amplitude"
	// Prepare waveform data (only first 10ms)
	var pts plotter.XYs
	maxTimeMs := 10.0
	for i := 0; i < windowLen; i++ {
		t := float64(i) * 1000.0 / float64(sampleRate) // ms
		if t > maxTimeMs {
			break
		}
		pts = append(pts, plotter.XY{X: t, Y: irWin[i]})
	}
	wline, err := plotter.NewLine(pts)
	if err != nil {
		return err
	}
	wline.Color = color.RGBA{R: 30, G: 100, B: 220, A: 255}
	p2.Add(wline)
	p2.X.Min = 0
	p2.X.Max = maxTimeMs
	// Y range auto

	// --- Compose both plots vertically ---
	const width = 6 * vg.Inch
	const height = 8 * vg.Inch                // increased height for frequency diagram
	img := vgimg.New(width, height+1*vg.Inch) // extra space for logo/headline
	dc := draw.New(img)

	// Draw logo at the top left, headline to the right, IR filename below
	// Logo is embedded in the binary
	logoW := 2.4 * vg.Inch  // doubled size
	logoH := 0.68 * vg.Inch // doubled size
	logoX := 0.3 * vg.Inch
	logoY := height + 0.2*vg.Inch // move logo down by an additional ~10px
	logoDrawn := false
	if len(logoData) > 0 {
		logoImg, err := png.Decode(bytes.NewReader(logoData))
		if err == nil {
			rect := vg.Rectangle{
				Min: vg.Point{X: logoX, Y: logoY},
				Max: vg.Point{X: logoX + logoW, Y: logoY + logoH},
			}
			dc.DrawImage(rect, logoImg)
			logoDrawn = true
		}
	}
	// Draw headline (bold, larger) to the right of the logo
	headline := "Valhallir Deconvolver IR Analysis"
	fntSize := vg.Points(14) // Same as IR filename
	if logoDrawn {
		headlineX := logoX + logoW + 0.3*vg.Inch
		headlineY := logoY + logoH - vg.Points(16) - vg.Points(5) // move headline up by ~10px
		boldFont := plot.DefaultFont
		boldFont.Weight = 3 // font.WeightBold is 3 in gonum/plot/font
		boldFace := font.DefaultCache.Lookup(boldFont, fntSize)
		dc.SetColor(color.Black)
		dc.FillString(boldFace, vg.Point{X: headlineX, Y: headlineY}, headline)
		// Draw IR filename below headline, left-aligned, standard font
		fileLabel := "IR-File: " + filepath.Base(irFileName)
		fileY := headlineY - fntSize - vg.Points(6)
		fileFace := font.DefaultCache.Lookup(plot.DefaultFont, vg.Points(10))
		dc.FillString(fileFace, vg.Point{X: headlineX, Y: fileY}, fileLabel)
	}

	// Custom tile arrangement: frequency diagram gets more height, waveform gets less
	tiles := draw.Tiles{
		Rows:   2,
		Cols:   1,
		PadX:   vg.Millimeter,
		PadY:   20 * vg.Millimeter, // more space between plots to emphasize frequency diagram
		PadTop: vg.Points(15),      // move diagrams down by ~20px
	}

	// Offset the plots down by 1 inch to make space for logo/headline
	imgPlots := vgimg.New(width, height)
	dcPlots := draw.New(imgPlots)
	canvases := plot.Align([][]*plot.Plot{{p}, {p2}}, tiles, dcPlots)
	p.Draw(canvases[0][0])
	p2.Draw(canvases[1][0])
	dc.DrawImage(vg.Rectangle{Min: vg.Point{X: 0, Y: 0}, Max: vg.Point{X: width, Y: height}}, imgPlots.Image())

	// Save as PNG in the same directory as the IR file
	irDir := filepath.Dir(irFileName)
	irBase := filepath.Base(irFileName)
	irNameWithoutExt := strings.TrimSuffix(irBase, filepath.Ext(irBase))
	plotFileName := filepath.Join(irDir, irNameWithoutExt+".png")

	plotFile, err := os.Create(plotFileName)
	if err != nil {
		return err
	}
	defer plotFile.Close()
	_, err = vgimg.PngCanvas{Canvas: img}.WriteTo(plotFile)
	return err
}