sig-pub/internal/trading/sig/indicator_context.go

package sig

import (
	"fmt"
	"sig-pub/pkg/indicator"
	"sig-pub/pkg/types"
	"sig-pub/pkg/types/series"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/zlog"
	"slices"
	"strings"

	"github.com/spf13/cast"
)

type IndicatorStates map[string]*IndicatorState

func NewIndicatorStates() IndicatorStates {
	return make(IndicatorStates)
}

type IOffsetIndicatorContext interface {
	indicator.IIndicatorContext
	SetOffset(offset int16)
	AddOffset(offset int16)
	GetOffset() (offset int16)
}

// IndicatorContext 指标上下文, 提供k线序列给指标计算使用
type IndicatorContext struct {
	IOffsetIndicatorContext
	indicator       indicator.IIndicator
	indicatorsReg   *indicator.IndicatorRegistry
	input           types.Input
	indicatorStates map[string]*IndicatorState // {macd{window:0,fast:9,slow:21,single:10}: state} 初始化时与KlineSeries周期同步
	kSeries         *KlineSeries
	offset          int16
	indicatorTrace  []string // 指标调用链避免指标循环引用
}

func NewIndicatorContext(indicator indicator.IIndicator, input types.Input, indicatorStates IndicatorStates, kSeries *KlineSeries, indicatorsReg *indicator.IndicatorRegistry) *IndicatorContext {
	return &IndicatorContext{
		indicator:       indicator,
		input:           input,
		indicatorStates: indicatorStates,
		kSeries:         kSeries,
		indicatorsReg:   indicatorsReg,
		indicatorTrace:  []string{indicator.Meta().Name},
	}
}

func (c *IndicatorContext) SetOffset(offset int16) {
	c.offset = offset
}

func (c *IndicatorContext) AddOffset(offset int16) {
	c.offset += offset
}

func (c *IndicatorContext) GetOffset() (offset int16) {
	return c.offset
}

func (c *IndicatorContext) Input() (in types.Input) {
	return c.input
}

func (c *IndicatorContext) Get(offset int16) (kline types.Kline) {
	offset += c.offset
	k, err := c.kSeries.Get(offset)
	if err != nil {
		lastTs := c.kSeries.LastTs()
		zlog.Warningf("get kline series offset out of range: offset=%d, length=%d, lastTs=%d", offset, c.kSeries.Length(), lastTs)
		panic(err)
	}
	return k
}

func (c *IndicatorContext) Series(offset, count int16) (klines series.Klines) {
	offset += c.offset
	ks, err := c.kSeries.Series(offset, count)
	if err != nil {
		zlog.Warningf("get kline series offset out of range: offset=%d, count=%d, length=%d, lastTs=%d", offset, count, c.kSeries.Length(), c.kSeries.LastTs())
		panic(err)
	}
	return ks
}

// func (c *IndicatorContext) State(effects ...any)
// 窗口/参数
// tradingPlan -> interval -> context -> {macd{window:0,fast:9,slow:21,single:10}: state, ema21:state} -> state[[ema]{1.1, 1.2}, [ema]{1.3, 1.4}]
func (c *IndicatorContext) State() indicator.IIndicatorState {
	inputs := collect.Mapping(c.indicator.Meta().Input, func(in types.InputArg) string {
		return c.Input().String(in.Name)
	})
	stateKey := fmt.Sprintf("%s{%s}", c.indicator.Meta().Name, strings.Join(inputs, ","))

	state, ok := c.indicatorStates[stateKey]
	if !ok {
		state = NewIndicatorState(c.kSeries.Interval)
		c.indicatorStates[stateKey] = state
		// 从头KlineSeries跑一遍, 针对ema,macd等回溯迭代指标, 将state与KlineSeries对齐
		c.backtrackIndicatorState(c.indicator)
	}
	state.SetIndicatorContext(c)
	return state
}

// 当首次初始化某个state后, 把indicator在klineSeries从头跑一遍
func (c *IndicatorContext) backtrackIndicatorState(indicator indicator.IIndicator) {
	_offset := c.offset
	defer c.SetOffset(_offset)

	candleLength := c.kSeries.Length()
	candlePeriods := indicator.CandlePeriods(c)
	for i := int(candlePeriods) - 1; i < candleLength; i++ {
		offset := int16(candleLength - i - 1)
		c.SetOffset(offset)
		// 计算指标值，指标内部会更新state
		indicator.Calculate(c)
	}
}

// 获取窗口类型指标
func (c *IndicatorContext) Indicator(name string, args ...any) (series indicator.IIndicatorSeries) {
	indicator, ok := c.indicatorsReg.Indicator(name)
	if !ok {
		panic(fmt.Errorf("indicator %s not exists", name))
	}

	// 避免循环依赖
	if slices.Contains(c.indicatorTrace, name) {
		panic(fmt.Errorf("indicator %s recursive call", name))
	}

	input := matchIndicatorArgs(args...)

	indicatorContext := NewIndicatorContext(indicator, input, c.indicatorStates, c.kSeries, c.indicatorsReg)
	indicatorContext.offset = c.offset
	indicatorContext.indicatorTrace = append(c.indicatorTrace, name)
	return NewWindowIndicatorSeries(indicator, indicatorContext)
}

func matchIndicatorArgs(args ...any) (input types.Input) {
inputLoop:
	for _, arg := range args {
		switch v := arg.(type) {
		case types.Input:
			input = v
			break inputLoop
		}
	}
windowLoop:
	for _, arg := range args {
		switch v := arg.(type) {
		case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
			window := cast.ToInt16(v)
			if input == nil {
				input = make(types.Input)
			}
			input["window"] = window
			break windowLoop
		}
	}
	return
}