lib/heatmap.go

// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package lib

import (
	"fmt"
	im "image"

	ui "github.com/gizak/termui/v3"
	rw "github.com/mattn/go-runewidth"
)

type HeatMap struct {
	ui.Block
	XLabelStyles []ui.Style
	CellColors   [][]ui.Color
	NumStyles    [][]ui.Style
	Formatter    func(float64) string
	Data         [][]float64
	XLabels      []string
	YLabels      []string
	CellWidth    int
	CellGap      int
}

func NewHeatMap() *HeatMap {
	return &HeatMap{
		Block:        *ui.NewBlock(),
		CellColors:   [][]ui.Color{ui.StandardColors, ui.StandardColors},
		NumStyles:    [][]ui.Style{ui.StandardStyles, ui.StandardStyles},
		Formatter:    func(n float64) string { return fmt.Sprint(n) },
		XLabelStyles: ui.StandardStyles,
		CellGap:      1,
		CellWidth:    3,
	}
}

func (hm *HeatMap) Draw(buffer *ui.Buffer) {
	hm.Block.Draw(buffer)

	cellX := hm.Inner.Min.X

	for i, column := range hm.Data {
		cellY := 0
		for j, datum := range column {
			buffer.SetString(
				hm.YLabels[j],
				ui.StyleClear,
				im.Pt(hm.Inner.Min.X, (hm.Inner.Max.Y-2)-cellY),
			)
			for x := cellX + 5; x < ui.MinInt(cellX+hm.CellWidth, hm.Inner.Max.X)+5; x++ {
				for y := (hm.Inner.Max.Y - 2) - cellY; y > (hm.Inner.Max.Y-2)-cellY-1; y-- {
					cell := ui.NewCell(' ', ui.NewStyle(ui.ColorClear, color(hm.CellColors, i, j)))
					buffer.SetCell(cell, im.Pt(x, y))
				}
			}

			if hm.Formatter != nil {
				hm.drawNumber(buffer, datum, i, j, cellX+5, cellY)
			}

			cellY++
		}

		if i < len(hm.XLabels) {
			hm.drawLabel(buffer, cellX+5, i)
		}

		cellX += hm.CellWidth + hm.CellGap
	}
}

func (hm *HeatMap) drawLabel(buffer *ui.Buffer, cellX, i int) {
	labelX := cellX + ui.MaxInt(int(float64(hm.CellWidth)/2)-int(float64(rw.StringWidth(hm.XLabels[i]))/2), 0)
	buffer.SetString(
		ui.TrimString(hm.XLabels[i], hm.CellWidth),
		ui.SelectStyle(hm.XLabelStyles, i),
		im.Pt(labelX, hm.Inner.Max.Y-1),
	)
}

func (hm *HeatMap) drawNumber(buffer *ui.Buffer, datum float64, i, j, cellX, cellY int) {
	x := cellX + int(float64(hm.CellWidth)/2) - 1
	numberStyle := style(hm.NumStyles, i, j)
	cellColor := color(hm.CellColors, i, j)
	buffer.SetString(
		hm.Formatter(datum),
		ui.NewStyle(numberStyle.Fg, cellColor, numberStyle.Modifier),
		im.Pt(x, (hm.Inner.Max.Y-2)-cellY),
	)
}

func color(colors [][]ui.Color, i, j int) ui.Color {
	return colors[i%len(colors)][j%len(colors)]
}

func style(styles [][]ui.Style, i, j int) ui.Style {
	return styles[i%len(styles)][j%len(styles)]
}