package indicator

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

// BollMB 布林带中轨
type BollMB struct {
}

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

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

func (c *BollMB) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	closeSeries := ctx.Series(0, int16(window)).Close()
	vector = closeSeries.Avg()
	return
}

// BollUB 布林带上轨
type BollUB struct {
}

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

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

func (c *BollUB) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	closeSeries := ctx.Series(0, int16(window)).Close()
	mb := closeSeries.Avg()

	// 标准差σ_t = sqrt(∑(P-MB)^2 / (n-1))
	sst := float64(0)
	for _, p := range closeSeries {
		sst += math.Pow(p-mb, 2)
	}
	sigma := math.Sqrt(sst / float64(window-1))

	vector = mb + 2*sigma
	return
}

// BollLB 布林带下轨
type BollLB struct {
}

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

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

func (c *BollLB) Calculate(ctx IIndicatorContext) (vector float64) {
	window := ctx.Input().Int16("window")
	closeSeries := ctx.Series(0, int16(window)).Close()
	mb := closeSeries.Avg()

	// 标准差σ_t = sqrt(∑(P-MB)^2 / (n-1))
	sst := float64(0)
	for _, p := range closeSeries {
		sst += math.Pow(p-mb, 2)
	}
	sigma := math.Sqrt(sst / float64(window-1))

	vector = mb - 2*sigma
	return
}