package indicator

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

type KDJ struct {
}

func (c *KDJ) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "KDJ",
		Input: []types.InputArg{
			{Name: "rsvWindow", Type: types.InputTypeUInt, Desc: "周期"}, // 9
			{Name: "kWindow", Type: types.InputTypeUInt, Desc: "K平滑"},  // 3
			{Name: "dWindow", Type: types.InputTypeUInt, Desc: "D平滑"},  // 3
		},
		State: []string{"k", "d", "j"},
		Plots: []Plot{
			{State: "k", Type: PlotLine, Props: PlotProps{"color": ColorBlue}},
			{State: "d", Type: PlotLine, Props: PlotProps{"color": ColorYellow}},
			{State: "j", Type: PlotLine, Props: PlotProps{"color": ColorPurple}},
		},
	}
}

func (c *KDJ) CandlePeriods(ctx IIndicatorContext) int16 {
	return ctx.Input().Int16("rsvWindow")
}

func (c *KDJ) Calculate(ctx IIndicatorContext) (vector float64) {
	rsvWindow := ctx.Input().Int16("rsvWindow")
	kWindow := float64(ctx.Input().Int16("kWindow"))
	dWindow := float64(ctx.Input().Int16("dWindow"))

	klines := ctx.Series(0, rsvWindow)
	if len(klines) < int(rsvWindow) {
		return
	}

	closePx := klines[0].CloseF64()
	low := klines.Low().Min()
	high := klines.High().Max()

	var rsv float64
	if high == low {
		rsv = 50
	} else {
		rsv = (closePx - low) / (high - low) * 100
	}

	// K
	prevK, ok := ctx.State().Get("k", 1)
	if !ok {
		prevK = 50
	}
	k := (kWindow-1)/kWindow*prevK + 1/kWindow*rsv
	ctx.State().Set("k", k)

	// D
	prevD, ok := ctx.State().Get("d", 1)
	if !ok {
		prevD = 50
	}
	d := (dWindow-1)/dWindow*prevD + 1/dWindow*k
	ctx.State().Set("d", d)

	// J
	j := 3*k - 2*d
	ctx.State().Set("j", j)

	return j
}