sig-pub/pkg/types/interval.go

package types

import (
	"fmt"
	"sort"
	"time"
)

var LossEmoji = "🔥"
var ProfitEmoji = "💰"

type Interval string

// AddMul 对指定毫秒时间戳增加周期数
func (i Interval) AddMul(ts, mul int64) (int64, bool) {
	c, ok := SupportedIntervals[i]
	if !ok {
		return ts, false
	}
	return c(ts, mul), true
}

func (i Interval) MustAddMul(ts, mul int64) int64 {
	ts, ok := i.AddMul(ts, mul)
	if !ok {
		panic(fmt.Errorf("unsupport interval AddMul: interval=%s, ts=%d, mul=%d", i, ts, mul))
	}
	return ts
}

const (
	Interval1s  = Interval("1s")
	Interval1m  = Interval("1m")
	Interval3m  = Interval("3m")
	Interval5m  = Interval("5m")
	Interval15m = Interval("15m")
	Interval30m = Interval("30m")
	Interval1h  = Interval("1h")
	Interval2h  = Interval("2h")
	Interval4h  = Interval("4h")
	Interval6h  = Interval("6h")
	Interval12h = Interval("12h")
	Interval1d  = Interval("1d")
	Interval2d  = Interval("2d")
	Interval3d  = Interval("3d")
	Interval5d  = Interval("5d")
	Interval1w  = Interval("1w")
	Interval1mo = Interval("1mo")
	Interval3mo = Interval("3mo")
)

// IntervalWindow is used by the indicators
type IntervalWindow struct {
	// The interval of kline
	Interval Interval `json:"interval"`

	// The windows size of the indicator (for example, EWMA and SMA)
	Window int `json:"window"`

	// RightWindow is used by the pivot indicator
	RightWindow *int `json:"rightWindow"`
}

type IntervalMap map[Interval]IntervalAdder

type IntervalAdder func(ts, mul int64) (ret int64)

// todo 支持其他周期方案:
// 1.订阅1m实时k线和拉取历史k线时进行聚合
// 2.PromQL: open/close -> first_over_time/last_over_time, vol->sum_over_time
var SupportedIntervals = IntervalMap{
	// Interval1s:  func(ts, mul int64) (ret int64) { return ts + (1000 * mul) },
	Interval1m:  func(ts, mul int64) (ret int64) { return ts + (1 * 60 * 1000 * mul) },
	Interval3m:  func(ts, mul int64) (ret int64) { return ts + (3 * 60 * 1000 * mul) },
	Interval5m:  func(ts, mul int64) (ret int64) { return ts + (5 * 60 * 1000 * mul) },
	Interval15m: func(ts, mul int64) (ret int64) { return ts + (15 * 60 * 1000 * mul) },
	Interval30m: func(ts, mul int64) (ret int64) { return ts + (30 * 60 * 1000 * mul) },
	Interval1h:  func(ts, mul int64) (ret int64) { return ts + (60 * 60 * 1000 * mul) },
	Interval2h:  func(ts, mul int64) (ret int64) { return ts + (2 * 60 * 60 * 1000 * mul) },
	Interval4h:  func(ts, mul int64) (ret int64) { return ts + (4 * 60 * 60 * 1000 * mul) },
	Interval6h:  func(ts, mul int64) (ret int64) { return ts + (6 * 60 * 60 * 1000 * mul) },
	Interval12h: func(ts, mul int64) (ret int64) { return ts + (12 * 60 * 60 * 1000 * mul) },
	Interval1d:  func(ts, mul int64) (ret int64) { return ts + (24 * 60 * 60 * 1000 * mul) },
	Interval2d:  func(ts, mul int64) (ret int64) { return ts + (2 * 24 * 60 * 60 * 1000 * mul) },
	Interval3d:  func(ts, mul int64) (ret int64) { return ts + (3 * 24 * 60 * 60 * 1000 * mul) },
	Interval5d:  func(ts, mul int64) (ret int64) { return ts + (5 * 24 * 60 * 60 * 1000 * mul) },
	Interval1w:  func(ts, mul int64) (ret int64) { return ts + (7 * 24 * 60 * 60 * 1000 * mul) },
	Interval1mo: func(ts, mul int64) (ret int64) { return time.UnixMilli(ts).AddDate(0, int(mul), 0).UnixMilli() },
	Interval3mo: func(ts, mul int64) (ret int64) { return time.UnixMilli(ts).AddDate(0, int(3*mul), 0).UnixMilli() },
}

var (
	iotasIntervals  []Interval
	intervalIotaMax int
	intervalIotas   = map[Interval]int{}
)

func init() {
	iotasIntervals = make([]Interval, len(SupportedIntervals))
	var tss = make([]int64, len(SupportedIntervals))
	var i = 0
	for interval, adder := range SupportedIntervals {
		iotasIntervals[i] = interval
		tss[i] = adder(0, 1)
		i++
	}
	sort.Slice(iotasIntervals, func(i, j int) bool {
		return tss[i] < tss[j]
	})
	for i, interval := range iotasIntervals {
		index := i + 1 // 0保留
		intervalIotas[interval] = index
		intervalIotaMax = max(intervalIotaMax, index)
	}
	// zlog.Debugf("init intervals: %#v", iotasIntervals)
	// zlog.Debugf("init intervalsIotas: %#v", intervalIotas)
}

// 对周期进行排序
func IntervalsSort(intervals []Interval) {
	if len(intervals) < 2 {
		return
	}
	sort.SliceStable(intervals, func(i, j int) bool {
		return intervalIotas[intervals[i]] < intervalIotas[intervals[j]]
	})
}

type IntervalState[T any] struct {
	state []T
}

func NewIntervalState[T any]() *IntervalState[T] {
	return &IntervalState[T]{
		state: make([]T, intervalIotaMax+1),
	}
}

func (s *IntervalState[T]) Get(interval Interval) T {
	i := intervalIotas[interval]
	return s.state[i]
}

func (s *IntervalState[T]) Set(interval Interval, v T) {
	i := intervalIotas[interval]
	s.state[i] = v
}

func (s *IntervalState[T]) Range(f func(interval Interval, v T)) {
	for i, interval := range iotasIntervals {
		index := i + 1 // 0保留
		v := s.state[index]
		f(interval, v)
	}
}

func (s *IntervalState[T]) SetIf(interval Interval, v T, cond func(old T) bool) {
	i := intervalIotas[interval]
	old := s.state[i]
	if cond(old) {
		s.state[i] = v
	}
}