package main

import (
	"fmt"
	"git.beejay.kim/Gshopper/sentio"
	"git.beejay.kim/Gshopper/sentio/indicator"
	"time"
)

const (
	ATR_MULTIPLIER         = 1.5
	ATR_PERIOD             = 5
	ATR_HHV                = 4
	ATR_STOPLOSS_THRESHOLD = 1

	EXTRA_POSITION_THRESHOLD = 1.002
)

var Strategy = alpacaQQQ{}

type alpacaQQQ struct{}

func (s alpacaQQQ) Name() string {
	return "Alpaca: QQQ [TQQQ : SQQQ]"
}

func (s alpacaQQQ) Model() string {
	return "qqq15"
}

func (s alpacaQQQ) MarketId() string {
	return "alpaca"
}

func (s alpacaQQQ) PositionSymbols() map[sentio.Side]string {
	return map[sentio.Side]string{
		sentio.BASE:  "QQQ",
		sentio.LONG:  "TQQQ",
		sentio.SHORT: "SQQQ",
	}
}

func (s alpacaQQQ) Interval() uint8 {
	return 5
}

func (s alpacaQQQ) Cooldown(periods uint8) time.Duration {
	return time.Minute * time.Duration(s.Interval()*periods)
}

func (s alpacaQQQ) Handle(market sentio.Market, ts time.Time, proba float64) ([]sentio.StrategyOrder, error) {
	var (
		portfolio sentio.Portfolio
		orders    []sentio.StrategyOrder
		now       = market.Time().Now()
		err       error
	)

	// skip too early trades
	if now.Hour() == 9 && now.Minute() < 45 {
		return orders, nil
	}

	if portfolio, err = market.Portfolio(); err != nil {
		return nil, err
	}

	for side, symbol := range s.PositionSymbols() {
		var (
			position   sentio.Position
			quote      *sentio.Quote
			bars       []sentio.Bar
			stoploss   float64
			uptrending bool
			ok         bool
		)

		// no need to trade BASE quote
		if sentio.BASE == side {
			continue
		}

		if quote, err = market.Quote(symbol); err != nil {
			return nil, err
		}

		fmt.Printf("Quotes: %s [BID: %f] [ASK: %s]\n", symbol, quote.BidPrice, quote.AskPrice)

		if bars, err = market.HistoricalBars(symbol, time.Minute, nil); err == nil && len(bars) > 0 {
			h := make([]float64, len(bars))
			l := make([]float64, len(bars))
			c := make([]float64, len(bars))

			for i := range bars {
				h[i] = bars[i].High
				l[i] = bars[i].Low
				c[i] = bars[i].Close
			}

			trailing := indicator.AtrTrailingStopLoss(h, l, c, ATR_PERIOD, ATR_MULTIPLIER, ATR_HHV)
			stoploss = trailing[len(trailing)-1]
			uptrending = trailing[len(trailing)-1] > trailing[len(trailing)-2]

			fmt.Printf("ATR Trailing StopLoss: [%s: %f; uptrading: %v]\n", symbol, stoploss, uptrending)
		}

		if portfolio != nil {
			position, ok = portfolio.Get(symbol)
			ok = ok && position != nil && position.GetSize() != 0
		} else {
			ok = false
		}

		// Close positions before market closed
		if ok && now.Hour() == 15 && now.Minute() > 55 {
			return []sentio.StrategyOrder{{
				Symbol: symbol,
				Action: sentio.OrderSell,
				Ratio:  1,
			}}, nil
		}

		// Close position if BidPrice less than StopLoss
		if ok && !uptrending && quote.BidPrice/stoploss < ATR_STOPLOSS_THRESHOLD {
			return []sentio.StrategyOrder{{
				Symbol: symbol,
				Action: sentio.OrderSell,
				Ratio:  1,
			}}, nil
		}

		if proba < 0 {
			continue
		}

		// Close LONG position
		if ok && sentio.LONG == side && proba < 1.0008 {
			return []sentio.StrategyOrder{{
				Symbol: symbol,
				Action: sentio.OrderSell,
				Ratio:  1,
			}}, nil
		}

		// Close SHORT position
		if ok && sentio.SHORT == side && proba > .9998 {
			return []sentio.StrategyOrder{{
				Symbol: symbol,
				Action: sentio.OrderSell,
				Ratio:  1,
			}}, nil
		}

		// Prevent BUYs for delayed probas
		if ts.Add(time.Minute * time.Duration(int64(s.Interval()/2))).Before(now.Round(time.Minute)) {
			continue
		}

		// Prevent BUYs on closing market
		if now.Hour() == 15 && now.Minute() > 46 {
			continue
		}

		if quote.BidPrice/stoploss < ATR_STOPLOSS_THRESHOLD {
			continue
		}

		if sentio.LONG == side && proba > 1 {
			size := float64(0)

			if ok && position.GetAvgPrice()/position.GetCurrentPrice() > EXTRA_POSITION_THRESHOLD {
				// extra position with cheaper price
				size = .5
			} else if !ok && proba > 1.0008 {
				// new trade position
				size = .7
			}

			if size > 0 {
				orders = append(orders, sentio.StrategyOrder{
					Symbol: symbol,
					Action: sentio.OrderBuy,
					Ratio:  size,
				})
			}

			continue
		}

		if sentio.SHORT == side && proba < 1 {
			size := float64(0)

			if ok && position.GetAvgPrice()/position.GetCurrentPrice() > EXTRA_POSITION_THRESHOLD {
				// extra position with cheaper price
				size = .5
			} else if !ok && proba < .9990 {
				// new trade position
				size = .7
			}

			if size > 0 {
				orders = append(orders, sentio.StrategyOrder{
					Symbol: symbol,
					Action: sentio.OrderBuy,
					Ratio:  size,
				})
			}

			continue
		}
	}

	return orders, nil
}