package com

import (
	"github.com/banbox/banbot/strat"
	ta "github.com/banbox/banta"
	"math"
)

type TRS struct {
	MaxLow   float64
	MinHigh  float64
	SLLong   float64 // 止损的tr倍率，有持仓时，逐渐减小
	SLShort  float64 // 止损的tr倍率，有持仓时，逐渐减小
	InitSLTR float64
	minSLTR  float64
	stepSLTR float64
}

func NewTRS(initSLTR, minRate, costStep float64) *TRS {
	var minSLTR = initSLTR * minRate
	var stepSLTR = (initSLTR - minSLTR) / costStep
	return &TRS{
		InitSLTR: initSLTR,
		minSLTR:  minSLTR,
		stepSLTR: stepSLTR,
	}
}

func (m *TRS) TrialLong(price float64) {
	m.SLLong = m.InitSLTR
	m.MaxLow = price
}

func (m *TRS) TrialShort(price float64) {
	m.SLShort = m.InitSLTR
	m.MinHigh = price
}

func (m *TRS) setTRS(e *ta.BarEnv, longNum, shortNum int) {
	h, l := e.High.Get(0), e.Low.Get(0)
	if longNum == 0 {
		if m.SLLong != m.InitSLTR {
			m.MaxLow = 0
		}
	} else {
		m.MaxLow = max(m.MaxLow, l)
		if m.SLLong > m.minSLTR+0.01 {
			m.SLLong -= m.stepSLTR
		}
	}
	if shortNum == 0 {
		if m.SLShort != m.InitSLTR {
			m.MinHigh = 0
		}
	} else {
		m.MinHigh = min(m.MinHigh, h)
		if m.SLShort > m.minSLTR+0.01 {
			m.SLShort -= m.stepSLTR
		}
	}
}

func (m *TRS) Calc(s *strat.StratJob, atr float64) (float64, float64) {
	m.setTRS(s.Env, len(s.LongOrders), len(s.ShortOrders))
	var stopLong, stopShort float64
	if m.MaxLow > 0 {
		stopLong = m.MaxLow - atr*m.SLLong
	}
	if m.MinHigh > 0 {
		stopShort = m.MinHigh + atr*m.SLShort
	}
	return stopLong, stopShort
}

func (m *TRS) CalcSL(s *strat.StratJob, atr *ta.Series, lenAtr, lenAtrB int, pow, rate float64) (float64, float64) {
	atrBase := ta.Lowest(ta.Highest(atr, lenAtr), lenAtrB).Get(0)
	atrRate := math.Pow(atr.Get(0)/atrBase, pow)
	return m.Calc(s, atrBase*atrRate*rate)
}

func (m *TRS) CalcSL2(s *strat.StratJob, lenAtr, lenAtrB int, pow, rate float64) (float64, float64) {
	e := s.Env
	atr := ta.ATR(e.High, e.Low, e.Close, lenAtr)
	return m.CalcSL(s, atr, lenAtr, lenAtrB, pow, rate)
}