sig-pub/pkg/indicator/super_trend.go

package indicator

import "sig-pub/pkg/types"

// SuperTrend 超级趋势
type SuperTrend struct {
}

func (c SuperTrend) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "SuperTrend",
		Input: []types.InputArg{
			{Name: "window", Type: types.InputTypeUInt, Desc: "ATR周期(7/14)"},
			{Name: "mul", Type: types.InputTypeUInt, Desc: "乘数(建议2-4)"},
		},
		State: []string{"direction"},
		Plots: []Plot{
			{State: "vector", Type: PlotLine, Props: PlotProps{"color": ColorGreen}, Exps: []PlotExp{
				{Exp: "direction == -1", Props: PlotProps{"color": ColorRed}},
				{Exp: "direction == 1", Props: PlotProps{"color": ColorGreen}},
			}},
		},
	}
}

func (c SuperTrend) CandlePeriods(ctx IIndicatorContext) int16 {
	return ctx.Indicator("ATR", ctx.Input().Int16("window")).CandlePeriods()
}

func (c SuperTrend) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	mul := ctx.Input().Float("mul")

	atr := ctx.Indicator("ATR", window).Get(0)
	hl2 := ctx.Get(0).HL2()
	closeP := ctx.Get(0).CloseF64()

	upper := hl2 + mul*atr // 潛在上漲時的阻力位
	lower := hl2 - mul*atr // 潛在下跌時的支撐位

	prevTrend, ok := ctx.State().Get("_trend", 1)
	prevDirection, _ := ctx.State().Get("direction", 1) // 方向: 1.up, -1.down

	// 1.初始化
	if !ok {
		if closeP > upper {
			prevTrend = lower
			prevDirection = 1
		} else {
			prevTrend = upper
			prevDirection = -1
		}
		ctx.State().Set("_trend", prevTrend)
		ctx.State().Set("direction", prevDirection)
		return prevTrend
	}

	// 2.迭代计算
	trend, direction := prevTrend, prevDirection
	if prevDirection == 1 {
		// uptrend
		trend = max(lower, prevTrend)
		if closeP < prevTrend {
			trend = upper  // 取上轨作为新红线
			direction = -1 // 转下跌趋势
		}
	} else {
		// downtrend
		trend = min(upper, prevTrend)
		if closeP > prevTrend {
			trend = lower // 取下轨作为新绿线
			direction = 1 // 转上升趋势
		}
	}
	ctx.State().Set("_trend", trend)
	ctx.State().Set("direction", direction)
	return trend
}