package main

import (
	"encoding/csv"
	"image/color"
	"os"
	"sort"
	"strconv"

	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
	"gonum.org/v1/plot/vg/draw"
	"gonum.org/v1/plot/vg/vgimg"
	"gorgonia.org/tensor"
	"gorgonia.org/tensor/native"
)

func main() {
	sizeCPU := func(a Bench) (float64, float64) { return a.Size, a.CPU }
	// median := func(a []float64) float64 { return ntile(a, 0.5) }
	all := parse(ingest("data.csv"))
	list := []string{
		"gcc", "ocaml", "node",
		"java", "go", "ghc",
		"rust", "julia", "swift",
	}

	ps := make([]*plot.Plot, 0, len(list))
	for _, l := range list {
		lang := filter(all, func(a Bench) bool { return a.Language == l })
		lang = argminByName(lang, func(a Bench) float64 { return a.CPU })
		p := plotAll(all, sizeCPU)
		p.BackgroundColor = color.RGBA{R: 255}
		p.X.Max = 2000
		p.Y.Max = 150
		p.Title.Text = l

		plotStar(p, lang, sizeCPU, mean)
		ps = append(ps, p)
	}
	// Draw plots in a tiled fashion
	cols := 3
	t := tensor.New(tensor.WithBacking(ps), tensor.WithShape(len(list)/cols, cols))
	matUgh, err := native.Matrix(t)
	dieIfErr(err)
	mat := matUgh.([][]*plot.Plot)
	tiles := draw.Tiles{Rows: t.Shape()[0], Cols: t.Shape()[1]}
	img := vgimg.New(60*vg.Centimeter, 60*vg.Centimeter)
	canvas := draw.New(img)
	mini := plot.Align(mat, tiles, canvas)
	for i := 0; i < tiles.Rows; i++ {
		for j := 0; j < tiles.Cols; j++ {
			mat[i][j].Draw(mini[i][j])
		}
	}
	w, err := os.Create("gb.png")
	dieIfErr(err)
	png := vgimg.PngCanvas{Canvas: img}
	png.WriteTo(w)
}

func ingest(filename string) [][]string {
	f, err := os.Open(filename)
	dieIfErr(err)
	r := csv.NewReader(f)
	records, err := r.ReadAll()
	dieIfErr(err)
	return records
}

type Bench struct {
	Name     string
	Language string
	Size     float64
	CPU, Mem float64
}

func parse(recs [][]string) []Bench {
	retVal := make([]Bench, 0, len(recs))
	for i, r := range recs {
		if i == 0 {
			continue
		}
		size, err := strconv.ParseFloat(r[4], 64)
		dieIfErr(err)
		cpu, err := strconv.ParseFloat(r[5], 64)
		dieIfErr(err)
		mem, err := strconv.ParseFloat(r[6], 64)
		dieIfErr(err)

		b := Bench{
			Name:     r[0],
			Language: r[1],
			Size:     size,
			CPU:      cpu,
			Mem:      mem,
		}
		retVal = append(retVal, b)
	}
	return retVal
}

// filter filters a list of Bench according to the criteria given in f
func filter(a []Bench, f func(a Bench) bool) (retVal []Bench) {
	for i := range a {
		if f(a[i]) {
			retVal = append(retVal, a[i])
		}
	}
	return retVal
}

// reduce2 reduces a pair of numbers given by f, using the reduction function g
func reduce2(a []Bench, f func(a Bench) (float64, float64), g func([]float64) float64) (float64, float64) {
	xs, ys := make([]float64, 0, len(a)), make([]float64, 0, len(a))
	for i := range a {
		x, y := f(a[i])
		xs = append(xs, x)
		ys = append(ys, y)
	}
	return g(xs), g(ys)
}

func argminByName(a []Bench, f func(a Bench) float64) (retVal []Bench) {
	m := make(map[string]Bench)
	for _, b := range a {
		if v, ok := m[b.Name]; ok {
			if f(b) < f(v) {
				m[b.Name] = b
			}
			continue
		}
		m[b.Name] = b
	}
	for _, v := range m {
		retVal = append(retVal, v)
	}
	return retVal
}

// makeXYs extracts pairs of desired numbers given by f
func makeXYs(a []Bench, f func(a Bench) (float64, float64)) plotter.XYs {
	retVal := make(plotter.XYs, len(a))
	for i := range a {
		x, y := f(a[i])
		retVal[i].X = x
		retVal[i].Y = y
	}
	return retVal
}

// plotAll plots all the points of the given []Bench
func plotAll(a []Bench, f func(a Bench) (float64, float64)) *plot.Plot {
	p, err := plot.New()
	dieIfErr(err)
	s, err := plotter.NewScatter(makeXYs(a, f))
	dieIfErr(err)
	p.Add(s)
	return p
}

// plotstar adds a line star thing into the plot
func plotStar(p *plot.Plot, a []Bench, f func(a Bench) (float64, float64), g func([]float64) float64) {
	s := new(star)
	s.XYs = makeXYs(a, f)
	s.mx, s.my = reduce2(a, f, g)
	s.trx, s.try = p.X.Max, p.Y.Max
	s.LineStyle = plotter.DefaultLineStyle
	s.LineStyle.Color = color.RGBA{R: 255, A: 255}
	p.Add(s)
}

// star is a data structure used for plotting line stars
type star struct {
	plotter.XYs
	draw.LineStyle
	mx, my   float64
	trx, try float64 // truncate at
}

func (s *star) Plot(c draw.Canvas, p *plot.Plot) {
	tx, ty := p.Transforms(&c)
	trx, try := tx(s.trx), ty(s.try)
	ls := s.LineStyle
	mx, my := tx(s.mx), ty(s.my)
	for _, xy := range s.XYs {
		x := tx(xy.X)
		y := ty(xy.Y)
		if x > trx {
			x = trx
		}
		if y > try {
			y = try
		}
		c.StrokeLine2(ls, mx, my, x, y)
	}
}

func dieIfErr(err error) {
	if err != nil {
		panic(err)
	}
}

func mean(a []float64) float64 {
	var s float64
	for _, v := range a {
		s += v
	}
	return s / float64(len(a))
}

func ntile(a []float64, p float64) float64 {
	sort.Float64s(a)
	at := int(float64(len(a)) * p)
	return a[at]
}