package indicator

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

// RVI Relative Vigor Index
// RVI = SMA(Num, N) / SMA(Denom, N)
// Num = (Close-Open) + 2*(Close_1-Open_1) + 2*(Close_2-Open_2) + (Close_3-Open_3)
// Denom = (High-Low) + 2*(High_1-Low_1) + 2*(High_2-Low_2) + (High_3-Low_3)
// Signal = (RVI + 2*RVI_1 + 2*RVI_2 + RVI_3) / 6
type RVI struct {
}

func (c *RVI) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "RVI",
		Input: []types.InputArg{
			{Name: "window", Type: types.InputTypeUInt, Desc: "窗口大小", Default: 10},
		},
		State: []string{"sig"},
		Plots: []Plot{
			{Name: "RVI", State: "vector", Type: PlotLine, Props: PlotProps{"color": ColorGreen}},
			{Name: "Signal", State: "sig", Type: PlotLine, Props: PlotProps{"color": ColorRed}},
		},
	}
}

func (c *RVI) CandlePeriods(ctx IIndicatorContext) int16 {
	return ctx.Input().Int16("window") + 3
}

func (c *RVI) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")

	count := window + 3
	klines := ctx.Series(0, count)

	if len(klines) < int(count) {
		return
	}

	var numSum, denomSum float64

	for i := int16(0); i < window; i++ {
		k0 := klines[i]
		k1 := klines[i+1]
		k2 := klines[i+2]
		k3 := klines[i+3]

		// CO = Close - Open
		co0 := k0.CloseF64() - k0.OpenF64()
		co1 := k1.CloseF64() - k1.OpenF64()
		co2 := k2.CloseF64() - k2.OpenF64()
		co3 := k3.CloseF64() - k3.OpenF64()

		// HL = High - Low
		hl0 := k0.HighF64() - k0.LowF64()
		hl1 := k1.HighF64() - k1.LowF64()
		hl2 := k2.HighF64() - k2.LowF64()
		hl3 := k3.HighF64() - k3.LowF64()

		val1 := (co0 + 2*co1 + 2*co2 + co3) / 6.0
		val2 := (hl0 + 2*hl1 + 2*hl2 + hl3) / 6.0

		numSum += val1
		denomSum += val2
	}

	var rvi float64
	if denomSum != 0 {
		rvi = numSum / denomSum
	}

	ctx.State().Set("rvi", rvi)

	// Calculate Signal
	// Signal = (RVI + 2*RVI_1 + 2*RVI_2 + RVI_3) / 6
	rvi1, ok1 := ctx.State().Get("rvi", 1)
	rvi2, ok2 := ctx.State().Get("rvi", 2)
	rvi3, ok3 := ctx.State().Get("rvi", 3)

	var sig float64
	if ok1 && ok2 && ok3 {
		sig = (rvi + 2*rvi1 + 2*rvi2 + rvi3) / 6.0
	} else {
		sig = rvi
	}
	ctx.State().Set("sig", sig)

	return rvi
}