From 931ce434ef102923aac28b11db3daa8a1fb0c42a Mon Sep 17 00:00:00 2001
From: c9s <yoan@maicoin.com>
Date: Mon, 12 Oct 2020 17:07:50 +0800
Subject: [PATCH] improve baseline calculator

---
 cmd/transfers.go | 75 +++++++++++++++++++++++++++++-------------------
 1 file changed, 46 insertions(+), 29 deletions(-)

diff --git a/cmd/transfers.go b/cmd/transfers.go
index b541b1aea..088b0daa5 100644
--- a/cmd/transfers.go
+++ b/cmd/transfers.go
@@ -13,19 +13,17 @@ import (
 
 func init() {
 	transferHistoryCmd.Flags().String("exchange", "", "target exchange")
-	transferHistoryCmd.Flags().String("asset", "BTC", "trading symbol")
+	transferHistoryCmd.Flags().String("asset", "", "trading symbol")
 	transferHistoryCmd.Flags().String("since", "", "since time")
 	RootCmd.AddCommand(transferHistoryCmd)
 }
 
-
-
-type TimeRecord struct {
+type timeRecord struct {
 	Record interface{}
-	Time time.Time
+	Time   time.Time
 }
 
-type timeSlice []TimeRecord
+type timeSlice []timeRecord
 
 func (p timeSlice) Len() int {
 	return len(p)
@@ -39,13 +37,9 @@ func (p timeSlice) Swap(i, j int) {
 	p[i], p[j] = p[j], p[i]
 }
 
-
-
-
-
 var transferHistoryCmd = &cobra.Command{
-	Use:          "transfer-history",
-	Short:        "show transfer history",
+	Use:   "transfer-history",
+	Short: "show transfer history",
 
 	SilenceUsage: true,
 	RunE: func(cmd *cobra.Command, args []string) error {
@@ -89,7 +83,6 @@ var transferHistoryCmd = &cobra.Command{
 			}
 		}
 
-
 		exchange := newExchange(exchangeName)
 
 		var records timeSlice
@@ -99,7 +92,7 @@ var transferHistoryCmd = &cobra.Command{
 			return err
 		}
 		for _, d := range deposits {
-			records = append(records, TimeRecord{
+			records = append(records, timeRecord{
 				Record: d,
 				Time:   d.EffectiveTime(),
 			})
@@ -110,7 +103,7 @@ var transferHistoryCmd = &cobra.Command{
 			return err
 		}
 		for _, w := range withdraws {
-			records = append(records, TimeRecord{
+			records = append(records, timeRecord{
 				Record: w,
 				Time:   w.EffectiveTime(),
 			})
@@ -134,39 +127,63 @@ var transferHistoryCmd = &cobra.Command{
 		}
 
 		stats := calBaselineStats(asset, deposits, withdraws)
-		log.Infof("total %s deposit: %f (x %d)", asset, stats.TotalDeposit, stats.NumOfDeposit)
-		log.Infof("total %s withdraw: %f (x %d)", asset, stats.TotalWithdraw, stats.NumOfWithdraw)
-		log.Infof("baseline %s balance: %f", asset, stats.BaselineBalance)
+		for asset, quantity := range stats.TotalDeposit {
+			log.Infof("total %s deposit: %f", asset, quantity)
+		}
+
+		for asset, quantity := range stats.TotalWithdraw {
+			log.Infof("total %s withdraw: %f", asset, quantity)
+		}
+
+		for asset, quantity := range stats.BaselineBalance {
+			log.Infof("baseline %s balance: %f", asset, quantity)
+		}
+
 		return nil
 	},
 }
 
 type BaselineStats struct {
-	Asset string
-	NumOfDeposit int
-	NumOfWithdraw int
-	TotalDeposit float64
-	TotalWithdraw float64
-	BaselineBalance float64
+	Asset           string
+	TotalDeposit    map[string]float64
+	TotalWithdraw   map[string]float64
+	BaselineBalance map[string]float64
 }
 
 func calBaselineStats(asset string, deposits []types.Deposit, withdraws []types.Withdraw) (stats BaselineStats) {
 	stats.Asset = asset
-	stats.NumOfDeposit = len(deposits)
-	stats.NumOfWithdraw = len(withdraws)
+	stats.TotalDeposit = make(map[string]float64)
+	stats.TotalWithdraw = make(map[string]float64)
+	stats.BaselineBalance = make(map[string]float64)
 
 	for _, deposit := range deposits {
 		if deposit.Status == types.DepositSuccess {
-			stats.TotalDeposit += deposit.Amount
+			if _, ok := stats.TotalDeposit[deposit.Asset]; !ok {
+				stats.TotalDeposit[deposit.Asset] = 0.0
+			}
+
+			stats.TotalDeposit[deposit.Asset] += deposit.Amount
 		}
 	}
 
 	for _, withdraw := range withdraws {
 		if withdraw.Status == "completed" {
-			stats.TotalWithdraw += withdraw.Amount
+			if _, ok := stats.TotalWithdraw[withdraw.Asset]; !ok {
+				stats.TotalWithdraw[withdraw.Asset] = 0.0
+			}
+
+			stats.TotalWithdraw[withdraw.Asset] += withdraw.Amount
 		}
 	}
 
-	stats.BaselineBalance = stats.TotalDeposit - stats.TotalWithdraw
+	for asset, deposit := range stats.TotalDeposit {
+		withdraw, ok := stats.TotalWithdraw[asset]
+		if !ok {
+			withdraw = 0.0
+		}
+
+		stats.BaselineBalance[asset] = deposit - withdraw
+	}
+
 	return stats
 }