// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package plot

import (
	"bytes"
	"cmp"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"slices"
	"strconv"
	"strings"
)

type gnuplotter struct {
	*Plot
	code bytes.Buffer

	confidence float64

	colorScale, pointScale func(point) int
	colorRev, pointRev     func(int) value

	nColors, nPoints int
}

func (p *Plot) Gnuplot(term string, out io.Writer) error {
	pl := gnuplotter{Plot: p}
	if err := pl.plot(term); err != nil {
		return err
	}

	// TODO: Only do this if we're launching a windowed gnuplot
	if false {
		fmt.Fprintf(&pl.code, "pause mouse close\n")
	}
	code := pl.code.Bytes()

	switch term {
	case "":
		_, err := out.Write(code)
		return err
	case "png":
		cmd := exec.Command("gnuplot")
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return fmt.Errorf("creating pipe to gnuplot: %w", err)
		}
		cmd.Stdout = out
		cmd.Stderr = os.Stderr
		if err := cmd.Start(); err != nil {
			return fmt.Errorf("starting gnuplot: %w", err)
		}
		defer cmd.Process.Kill()
		if _, err := stdin.Write(code); err != nil {
			return fmt.Errorf("writing to gnuplot: %w", err)
		}
		stdin.Close()
		if err := cmd.Wait(); err != nil {
			return fmt.Errorf("gnuplot failed: %w", err)
		}
	}
	return nil
}

func (p *gnuplotter) plot(term string) error {
	pts := p.points
	p.confidence = 0.95

	if len(pts) == 0 {
		return fmt.Errorf("no data")
	}

	if pointsKinds(pts, AesX)&kindContinuous == 0 {
		// TODO: Bar chart
		//
		// TODO: Say what type it *is* and preferably an example of what made is
		// non-continuous.
		return fmt.Errorf("non-numeric X data not supported")
	}
	if pointsKinds(pts, AesY)&kindContinuous == 0 {
		// TODO: Horizontal bar chart?
		return fmt.Errorf("non-numeric Y data not supported")
	}
	rowScale, rowRev, nRows := ordScale(pts, AesRow)
	colScale, colRev, nCols := ordScale(pts, AesCol)
	multiplot := nRows > 1 || nCols > 1
	p.colorScale, p.colorRev, p.nColors = ordScale(pts, AesColor)
	p.pointScale, p.pointRev, p.nPoints = ordScale(pts, AesPoint)

	switch term {
	case "":
		// Just code
	case "png":
		fmt.Fprintf(&p.code, "set terminal pngcairo size %d,%d\n", nCols*640, nRows*480)
	default:
		return fmt.Errorf("unknown output type %s", term)
	}

	// Construct multiplot row and column labels
	rowLabels := make([]string, nRows)
	colLabels := make([]string, nCols)
	if multiplot {
		// Configure multiplot
		fmt.Fprintf(&p.code, "set multiplot layout %d,%d columnsfirst margins char 12,char 3,char 4,char 2 spacing char 10, char 4\n", nRows, nCols)
		// TODO: Use a benchproc.KeyHeader to display this more nicely. It may
		// be really hard to center a label across a multiple graphs, though.
		for row := range nRows {
			rowLabels[row] = rowRev(row).StringValues()
		}
		for col := range nCols {
			colLabels[col] = colRev(col).StringValues()
		}
	}

	// Set log scales
	setLogScale := func(aes Aes, name string) {
		if base := p.logScale.Get(aes); base != 0 {
			fmt.Fprintf(&p.code, "set logscale %s %d\n", name, base)
		}
	}
	setLogScale(AesX, "x")
	setLogScale(AesY, "y")

	// Sort the points in the order the data must be emitted.
	slices.SortFunc(pts, func(a, b point) int {
		if c := a.Get(AesCol).compare(b.Get(AesCol)); c != 0 {
			return c
		}
		if c := a.Get(AesRow).compare(b.Get(AesRow)); c != 0 {
			return c
		}
		if c := a.Get(AesColor).compare(b.Get(AesColor)); c != 0 {
			return c
		}
		if c := a.Get(AesPoint).compare(b.Get(AesPoint)); c != 0 {
			return c
		}
		// For a line plot, X must be sorted numerically.
		return cmp.Compare(a.Get(AesX).val, b.Get(AesX).val)
	})

	// Emit plots
	type rowCol struct{ row, col int }
	plots, _ := groupBy(pts, func(pt point) rowCol {
		return rowCol{rowScale(pt), colScale(pt)}
	})
	for col := range nCols {
		for row := range nRows {
			pts := plots[rowCol{row, col}]
			if multiplot && col == 0 {
				// Label this row.
				fmt.Fprintf(&p.code, "set label 1 %s at char 2, graph 0.5 center rotate by 90\n", gpString(rowLabels[row]))
			}
			if multiplot && row == 0 {
				// Label this column.
				fmt.Fprintf(&p.code, "set title %s\n", gpString(colLabels[col]))
			}
			p.onePlot(pts, rowLabels[row], colLabels[col])
			fmt.Fprintf(&p.code, "unset label 1\n")
			fmt.Fprintf(&p.code, "unset title\n")
		}
	}

	if multiplot {
		fmt.Fprintf(&p.code, "unset multiplot\n")
	}

	return nil
}

func pointAesGetter2(aes1, aes2 Aes) func(pt point) [2]value {
	return func(pt point) [2]value {
		return [2]value{pt.Get(aes1), pt.Get(aes2)}
	}
}

func isPowerOf2(x int) bool {
	return x > 0 && x&(x-1) == 0
}

func (p *gnuplotter) onePlot(pts []point, rowLabel, colLabel string) {
	if len(pts) == 0 {
		// Skip this plot.
		fmt.Fprintf(&p.code, "set multiplot next\n")
		return
	}

	var reset strings.Builder

	// Let gnuplot print scientific values on tick marks. This is much nicer
	// than putting it on the unit.
	setFormat := func(axis string, aes Aes) (scale func(float64) float64, label string) {
		kinds := pointsKinds(pts, aes)
		// Scale the values. We ask for no rescaling because we'll configure
		// gnuplot to do the scientific scaling for us.
		scale, _, _, label, _ = p.continuousScale(pts, aes, false)
		if kinds&kindRatio != 0 {
			// Format ratios as a percent delta.
			fmt.Fprintf(&p.code, "set format %s '%%+h%%%%'\n", axis)
			scale = func(x float64) float64 { return (x - 1) * 100 }
			label = "delta " + label

			// Always include 0.
			fmt.Fprintf(&p.code, "set %srange [*<0:0<*]\n", axis)

			// Draw a line at 0. The command uses the opposite axis.
			za := "x"
			if aes == AesX {
				za = "y"
			}
			fmt.Fprintf(&p.code, "set %szeroaxis dt 2\n", za)
			fmt.Fprintf(&reset, "unset %szeroaxis\n", za)
		} else {
			if base := p.logScale.Get(aes); base != 0 && isPowerOf2(base) {
				fmt.Fprintf(&p.code, "set format %s '%%.0b%%B'\n", axis)
			} else {
				fmt.Fprintf(&p.code, "set format %s '%%.0s%%c'\n", axis)
			}
		}
		return
	}
	xScale, xLabel := setFormat("x", AesX)
	yScale, yLabel := setFormat("y", AesY)

	// Set axis labels. Skip these if they would just duplicate the row/column
	// labels. This often happens with the default -row .unit -y .value
	// configuration, which will produce the same labels for the Y axis and the
	// row.
	//
	// TODO: If all X axes in a column have the same labels, but they differ
	// from the column label, consider showing only the X axis label in the
	// bottom-most graph.
	if colLabel != xLabel {
		fmt.Fprintf(&p.code, "set xlabel %s\n", gpString(xLabel))
	}
	if rowLabel != yLabel {
		fmt.Fprintf(&p.code, "set ylabel %s\n", gpString(yLabel))
	}

	// TODO: Should this be done up front? Then continuousScale would
	// have to understand summaries, but that's fine.
	//
	// TODO: Do something with the warnings. Allow configuring
	// confidence.
	pts, _ = transformSummarize(pts, AesY, p.confidence)

	// Set up for plotting ratios.
	kinds := pointsKinds(pts, AesY)
	var ratioPos, ratioNeg string
	if p.dvAes == AesY && kinds&kindRatio != 0 {
		// Check that all units have the same "better" direction.
		better := 0
		for i, unitName := range p.pointsUnits(pts) {
			better1 := p.units.GetBetter(unitName)
			if i == 0 {
				better = better1
			} else if better != better1 {
				better = 0
				break
			}
		}

		if better > 0 {
			ratioPos, ratioNeg = "red", "green"
		} else if better < 0 {
			ratioPos, ratioNeg = "green", "red"
		}
	}

	// Construct the key.
	//
	// TODO: There's some general n-way formulation of this, but it's not coming
	// to me right now. I think it's something like "everything is lp, but we
	// have a default linecolor/pointtype/pointsize when we're not iterating
	// over that dimension".
	var key grid2D[string]
	if p.nColors > 1 && p.nPoints > 1 {
		// Two separate columns for colors and points.
		//
		// TODO: Maybe not if they're totally correlated?
		//
		// TODO: Maybe only for values that appear in this plot?
		for color := range p.nColors {
			label := p.colorRev(color).StringValues()
			sample := fmt.Sprintf("with l title %s linecolor %d", gpString(label), color+1)
			key.Set(0, color, sample)
		}
		for point := range p.nPoints {
			label := p.pointRev(point).StringValues()
			sample := fmt.Sprintf("with p title %s lc 0 pointtype %d ps 2", gpString(label), point+1)
			key.Set(1, point, sample)
		}
	} else {
		// One key that combines colors and points.
		i := 0
		sliceBy(pts, pointAesGetter2(AesColor, AesPoint),
			func(seriesKey [2]value, pts []point) {
				title := ""
				for _, k := range seriesKey {
					str := k.StringValues()
					if title != "" && str != "" {
						title += " "
					}
					title += str
				}
				colorIdx := p.colorScale(pts[0]) + 1
				pointIdx := p.pointScale(pts[0]) + 1
				sample := fmt.Sprintf("with lp title %s linecolor %d pointtype %d ps 2", gpString(title), colorIdx, pointIdx)
				key.Set(0, i, sample)
				i++
			})
	}

	// Emit point data and build plot command. We build this up in several layers.
	const (
		layerPos = iota
		layerNeg
		layerRange
		layerCenter
		maxLayers
	)
	var plotArgs []string
	var data strings.Builder
	anyRange := false
	// TODO: Gnuplot does support "variable" point type, so it doesn't
	// necessarily have to separate series. But if we went down that road you'd
	// need some way to separate series that have the same color but different
	// points, which is probably not the default you want. We could have
	// different series separating and non-separating aesthetics even if they
	// control the same thing, or maybe a way to specify options on an aesthetic
	// for whether it separates series (some of them would always have to be one
	// way or another, possibly depending on the plot target), or maybe a way to
	// specify a particular mark with options specific to that mark (though that
	// complicates all sorts of things and makes the whole tool less automagic).
	for layer := range maxLayers {
		sliceBy(pts, pointAesGetter2(AesColor, AesPoint),
			func(seriesKey [2]value, pts []point) {
				colorIdx := p.colorScale(pts[0]) + 1
				pointIdx := p.pointScale(pts[0]) + 1

				if layer == layerRange {
					haveRange := false
					for _, pt := range pts {
						if !math.IsInf(pt.Get(AesY).summary.Lo, 0) {
							haveRange, anyRange = true, true
							break
						}
					}
					if !haveRange {
						return
					}

					// Emit range
					plotArg := fmt.Sprintf("'-' using 1:2:3 with filledcurves notitle fc linetype %d fs transparent solid 0.25", colorIdx)
					plotArgs = append(plotArgs, plotArg)

					for _, pt := range pts {
						x := pt.Get(AesX).val
						y := pt.Get(AesY).summary
						if !math.IsInf(y.Lo, 0) {
							fmt.Fprintf(&data, "%g %g %g\n", xScale(x), yScale(y.Lo), yScale(y.Hi))
						}
					}
					fmt.Fprintf(&data, "e\n")
					return
				}

				// The other layers all use the same data and just different
				// plot options.
				plotArg := "'-' using 1:2"
				switch layer {
				case layerPos:
					if ratioPos == "" {
						return
					}
					plotArg += " with filledcurves below notitle fs transparent solid 0.1 fc '" + ratioPos + "' lw 0"
				case layerNeg:
					if ratioNeg == "" {
						return
					}
					plotArg += " with filledcurves above notitle fs transparent solid 0.1 fc '" + ratioNeg + "' lw 0"
				case layerCenter:
					plotArg += fmt.Sprintf(" with lp notitle linecolor %d pointtype %d ps 2", colorIdx, pointIdx)
				}

				// Emit center curve.
				plotArgs = append(plotArgs, plotArg)
				for _, pt := range pts {
					fmt.Fprintf(&data, "%g %g\n", xScale(pt.Get(AesX).val), yScale(pt.Get(AesY).val))
				}
				fmt.Fprintf(&data, "e\n")
			})
	}

	if anyRange {
		// Add a legend entry for the range.
		sample := fmt.Sprintf("with filledcurves title '%v%% confidence' fc linetype 0 fs transparent solid 0.25", p.confidence*100)
		key.Set(0, key.h, sample)
	}

	// Configure key.
	if key.w > 1 {
		fmt.Fprintf(&p.code, "set key maxrows %d\n", key.h)
		defer fmt.Fprintf(&p.code, "set key default\n")
	}
	for x := range key.w {
		for y := range key.h {
			sample, ok := key.Get(x, y)
			if ok {
				plotArgs = append(plotArgs, "keyentry "+sample)
			} else {
				// Gnuplot will trim away any notitle keys at the end, so we
				// have to produce "real" key entries, at least for the very
				// last one.
				plotArgs = append(plotArgs, "keyentry title ' '")
			}
		}
	}

	fmt.Fprintf(&p.code, "plot %s\n", strings.Join(plotArgs, ", "))

	p.code.WriteString(data.String())

	p.code.WriteString(reset.String())
}

// gpString returns s escaped for Gnuplot
func gpString(s string) string {
	// I can't find any documentation on Gnuplot's escape syntax, but as far as
	// I can tell, it's compatible with Go's escaping rules.
	return strconv.Quote(s)
}

// grid2D is a sparse mapping from 2-D indexes to V.
type grid2D[V any] struct {
	w, h int
	vals map[grid2DKey]V
}

type grid2DKey struct{ x, y int }

func (g *grid2D[V]) Set(x, y int, v V) {
	if g.vals == nil {
		g.vals = make(map[grid2DKey]V)
	}
	g.vals[grid2DKey{x, y}] = v
	g.w = max(g.w, x+1)
	g.h = max(g.h, y+1)
}

func (g *grid2D[V]) Get(x, y int) (V, bool) {
	v, ok := g.vals[grid2DKey{x, y}]
	return v, ok
}
