1package chart23import (4 "image"5 "math"67 "github.com/charmbracelet/x/ansi"89 "src.rest/cat/012/internal/sheet"10)1112// columnPlan lays out column, line and area charts: a value axis on the13// left with labeled ticks, the plot, and a category axis and labels14// below it. Draw adds the legend.15//16// 1,500 ┤ █▆17// 1,000 ┤ ██ ▆18// 500 ┤ ██ █19// 0 ┼─▀▀─▀────20// Jan Feb21type columnPlan struct {22 d sheet.ChartData23 kind columnKind24 stacked bool25 grid bool26 msg string27 sc scale28 labelW int // width of the tick labels29 axisX int // column of the value axis30 axisRow int // row of the category axis31 catRow int32 plot image.Rectangle33 shown int // categories drawn34 // Columns: each category is a group gw cells wide, starting off cells35 // into the plot; each series' bar in it (the stack, when stacked) is36 // bw wide.37 gw, off, bw, inset int38}3940// columnKind is what a columnPlan draws at each category.41type columnKind uint84243const (44 columnBars columnKind = iota45 columnLine46 columnArea47)4849func newColumnPlan(d sheet.ChartData, w, h int, kind columnKind, o sheet.ChartOptions) *columnPlan {50 stacked := o.Stack != sheet.StackNone && kind != columnLine51 d = prepare(d, o, kind != columnLine)52 p := &columnPlan{d: d, kind: kind, stacked: stacked, grid: !o.NoGrid}53 lo, hi, ok := finite(d.Series)54 if stacked {55 lo, hi, ok = stackedRange(d, len(d.Categories))56 }57 if !ok || len(d.Categories) == 0 {58 p.msg = "No numbers to chart"59 return p60 }61 if kind != columnLine && !o.Log {62 lo, hi = min(lo, 0), max(hi, 0)63 }64 p.catRow = h - 165 p.axisRow = p.catRow - 166 if p.axisRow < 2 {67 p.msg = "Too small to chart"68 return p69 }70 p.sc = newAxis(lo, hi, p.axisRow, 2, 6, d.Format, o)71 for j := 0; j <= p.sc.n; j++ {72 p.labelW = max(p.labelW, ansi.StringWidth(p.sc.label(p.sc.tick(j))))73 }74 p.axisX = p.labelW + 175 p.plot = image.Rect(p.axisX+1, 0, w-1, p.axisRow)76 plotW := p.plot.Dx()77 if plotW < 2 {78 p.msg = "Too small to chart"79 return p80 }81 if kind != columnBars {82 p.shown = min(len(d.Categories), plotW*2)83 return p84 }85 slots := len(d.Series)86 if stacked {87 slots = 188 }89 p.shown = max(min(len(d.Categories), plotW/slots), 1)90 p.gw = plotW / p.shown91 gap := 092 if p.gw > slots {93 gap = max(1, p.gw/4)94 }95 p.bw = max((p.gw-gap)/slots, 1)96 p.inset = (p.gw - p.bw*slots) / 297 p.off = (plotW - p.shown*p.gw) / 298 return p99}100101// bar returns the columns of series j's bar in category i; stacked, all102// series share the first.103func (p *columnPlan) bar(i, j int) (x0, x1 int) {104 if p.stacked {105 j = 0106 }107 x0 = p.plot.Min.X + p.off + i*p.gw + p.inset + j*p.bw108 return x0, x0 + p.bw109}110111// piles returns category i's bars as piles of spans in values, each with112// the series whose slot it takes: one pile when stacked, one per series113// otherwise.114func (p *columnPlan) piles(i int, buf []span) []span {115 if p.stacked {116 return stack(p.d, i, buf)117 }118 buf = buf[:0]119 base := p.sc.base()120 for j, s := range p.d.Series {121 if v := value(s, i); !math.IsNaN(v) {122 buf = append(buf, span{min(v, base), max(v, base), j})123 }124 }125 return buf126}127128// dotX returns the x of category i's point in braille dots (two per cell).129func (p *columnPlan) dotX(i int) int {130 dots := p.plot.Dx() * 2131 if p.shown <= 1 {132 return dots / 2133 }134 return int(math.Round(float64(i) * float64(dots-1) / float64(p.shown-1)))135}136137// catCenter is the column under which category i's label is centered.138func (p *columnPlan) catCenter(i int) int {139 if p.kind != columnBars {140 return p.plot.Min.X + p.dotX(i)/2141 }142 return p.plot.Min.X + p.off + i*p.gw + p.gw/2143}144145func (p *columnPlan) draw(g *Grid, o Options) {146 drawValueAxis(g, p.sc, p.labelW, p.axisX, p.axisRow, p.plot.Min.X, p.plot.Max.X+1)147 p.drawCategories(g)148 if o.Image {149 return150 }151 if p.grid {152 hGridlines(g, p.sc, p.axisRow, p.plot)153 }154 switch p.kind {155 case columnLine:156 p.drawLines(g)157 case columnArea:158 p.drawArea(g)159 default:160 p.drawBars(g)161 }162}163164// drawValueAxis draws a vertical value axis at column axisX with a label165// at each tick, and the category axis along row axisRow from x0 to x1.166func drawValueAxis(g *Grid, sc scale, labelW, axisX, axisRow, x0, x1 int) {167 for y := axisRow - sc.cells(); y < axisRow; y++ {168 g.set(axisX, y, "│", Axis)169 }170 for x := x0; x < min(x1, g.W); x++ {171 g.set(x, axisRow, "─", Axis)172 }173 for j := 0; j <= sc.n; j++ {174 y := axisRow - j*sc.k175 label := sc.label(sc.tick(j))176 g.text(labelW-ansi.StringWidth(label), y, label, Label)177 mark := "┤"178 if j == 0 {179 mark = "┼"180 }181 g.set(axisX, y, mark, Axis)182 }183}184185// hGridlines draws a dashed line across the plot at each tick above the186// axis.187func hGridlines(g *Grid, sc scale, axisRow int, plot image.Rectangle) {188 for j := 1; j <= sc.n; j++ {189 y := axisRow - j*sc.k190 for x := plot.Min.X; x < plot.Max.X; x++ {191 g.set(x, y, "┈", Gridline)192 }193 }194}195196// drawCategories labels as many categories as fit without touching.197func (p *columnPlan) drawCategories(g *Grid) {198 room := p.gw - 1199 if p.kind != columnBars {200 room = p.plot.Dx()/max(p.shown-1, 1) - 1201 }202 room = max(room, 3)203 next := 0 // first free column204 for i := 0; i < p.shown; i++ {205 label := ansi.Truncate(p.d.Categories[i], room, "…")206 w := ansi.StringWidth(label)207 x := min(max(p.catCenter(i)-(w-1)/2, 0), g.W-w)208 if x < next {209 continue210 }211 g.text(x, p.catRow, label, Label)212 next = x + w + 1213 }214}215216func (p *columnPlan) drawBars(g *Grid) {217 drawPiles(p.sc, p.shown, p.stacked, p.piles, lowerEighths, "▀", func(i, j, r int, cell Cell) {218 x0, x1 := p.bar(i, j)219 for x := x0; x < x1; x++ {220 if at := g.At(x, p.axisRow-1-r); at != nil {221 *at = cell222 }223 }224 })225}226227// dotY maps a value to a braille dot row, four per cell, from the top:228// a value on a tick is on the tick's line, through the middle of its229// row, as bars are (see inCells), and the axis's value on the lowest230// dots above the axis row.231func (p *columnPlan) dotY(v float64) int {232 rows := p.axisRow * 4233 y := int(math.Floor((float64(p.axisRow) + 0.5 - p.sc.at(v)) * 4))234 return min(max(y, 0), rows-1)235}236237// drawLines draws each series as a braille line; where lines cross a238// cell, the later series' color wins.239func (p *columnPlan) drawLines(g *Grid) {240 b := newBraille(p.plot.Dx(), p.axisRow)241 for j, s := range p.d.Series {242 px, py, have := 0, 0, false243 for i := 0; i < p.shown && i < len(s.Values); i++ {244 v := s.Values[i]245 if math.IsNaN(v) {246 have = false247 continue248 }249 x, y := p.dotX(i), p.dotY(v)250 if have {251 b.line(px, py, x, y, j)252 } else {253 b.dot(x, y, j)254 }255 px, py, have = x, y, true256 }257 }258 b.draw(g, p.plot.Min.X, 0)259}260261// braille is a canvas of braille dots, two by four per cell, each cell262// colored by the series that drew in it last.263type braille struct {264 w, h int265 bits []uint8266 owner []int267}268269func newBraille(w, h int) *braille {270 return &braille{w: w, h: h, bits: make([]uint8, w*h), owner: make([]int, w*h)}271}272273func (b *braille) dot(x, y, series int) {274 cx, cy := x/2, y/4275 if x < 0 || y < 0 || cx >= b.w || cy >= b.h {276 return277 }278 b.bits[cy*b.w+cx] |= brailleBit[x%2][y%4]279 b.owner[cy*b.w+cx] = series280}281282func (b *braille) line(x0, y0, x1, y1, series int) {283 bresenham(x0, y0, x1, y1, func(x, y int) { b.dot(x, y, series) })284}285286// draw puts the dots on g with the top-left cell at x, y.287func (b *braille) draw(g *Grid, x, y int) {288 for cy := range b.h {289 for cx := range b.w {290 if bits := b.bits[cy*b.w+cx]; bits != 0 {291 g.set(x+cx, y+cy, string(rune(0x2800+int(bits))), SeriesRole(b.owner[cy*b.w+cx]))292 }293 }294 }295}296297// brailleBit is the bit of the braille dot at column x, row y of a cell.298var brailleBit = [2][4]uint8{{0x01, 0x02, 0x04, 0x40}, {0x08, 0x10, 0x20, 0x80}}299300func bresenham(x0, y0, x1, y1 int, plot func(x, y int)) {301 dx, dy := abs(x1-x0), -abs(y1-y0)302 sx, sy := sign(x1-x0), sign(y1-y0)303 e := dx + dy304 for {305 plot(x0, y0)306 if x0 == x1 && y0 == y1 {307 return308 }309 e2 := 2 * e310 if e2 >= dy {311 e += dy312 x0 += sx313 }314 if e2 <= dx {315 e += dx316 y0 += sy317 }318 }319}320321func abs(n int) int { return max(n, -n) }322323func sign(n int) int {324 switch {325 case n > 0:326 return 1327 case n < 0:328 return -1329 }330 return 0331}