package indicator

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

// ADX 趋势强度指标, 用于衡量价格趋势的强弱, 而不考虑趋势的方向
// 它通常与 +DI（正向方向性指标）和 -DI（负向方向性指标）一起使用，形成 DMI（Directional Movement Index，方向性运动指数）系统
type ADX struct {
}

// indicator interface
func (c *ADX) Meta() IndicatorMeta {
	return IndicatorMeta{
		Name: "ADX",
		Input: []types.InputArg{
			{Name: "window", Type: types.InputTypeUInt, Desc: "窗口大小"}, // 14
		},
	}
}

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

// Calculate 计算单根k线ADX指标
func (c *ADX) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	if window <= 0 {
		return
	}

	kls := ctx.Series(0, window+1)
	highs := kls.High()
	lows := kls.Low()
	closes := kls.Close()

	trs := make([]float64, 0, window)
	plusDMs := make([]float64, 0, window)
	minusDMs := make([]float64, 0, window)

	for i := range int(window) {
		high := highs[i]
		low := lows[i]
		prevHigh := highs[i+1]
		prevLow := lows[i+1]
		prevClose := closes[i+1]

		// 计算真实范围(True Range, TR)
		tr := math.Max(high-low, math.Max(math.Abs(high-prevClose), math.Abs(low-prevClose)))
		trs = append(trs, tr)

		// 计算方向运动 (Directional Movement, DM)
		upMove := high - prevHigh
		downMove := prevLow - low
		plusDM := 0.0
		minusDM := 0.0
		if upMove > downMove && upMove > 0 {
			plusDM = upMove
		}
		if downMove > upMove && downMove > 0 {
			minusDM = downMove
		}
		plusDMs = append(plusDMs, plusDM)
		minusDMs = append(minusDMs, minusDM)
	}

	sumTR := ta.Sum(trs)
	sumPlus := ta.Sum(plusDMs)
	sumMinus := ta.Sum(minusDMs)

	// Wilder smoothing
	prevSmTR, ok := ctx.State().Get("_smTR", 1)
	if !ok {
		prevSmTR = sumTR
	}
	smTR := prevSmTR - prevSmTR/float64(window) + trs[0]
	ctx.State().Set("_smTR", smTR)

	prevSmPlus, ok := ctx.State().Get("_smPlus", 1)
	if !ok {
		prevSmPlus = sumPlus
	}
	smPlus := prevSmPlus - prevSmPlus/float64(window) + plusDMs[0]
	ctx.State().Set("_smPlus", smPlus)

	prevSmMinus, ok := ctx.State().Get("_smMinus", 1)
	if !ok {
		prevSmMinus = sumMinus
	}
	smMinus := prevSmMinus - prevSmMinus/float64(window) + minusDMs[0]
	ctx.State().Set("_smMinus", smMinus)

	pdi := 0.0
	mdi := 0.0
	if smTR > 0 {
		pdi = 100 * smPlus / smTR
		mdi = 100 * smMinus / smTR
	}
	// ctx.State().Set("pdi", pdi)
	// ctx.State().Set("mdi", mdi)

	dx := 0.0
	if (pdi + mdi) > 0 {
		dx = 100 * math.Abs(pdi-mdi) / (pdi + mdi)
	}

	prevAdx, ok := ctx.State().Get("_adx", 1)
	if !ok {
		// first time: fall back to current DX as initial ADX
		prevAdx = dx
	}
	adx := ((prevAdx*float64(window-1) + dx) / float64(window))
	ctx.State().Set("_adx", adx)

	vector = adx
	return
}