1 files changed, 48 insertions, 13 deletions

diff --git a/markets/poloniex.go b/markets/poloniex.go
index 34ebb7e..9aaeafb 100644
--- a/markets/poloniex.go
+++ b/markets/poloniex.go
@@ -52,40 +52,75 @@ func NewPoloniex() *Poloniex {
         }
 }
 
-func (p *Poloniex) GetBalance(apiKey, apiSecret string) (map[string]decimal.Decimal, error) {
+func (p *Poloniex) GetBalance(apiKey, apiSecret string) (Summary, error) {
         client, _ := poloniex.NewClient(apiKey, apiSecret)
+        var summary Summary
 
         accounts, err := client.TradeReturnAvailableAccountBalances()
         if poloniexInvalidCredentialsError(err) {
-                return nil, utils.Error{InvalidCredentials, "invalid poloniex credentials"}
+                return Summary{}, utils.Error{InvalidCredentials, "invalid poloniex credentials"}
         }
 
+        if poloniexRestrictedIPError(err) {
+                return Summary{}, utils.Error{IPRestricted, "IP restricted api key"}
+        }
+
+        if err != nil {
+                return Summary{}, err
+        }
+
+        positions, err := client.TradeGetMarginPosition()
         if err != nil {
-                return nil, err
+                return Summary{}, err
         }
 
-        balances := make(map[string]decimal.Decimal)
-        for currency, balance := range accounts.Margin {
-                balances[currency] = balances[currency].Add(balance)
+        marginAccount, err := client.TradeReturnMarginAccountSummary()
+        if err != nil {
+                return Summary{}, err
         }
 
-        for currency, balance := range accounts.Exchange {
-                balances[currency] = balances[currency].Add(balance)
+        summary.Balances = make(map[string]Balance)
+        for currency, amount := range accounts.Exchange {
+                balance := summary.Balances[currency]
+                balance.Amount = balance.Amount.Add(amount)
+
+                summary.Balances[currency] = balance
         }
 
-        return balances, nil
+        summary.BTCValue, err = p.ComputeAccountBalanceValue(summary.Balances)
+        if err != nil {
+                return Summary{}, err
+        }
+
+        for currencyPair, position := range positions {
+                if position.Type == "none" {
+                        continue
+                }
+                currency := currencyPair[4:]
+                balance := summary.Balances[currency]
+                balance.Amount = balance.Amount.Add(position.Amount)
+                balance.BTCValue = balance.BTCValue.Add(position.Total.Add(position.PlusValue))
+
+                summary.Balances[currency] = balance
+        }
+
+        summary.BTCValue = summary.BTCValue.Add(marginAccount.NetValue)
+
+        return summary, nil
 }
 
-func (p *Poloniex) ComputeAccountBalanceValue(account map[string]decimal.Decimal, baseCurrency string) (decimal.Decimal, error) {
+func (p *Poloniex) ComputeAccountBalanceValue(account map[string]Balance) (decimal.Decimal, error) {
         var total decimal.Decimal
 
-        for currency, amount := range account {
-                pair, err := p.GetCurrencyPair(baseCurrency, currency)
+        for currency, balance := range account {
+                pair, err := p.GetCurrencyPair("BTC", currency)
                 if err != nil {
                         return decimal.Zero, err
                 }
 
-                total = total.Add(amount.Mul(pair.Rate))
+                total = total.Add(balance.Amount.Mul(pair.Rate))
+                balance.BTCValue = balance.BTCValue.Add(balance.Amount.Mul(pair.Rate))
+                account[currency] = balance
         }
 
         return total, nil