1 files changed, 39 insertions, 10 deletions

diff --git a/pkg/plot/barchart.go b/pkg/plot/barchart.go
index 6bcacd9..04a36be 100644
--- a/pkg/plot/barchart.go
+++ b/pkg/plot/barchart.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"math"
 
+	"github.com/sgreben/jp/pkg/data"
 	"github.com/sgreben/jp/pkg/draw"
 )
 
@@ -22,10 +23,10 @@ func NewBarChart(canvas draw.Canvas) *BarChart {
 }
 
 // Draw implements Chart
-func (c *BarChart) Draw(data *DataTable) string {
+func (c *BarChart) Draw(table *data.Table) string {
 	minY := math.Inf(1)
 	maxY := math.Inf(-1)
-	for _, row := range data.Rows {
+	for _, row := range table.Rows {
 		for _, y := range row {
 			if y < minY {
 				minY = y
@@ -35,12 +36,31 @@ func (c *BarChart) Draw(data *DataTable) string {
 			}
 		}
 	}
-	paddingX := 4
+	paddingX := 2
 	paddingY := 3
 	chartHeight := c.Size().Height - paddingY*c.RuneSize().Height
 	chartWidth := c.Size().Width - 2*paddingX*c.RuneSize().Width
+
+	labelsBelowBars := true
+	labelsRight := false
+	maxLabelLength := 0
+	totalLabelLength := 0
+	for _, group := range table.Columns {
+		totalLabelLength += len(group)
+		if len(group) > maxLabelLength {
+			maxLabelLength = len(group)
+		}
+	}
+	if totalLabelLength*c.RuneSize().Width > chartWidth {
+		labelsBelowBars = false
+		if len(table.Columns)*c.RuneSize().Height <= chartHeight {
+			labelsRight = true
+			chartWidth -= 3 + maxLabelLength*c.RuneSize().Width
+		}
+	}
+
 	scaleY := float64(chartHeight) / maxY
-	barPaddedWidth := chartWidth / len(data.Columns)
+	barPaddedWidth := chartWidth / len(table.Columns)
 	barWidth := barPaddedWidth - (c.BarPaddingX * c.RuneSize().Width)
 	if barPaddedWidth < c.RuneSize().Width {
 		barPaddedWidth = c.RuneSize().Width
@@ -51,11 +71,11 @@ func (c *BarChart) Draw(data *DataTable) string {
 
 	scaleY = float64(chartHeight) / maxY
 
-	for i, group := range data.Columns {
+	for i, group := range table.Columns {
 		barLeft := paddingX*c.RuneSize().Width + barPaddedWidth*i
 		barRight := barLeft + barWidth
 
-		y := data.Rows[0][i]
+		y := table.Rows[0][i]
 		barHeight := y * scaleY
 		barBottom := (paddingY - 1) * c.RuneSize().Height
 		barTop := barBottom + int(barHeight)
@@ -66,20 +86,29 @@ func (c *BarChart) Draw(data *DataTable) string {
 			}
 		}
 
-		// Group label
 		barMiddle := int(math.Floor(float64(barLeft+barRight) / float64(2*c.RuneSize().Width)))
-		c.GetBuffer().WriteCenter(0, barMiddle, []rune(group))
+
+		// Group label
+		if labelsBelowBars {
+			c.GetBuffer().WriteCenter(0, barMiddle, []rune(group))
+		} else {
+			c.GetBuffer().WriteCenter(0, barMiddle, []rune(Fi(i)))
+		}
 
 		// Count label
 		countLabelY := int(math.Ceil(float64(barTop)/float64(c.RuneSize().Height))) * c.RuneSize().Height
-
 		if countLabelY <= barBottom && y > 0 {
 			c.GetBuffer().SetRow(barTop/c.RuneSize().Height, barLeft/c.RuneSize().Width, barRight/c.RuneSize().Width, '▁')
 			countLabelY = 3 * c.RuneSize().Height
 		}
-
 		c.GetBuffer().WriteCenter(countLabelY/c.RuneSize().Height, barMiddle, Ff(y))
 	}
+	if labelsRight {
+		for i, group := range table.Columns {
+			c.GetBuffer().WriteRight(c.GetBuffer().Height-i, paddingX+1+chartWidth/c.RuneSize().Width, []rune(Fi(i)))
+			c.GetBuffer().WriteRight(c.GetBuffer().Height-i, paddingX+4+chartWidth/c.RuneSize().Width, []rune(group))
+		}
+	}
 
 	b := bytes.NewBuffer(nil)
 	c.GetBuffer().Render(b)