move cost distribution to the accounting package

cmd/pnl.go

@@ -9,8 +9,8 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/c9s/bbgo/cmd/cmdutil"
+	"github.com/c9s/bbgo/pkg/accounting"
 	"github.com/c9s/bbgo/pkg/accounting/pnl"
-	"github.com/c9s/bbgo/pkg/bbgo"
 	"github.com/c9s/bbgo/pkg/service"
 	"github.com/c9s/bbgo/pkg/types"
 )
@@ -95,7 +95,7 @@ var pnlCmd = &cobra.Command{
 
 		log.Infof("%d trades loaded", len(trades))
 
-		stockManager := &bbgo.StockDistribution{
+		stockManager := &accounting.StockDistribution{
 			Symbol:             symbol,
 			TradingFeeCurrency: tradingFeeCurrency,
 		}

pkg/accounting/cost_distribution.go

@@ -0,0 +1,248 @@
+package accounting
+
+import (
+	"fmt"
+	"math"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+
+	"github.com/c9s/bbgo/pkg/types"
+)
+
+func zero(a float64) bool {
+	return int(math.Round(a*1e8)) == 0
+}
+
+func round(a float64) float64 {
+	return math.Round(a*1e8) / 1e8
+}
+
+type Stock types.Trade
+
+func (stock *Stock) String() string {
+	return fmt.Sprintf("%f (%f)", stock.Price, stock.Quantity)
+}
+
+func (stock *Stock) Consume(quantity float64) float64 {
+	q := math.Min(stock.Quantity, quantity)
+	stock.Quantity = round(stock.Quantity - q)
+	return q
+}
+
+type StockSlice []Stock
+
+func (slice StockSlice) QuantityBelowPrice(price float64) (quantity float64) {
+	for _, stock := range slice {
+		if stock.Price < price {
+			quantity += stock.Quantity
+		}
+	}
+
+	return round(quantity)
+}
+
+func (slice StockSlice) Quantity() (total float64) {
+	for _, stock := range slice {
+		total += stock.Quantity
+	}
+
+	return round(total)
+}
+
+type StockDistribution struct {
+	mu sync.Mutex
+
+	Symbol             string
+	TradingFeeCurrency string
+	Stocks             StockSlice
+	PendingSells       StockSlice
+}
+
+type DistributionStats struct {
+	PriceLevels   []string              `json:"priceLevels"`
+	TotalQuantity float64               `json:"totalQuantity"`
+	Quantities    map[string]float64    `json:"quantities"`
+	Stocks        map[string]StockSlice `json:"stocks"`
+}
+
+func (m *StockDistribution) DistributionStats(level int) *DistributionStats {
+	var d = DistributionStats{
+		Quantities: map[string]float64{},
+		Stocks:     map[string]StockSlice{},
+	}
+
+	for _, stock := range m.Stocks {
+		n := math.Ceil(math.Log10(stock.Price))
+		digits := int(n - math.Max(float64(level), 1.0))
+		div := math.Pow10(digits)
+		priceLevel := math.Floor(stock.Price/div) * div
+		key := strconv.FormatFloat(priceLevel, 'f', 2, 64)
+
+		d.TotalQuantity += stock.Quantity
+		d.Stocks[key] = append(d.Stocks[key], stock)
+		d.Quantities[key] += stock.Quantity
+	}
+
+	var priceLevels []float64
+	for priceString := range d.Stocks {
+		price, _ := strconv.ParseFloat(priceString, 32)
+		priceLevels = append(priceLevels, price)
+	}
+	sort.Float64s(priceLevels)
+
+	for _, price := range priceLevels {
+		d.PriceLevels = append(d.PriceLevels, strconv.FormatFloat(price, 'f', 2, 64))
+	}
+
+	sort.Float64s(priceLevels)
+
+	return &d
+}
+
+func (m *StockDistribution) stock(stock Stock) error {
+	m.mu.Lock()
+	m.Stocks = append(m.Stocks, stock)
+	m.mu.Unlock()
+	return m.flushPendingSells()
+}
+
+func (m *StockDistribution) squash() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	var squashed StockSlice
+	for _, stock := range m.Stocks {
+		if !zero(stock.Quantity) {
+			squashed = append(squashed, stock)
+		}
+	}
+	m.Stocks = squashed
+}
+
+func (m *StockDistribution) flushPendingSells() error {
+	if len(m.Stocks) == 0 || len(m.PendingSells) == 0 {
+		return nil
+	}
+
+	pendingSells := m.PendingSells
+	m.PendingSells = nil
+
+	for _, sell := range pendingSells {
+		if err := m.consume(sell); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (m *StockDistribution) consume(sell Stock) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if len(m.Stocks) == 0 {
+		m.PendingSells = append(m.PendingSells, sell)
+		return nil
+	}
+
+	idx := len(m.Stocks) - 1
+	for ; idx >= 0; idx-- {
+		stock := m.Stocks[idx]
+
+		// find any stock price is lower than the sell trade
+		if stock.Price >= sell.Price {
+			continue
+		}
+
+		if zero(stock.Quantity) {
+			continue
+		}
+
+		delta := stock.Consume(sell.Quantity)
+		sell.Consume(delta)
+		m.Stocks[idx] = stock
+
+		if zero(sell.Quantity) {
+			return nil
+		}
+	}
+
+	idx = len(m.Stocks) - 1
+	for ; idx >= 0; idx-- {
+		stock := m.Stocks[idx]
+
+		if zero(stock.Quantity) {
+			continue
+		}
+
+		delta := stock.Consume(sell.Quantity)
+		sell.Consume(delta)
+		m.Stocks[idx] = stock
+
+		if zero(sell.Quantity) {
+			return nil
+		}
+	}
+
+	if sell.Quantity > 0.0 {
+		m.PendingSells = append(m.PendingSells, sell)
+	}
+
+	return nil
+}
+
+func (m *StockDistribution) AddTrades(trades []types.Trade) (checkpoints []int, err error) {
+	feeSymbol := strings.HasPrefix(m.Symbol, m.TradingFeeCurrency)
+	for idx, trade := range trades {
+		// for other market trades
+		// convert trading fee trades to sell trade
+		if trade.Symbol != m.Symbol {
+			if feeSymbol && trade.FeeCurrency == m.TradingFeeCurrency {
+				trade.Symbol = m.Symbol
+				trade.IsBuyer = false
+				trade.Quantity = trade.Fee
+				trade.Fee = 0.0
+			}
+		}
+
+		if trade.Symbol != m.Symbol {
+			continue
+		}
+
+		if trade.IsBuyer {
+			if idx > 0 && len(m.Stocks) == 0 {
+				checkpoints = append(checkpoints, idx)
+			}
+
+			stock := toStock(trade)
+			if err := m.stock(stock); err != nil {
+				return checkpoints, err
+			}
+		} else {
+			stock := toStock(trade)
+			if err := m.consume(stock); err != nil {
+				return checkpoints, err
+			}
+		}
+	}
+
+	err = m.flushPendingSells()
+
+	m.squash()
+
+	return checkpoints, err
+}
+
+func toStock(trade types.Trade) Stock {
+	if strings.HasPrefix(trade.Symbol, trade.FeeCurrency) {
+		if trade.IsBuyer {
+			trade.Quantity -= trade.Fee
+		} else {
+			trade.Quantity += trade.Fee
+		}
+		trade.Fee = 0.0
+	}
+	return Stock(trade)
+}

pkg/bbgo/context.go

@@ -3,6 +3,7 @@ package bbgo
 import (
 	"sync"
 
+	"github.com/c9s/bbgo/pkg/accounting"
 	"github.com/c9s/bbgo/pkg/accounting/pnl"
 	"github.com/c9s/bbgo/pkg/types"
 )
@@ -20,7 +21,7 @@ type Context struct {
 
 	Balances                map[string]types.Balance
 	ProfitAndLossCalculator *pnl.AverageCostCalculator
-	StockManager            *StockDistribution
+	StockManager            *accounting.StockDistribution
 }
 
 func (c *Context) SetCurrentPrice(price float64) {

pkg/bbgo/stock.go

@@ -1,248 +1,2 @@
 package bbgo
 
-import (
-	"fmt"
-	"math"
-	"sort"
-	"strconv"
-	"strings"
-	"sync"
-
-	"github.com/c9s/bbgo/pkg/types"
-)
-
-func zero(a float64) bool {
-	return int(math.Round(a*1e8)) == 0
-}
-
-func round(a float64) float64 {
-	return math.Round(a*1e8) / 1e8
-}
-
-type Stock types.Trade
-
-func (stock *Stock) String() string {
-	return fmt.Sprintf("%f (%f)", stock.Price, stock.Quantity)
-}
-
-func (stock *Stock) Consume(quantity float64) float64 {
-	q := math.Min(stock.Quantity, quantity)
-	stock.Quantity = round(stock.Quantity - q)
-	return q
-}
-
-type StockSlice []Stock
-
-func (slice StockSlice) QuantityBelowPrice(price float64) (quantity float64) {
-	for _, stock := range slice {
-		if stock.Price < price {
-			quantity += stock.Quantity
-		}
-	}
-
-	return round(quantity)
-}
-
-func (slice StockSlice) Quantity() (total float64) {
-	for _, stock := range slice {
-		total += stock.Quantity
-	}
-
-	return round(total)
-}
-
-type StockDistribution struct {
-	mu sync.Mutex
-
-	Symbol             string
-	TradingFeeCurrency string
-	Stocks             StockSlice
-	PendingSells       StockSlice
-}
-
-type DistributionStats struct {
-	PriceLevels   []string              `json:"priceLevels"`
-	TotalQuantity float64               `json:"totalQuantity"`
-	Quantities    map[string]float64    `json:"quantities"`
-	Stocks        map[string]StockSlice `json:"stocks"`
-}
-
-func (m *StockDistribution) DistributionStats(level int) *DistributionStats {
-	var d = DistributionStats{
-		Quantities: map[string]float64{},
-		Stocks:     map[string]StockSlice{},
-	}
-
-	for _, stock := range m.Stocks {
-		n := math.Ceil(math.Log10(stock.Price))
-		digits := int(n - math.Max(float64(level), 1.0))
-		div := math.Pow10(digits)
-		priceLevel := math.Floor(stock.Price/div) * div
-		key := strconv.FormatFloat(priceLevel, 'f', 2, 64)
-
-		d.TotalQuantity += stock.Quantity
-		d.Stocks[key] = append(d.Stocks[key], stock)
-		d.Quantities[key] += stock.Quantity
-	}
-
-	var priceLevels []float64
-	for priceString := range d.Stocks {
-		price, _ := strconv.ParseFloat(priceString, 32)
-		priceLevels = append(priceLevels, price)
-	}
-	sort.Float64s(priceLevels)
-
-	for _, price := range priceLevels {
-		d.PriceLevels = append(d.PriceLevels, strconv.FormatFloat(price, 'f', 2, 64))
-	}
-
-	sort.Float64s(priceLevels)
-
-	return &d
-}
-
-func (m *StockDistribution) stock(stock Stock) error {
-	m.mu.Lock()
-	m.Stocks = append(m.Stocks, stock)
-	m.mu.Unlock()
-	return m.flushPendingSells()
-}
-
-func (m *StockDistribution) squash() {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	var squashed StockSlice
-	for _, stock := range m.Stocks {
-		if !zero(stock.Quantity) {
-			squashed = append(squashed, stock)
-		}
-	}
-	m.Stocks = squashed
-}
-
-func (m *StockDistribution) flushPendingSells() error {
-	if len(m.Stocks) == 0 || len(m.PendingSells) == 0 {
-		return nil
-	}
-
-	pendingSells := m.PendingSells
-	m.PendingSells = nil
-
-	for _, sell := range pendingSells {
-		if err := m.consume(sell); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (m *StockDistribution) consume(sell Stock) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if len(m.Stocks) == 0 {
-		m.PendingSells = append(m.PendingSells, sell)
-		return nil
-	}
-
-	idx := len(m.Stocks) - 1
-	for ; idx >= 0; idx-- {
-		stock := m.Stocks[idx]
-
-		// find any stock price is lower than the sell trade
-		if stock.Price >= sell.Price {
-			continue
-		}
-
-		if zero(stock.Quantity) {
-			continue
-		}
-
-		delta := stock.Consume(sell.Quantity)
-		sell.Consume(delta)
-		m.Stocks[idx] = stock
-
-		if zero(sell.Quantity) {
-			return nil
-		}
-	}
-
-	idx = len(m.Stocks) - 1
-	for ; idx >= 0; idx-- {
-		stock := m.Stocks[idx]
-
-		if zero(stock.Quantity) {
-			continue
-		}
-
-		delta := stock.Consume(sell.Quantity)
-		sell.Consume(delta)
-		m.Stocks[idx] = stock
-
-		if zero(sell.Quantity) {
-			return nil
-		}
-	}
-
-	if sell.Quantity > 0.0 {
-		m.PendingSells = append(m.PendingSells, sell)
-	}
-
-	return nil
-}
-
-func (m *StockDistribution) AddTrades(trades []types.Trade) (checkpoints []int, err error) {
-	feeSymbol := strings.HasPrefix(m.Symbol, m.TradingFeeCurrency)
-	for idx, trade := range trades {
-		// for other market trades
-		// convert trading fee trades to sell trade
-		if trade.Symbol != m.Symbol {
-			if feeSymbol && trade.FeeCurrency == m.TradingFeeCurrency {
-				trade.Symbol = m.Symbol
-				trade.IsBuyer = false
-				trade.Quantity = trade.Fee
-				trade.Fee = 0.0
-			}
-		}
-
-		if trade.Symbol != m.Symbol {
-			continue
-		}
-
-		if trade.IsBuyer {
-			if idx > 0 && len(m.Stocks) == 0 {
-				checkpoints = append(checkpoints, idx)
-			}
-
-			stock := toStock(trade)
-			if err := m.stock(stock); err != nil {
-				return checkpoints, err
-			}
-		} else {
-			stock := toStock(trade)
-			if err := m.consume(stock); err != nil {
-				return checkpoints, err
-			}
-		}
-	}
-
-	err = m.flushPendingSells()
-
-	m.squash()
-
-	return checkpoints, err
-}
-
-func toStock(trade types.Trade) Stock {
-	if strings.HasPrefix(trade.Symbol, trade.FeeCurrency) {
-		if trade.IsBuyer {
-			trade.Quantity -= trade.Fee
-		} else {
-			trade.Quantity += trade.Fee
-		}
-		trade.Fee = 0.0
-	}
-	return Stock(trade)
-}

pkg/bbgo/stock_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/c9s/bbgo/pkg/accounting"
 	"github.com/c9s/bbgo/pkg/types"
 )
 
@@ -21,7 +22,7 @@ func TestStockManager(t *testing.T) {
 		err = json.Unmarshal(out, &trades)
 		assert.NoError(t, err)
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -42,7 +43,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.01, IsBuyer: false},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -50,7 +51,7 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.AddTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 2)
-		assert.Equal(t, StockSlice{
+		assert.Equal(t, accounting.StockSlice{
 			{
 				Symbol:   "BTCUSDT",
 				Price:    9100.0,
@@ -75,7 +76,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.05, IsBuyer: false},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -93,7 +94,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 9200.0, Quantity: 0.05, IsBuyer: false},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -111,7 +112,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 8000.0, Quantity: 0.01, IsBuyer: false},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -119,7 +120,7 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.AddTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 1)
-		assert.Equal(t, StockSlice{
+		assert.Equal(t, accounting.StockSlice{
 			{
 				Symbol:   "BTCUSDT",
 				Price:    9100.0,
@@ -136,7 +137,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -144,7 +145,7 @@ func TestStockManager(t *testing.T) {
 		_, err := stockManager.AddTrades(trades)
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 1)
-		assert.Equal(t, StockSlice{
+		assert.Equal(t, accounting.StockSlice{
 			{
 				Symbol:   "BTCUSDT",
 				Price:    9100.0,
@@ -161,7 +162,7 @@ func TestStockManager(t *testing.T) {
 			{Symbol: "BTCUSDT", Price: 9100.0, Quantity: 0.05, IsBuyer: true},
 		}
 
-		var stockManager = &StockDistribution{
+		var stockManager = &accounting.StockDistribution{
 			TradingFeeCurrency: "BNB",
 			Symbol:             "BTCUSDT",
 		}
@@ -170,7 +171,7 @@ func TestStockManager(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Len(t, stockManager.Stocks, 0)
 		assert.Len(t, stockManager.PendingSells, 1)
-		assert.Equal(t, StockSlice{
+		assert.Equal(t, accounting.StockSlice{
 			{
 				Symbol:   "BTCUSDT",
 				Price:    9200.0,