package indicator

import (
	"sig-pub/pkg/types"
	"sig-pub/pkg/types/series"
)

// ATR = SMA(TR, N)
// 平均真实波幅 (ATR) atr define: https://www.investopedia.com/terms/a/atr.asp
type ATR struct {
}

// indicator interface
func (c *ATR) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "ATR",
		Input: []types.InputArg{
			{Name: "window", Type: types.InputTypeUInt, Desc: "窗口大小"},
		},
	}
}

func (c *ATR) CandlePeriods(ctx IIndicatorContext) int16 {
	return ctx.Input().Int16("window") + 1
}

// Calculate 计算单根k线rsi指标
func (c *ATR) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	klineSeries := ctx.Series(0, int16(window)+1)
	highs := klineSeries.High()
	lows := klineSeries.Low()
	closes := klineSeries.Close()

	trs := make([]float64, 0, window)
	for i := range window {
		high := highs[i]
		low := lows[i]
		close1 := closes[i+1]
		// 计算TR
		tr := max(high-low, high-close1, low-close1)
		trs = append(trs, tr)
	}
	// 计算TR平均值得到ATR
	seriesTR := series.NewFloats(trs...)
	atr := seriesTR.Avg()
	return atr
}