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

package sig

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

// IndicatorState
// ema, obv 指标递归计算时的状态存储
type IndicatorState struct {
	indicator.IIndicatorState
	indicatorContext indicator.IIndicatorContext
	interval         types.Interval
	intervalAdder    types.IntervalAdder
	state            map[string]*types.RingSeries[float64]
	lastTs           map[string]int64
}

func NewIndicatorState(interval types.Interval) *IndicatorState {
	intervalAdder, ok := types.SupportedIntervals[interval]
	if !ok {
		panic(fmt.Errorf("unsupport interval: %s", interval))
	}
	return &IndicatorState{
		interval:      interval,
		intervalAdder: intervalAdder,
		state:         make(map[string]*types.RingSeries[float64]),
		lastTs:        make(map[string]int64),
	}
}

func (s *IndicatorState) SetIndicatorContext(indicatorContext indicator.IIndicatorContext) {
	s.indicatorContext = indicatorContext
}

func (s *IndicatorState) ring(k string) *types.RingSeries[float64] {
	ring, ok := s.state[k]
	if !ok {
		ring = types.NewRingSeries[float64](indicator.MaxWindow, 8)
		s.state[k] = ring
	}
	return ring
}

func (s *IndicatorState) Set(k string, v float64) {
	ts := s.indicatorContext.Get(0).Ts
	if lastTs := s.lastTs[k]; lastTs < ts {
		// panic可替换为丢失指标用前一个值填充类似vmtsdb
		if expectTs := s.intervalAdder(lastTs, 1); expectTs != ts && lastTs != 0 {
			panic(fmt.Errorf("state 不连续: lastTs=%d, got=%d, expected=%d", lastTs, ts, expectTs))
		}
		s.ring(k).Push(v)
		s.lastTs[k] = ts
	}
}

func (s *IndicatorState) Get(k string, offset int16) (v float64, ok bool) {
	ts := s.indicatorContext.Get(0).Ts
	target := s.intervalAdder(ts, -int64(offset))
	ring := s.ring(k)
	for i := 0; i < ring.Length(); i++ {
		if s.intervalAdder(s.lastTs[k], -int64(i)) == target {
			return ring.Get(i)
		}
	}
	return
}

func (s *IndicatorState) Series(k string, offset, count int16) (v series.Floats, ok bool) {
	ts := s.indicatorContext.Get(0).Ts
	target := s.intervalAdder(ts, -int64(offset))
	ring := s.ring(k)
	for i := 0; i < ring.Length(); i++ {
		if s.intervalAdder(s.lastTs[k], -int64(i)) == target {
			return ring.Series(i, int(count))
		}
	}
	return
}