sig-pub/pkg/indicator/vrvp.go

package indicator

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

const (
	MaxRawBuckets = 2000 // 最大原始分桶数量，超过此数量将进行合并
)

// VRVP Volume Profile (Visible Range Volume Profile)
// 成交量分布图: 显示特定时间段内各价格水平的成交量分布
type VRVP struct {
	buckets  int     // 价格行数
	minPrice float64 // 最低价
	maxPrice float64 // 最高价

	// 流式计算状态
	baseStep   float64                     // 当前的基础步长
	rawBuckets map[int64]*types.VRVPBucket // 原始分桶数据, key = int64(price / baseStep)
	klines     int64                       // 累计K线数量
	skts, ekts int64                       // 开始时间戳, 结束时间戳
}

func (c *VRVP) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "VRVP",
		Desc: "成交量分布图",
		Input: []types.InputArg{
			{Name: "buckets", Type: types.InputTypeUInt, Desc: "价格分桶数量", Default: 24},
		},
	}
}

// Init 初始化参数
func (v *VRVP) Init(input types.Input) (err error) {
	v.buckets = input.Int("buckets")
	if v.buckets <= 0 {
		v.buckets = 24
	}
	v.minPrice = math.MaxFloat64
	v.maxPrice = -math.MaxFloat64
	v.rawBuckets = make(map[int64]*types.VRVPBucket)
	v.baseStep = 0
	return
}

func (v *VRVP) CandlePeriods(ctx IIndicatorContext) int16 {
	return 1
}

// Eliminate 清除这根K线的数据
func (v *VRVP) Eliminate(ctx IIndicatorContext) {
	k := ctx.Get(0)
	high := k.HighF64()
	low := k.LowF64()
	vol := k.VolF64()

	if vol <= 0 {
		return
	}

	v.klines--

	// 限制说明 ：在流式计算模式下，无法精确回滚 minPrice 、 maxPrice 以及 baseStep 的历史变化。
	// 这意味着在长期运行后，分布图的统计范围可能会比实际存在的 K 线范围略大，但这不影响分布形状的准确性。

	if v.baseStep <= 0 {
		return
	}

	// 将成交量从原始桶中移除
	startIdx := int64(math.Floor(low / v.baseStep))
	endIdx := int64(math.Floor(high / v.baseStep))

	coveredBuckets := float64(endIdx - startIdx + 1)
	volPerBucket := vol / coveredBuckets

	isUp := k.CloseF64() >= k.OpenF64()

	for i := startIdx; i <= endIdx; i++ {
		bucket, exists := v.rawBuckets[i]
		if !exists {
			continue
		}

		bucket.Volume -= volPerBucket
		if isUp {
			bucket.BuyVol -= volPerBucket
		} else {
			bucket.SellVol -= volPerBucket
		}

		// 如果桶的成交量归零（考虑浮点误差），则移除该桶以节省内存
		if bucket.Volume < 1e-8 {
			delete(v.rawBuckets, i)
		}
	}
}

// Accumulate 累积每根K线的数据
func (v *VRVP) Accumulate(ctx IIndicatorContext) {
	k := ctx.Get(0)
	high := k.HighF64()
	low := k.LowF64()
	vol := k.VolF64()
	v.klines++
	if v.skts == 0 {
		v.skts = k.Ts
	}
	v.ekts = k.Ts

	if vol <= 0 {
		return
	}

	// 更新全局极值
	if high > v.maxPrice {
		v.maxPrice = high
	}
	if low < v.minPrice {
		v.minPrice = low
	}

	// 初始化 baseStep
	if v.baseStep == 0 {
		// 如果 K 线有波动，使用波动的一小部分作为初始精度
		// 如果无波动(High==Low)，使用价格的万分之一
		if high > low {
			v.baseStep = (high - low) / 100.0
		} else {
			if low > 0 {
				v.baseStep = low * 0.0001
			} else {
				v.baseStep = 0.01 // 默认值
			}
		}
	}

	// 将成交量分配到原始桶中
	startIdx := int64(math.Floor(low / v.baseStep))
	endIdx := int64(math.Floor(high / v.baseStep))

	coveredBuckets := float64(endIdx - startIdx + 1)
	volPerBucket := vol / coveredBuckets

	isUp := k.CloseF64() >= k.OpenF64()

	for i := startIdx; i <= endIdx; i++ {
		bucket, exists := v.rawBuckets[i]
		if !exists {
			bucket = &types.VRVPBucket{
				Price: float64(i)*v.baseStep + v.baseStep/2, // 暂存中心价
			}
			v.rawBuckets[i] = bucket
		}

		bucket.Volume += volPerBucket
		if isUp {
			bucket.BuyVol += volPerBucket
		} else {
			bucket.SellVol += volPerBucket
		}
	}

	// 检查是否需要合并桶
	if len(v.rawBuckets) > MaxRawBuckets {
		v.halveResolution()
	}
}

// halveResolution 将分辨率减半（步长翻倍）
func (v *VRVP) halveResolution() {
	newBaseStep := v.baseStep * 2
	newBuckets := make(map[int64]*types.VRVPBucket, len(v.rawBuckets)/2+1)

	for key, bucket := range v.rawBuckets {
		// 计算新的 key
		// oldPrice ~= key * oldStep
		// newKey = floor(oldPrice / newStep) = floor(key * oldStep / (2 * oldStep)) = floor(key / 2)
		newKey := key >> 1 // key / 2

		newBucket, exists := newBuckets[newKey]
		if !exists {
			newBucket = &types.VRVPBucket{
				Price: float64(newKey)*newBaseStep + newBaseStep/2,
			}
			newBuckets[newKey] = newBucket
		}

		newBucket.Volume += bucket.Volume
		newBucket.BuyVol += bucket.BuyVol
		newBucket.SellVol += bucket.SellVol
	}

	v.baseStep = newBaseStep
	v.rawBuckets = newBuckets
}

// Summary 计算最终的成交量分布结果
// 在所有K线Accumulate完成后调用
func (v *VRVP) Summary(ctx IIndicatorContext) (summary any, ok bool) {
	// zlog.Debugf("VRVP Summary: klines=%d, skts=%d, ekts=%d, minPrice=%.2f, maxPrice=%.2f, baseStep=%.2f",
	// 	v.klines, v.skts, v.ekts, v.minPrice, v.maxPrice, v.baseStep)
	if v.minPrice >= v.maxPrice {
		return nil, false
	}

	// 计算最终的目标步长
	rangeHeight := v.maxPrice - v.minPrice
	finalStep := rangeHeight / float64(v.buckets)

	// 如果 finalStep 小于当前的 baseStep，说明数据太稀疏，无法满足 rowCount 的精度要求
	// 但通常情况下，由于我们只在桶过多时才合并，baseStep 应该相对较小

	// 初始化最终分桶
	buckets := make([]types.VRVPBucket, v.buckets)
	for i := 0; i < v.buckets; i++ {
		buckets[i].Price = v.minPrice + float64(i)*finalStep + finalStep/2
	}

	// 将 rawBuckets 聚合到最终 buckets
	for _, rawB := range v.rawBuckets {
		// 计算原始桶对应的价格范围中心
		// 注意：rawB.Price 在合并过程中可能不再准确，重新计算更稳妥，或者在合并时更新 Price
		// 这里我们使用 key 重新计算，更准确
		// 但由于 map key 不在 value 中，我们无法直接获取 key
		// 所以我们需要遍历 map 的 key
		// 修正：在上面的循环中我们无法获得 key，所以需要修改遍历方式
		// 或者我们在 PriceBucket 中存储准确的 Price
		// 在 halveResolution 中，我们更新了 Price，所以 rawB.Price 是当前 baseStep 下的中心价

		price := rawB.Price

		// 找到对应的最终桶索引
		idx := int((price - v.minPrice) / finalStep)

		if idx < 0 {
			idx = 0
		} else if idx >= v.buckets {
			idx = v.buckets - 1
		}

		buckets[idx].Volume += rawB.Volume
		buckets[idx].BuyVol += rawB.BuyVol
		buckets[idx].SellVol += rawB.SellVol
	}

	ok, summary = true, &types.VRVPSummary{
		Klines:   v.klines,
		Step:     finalStep,
		MinPrice: v.minPrice,
		MaxPrice: v.maxPrice,
		Buckets:  buckets,
	}
	return
}