package strategy

import (
	"github.com/shopspring/decimal"
	"math"
)

/************** O(1) 滾動標準差：Welford + ring buffer **************/

// 滾動統計（固定窗寬）：維持 mean 與 M2（平方離差和）
// 支援「移除最舊樣本」+「新增最新樣本」的 O(1) 更新
type rollingWelford struct {
	n     int               // 窗寬
	buf   []decimal.Decimal // 環狀緩衝（固定長度 n）
	head  int               // 下一個覆寫位置（最舊元素位置）
	size  int               // 當前已填入數量（<= n）
	mean  decimal.Decimal   // 目前視窗內的平均值
	m2    decimal.Decimal   // 目前視窗內的平方離差和（∑(x-mean)^2）
	ready bool              // 是否已填滿 n
}

func newRollingWelford(n int) *rollingWelford {
	if n <= 0 {
		panic("rollingWelford window must be > 0")
	}
	return &rollingWelford{
		n:   n,
		buf: make([]decimal.Decimal, n),
	}
}

// add 一個樣本（標準 Welford 加法）
func (rw *rollingWelford) add(x decimal.Decimal) {
	// n_old = size, n_new = size+1
	nOld := decimal.NewFromInt(int64(rw.size))
	nNew := nOld.Add(decimal.NewFromInt(1))

	delta := x.Sub(rw.mean)
	meanNew := rw.mean.Add(delta.Div(nNew))
	delta2 := x.Sub(meanNew)
	rw.m2 = rw.m2.Add(delta.Mul(delta2))
	rw.mean = meanNew

	if rw.size < rw.n {
		rw.size++
		if rw.size == rw.n {
			rw.ready = true
		}
	}
}

// remove 一個樣本（Welford 反向移除公式）
func (rw *rollingWelford) remove(x decimal.Decimal) {
	if rw.size <= 0 {
		return
	}
	// n_old = size, n_new = size-1
	nOld := decimal.NewFromInt(int64(rw.size))
	nNew := nOld.Sub(decimal.NewFromInt(1))
	if nNew.LessThanOrEqual(decimal.Zero) {
		// 視窗將清空
		rw.size = 0
		rw.mean = decimal.Zero
		rw.m2 = decimal.Zero
		rw.ready = false
		return
	}
	delta := x.Sub(rw.mean)                 // (x - mean_old)
	meanNew := rw.mean.Sub(delta.Div(nNew)) // mean_new = mean_old - delta / n_new
	// M2_new = M2_old - delta*(x - mean_new)
	rw.m2 = rw.m2.Sub(delta.Mul(x.Sub(meanNew)))
	rw.mean = meanNew
	rw.size--
	rw.ready = rw.size == rw.n // 正常來說 size==n-1，移除後會在 add 後再回到 n
}

// push：O(1) 更新（移除最舊 + 新增最新）
func (rw *rollingWelford) push(x decimal.Decimal) {
	if rw.size == rw.n {
		old := rw.buf[rw.head]
		rw.remove(old)
	}
	rw.buf[rw.head] = x
	rw.head = (rw.head + 1) % rw.n
	rw.add(x)
}

// mid() = 平均；std() = sqrt(variance)
// 這裡採「母體方差」分母 = n（與很多圖表預設一致；如需樣本用 n-1）
func (rw *rollingWelford) mid() decimal.Decimal { return rw.mean }

func (rw *rollingWelford) std() decimal.Decimal {
	if !rw.ready || rw.size == 0 {
		return decimal.Zero
	}
	nDec := decimal.NewFromInt(int64(rw.n))
	variance := rw.m2.Div(nDec)
	f, _ := variance.Float64()
	return decimal.NewFromFloat(math.Sqrt(f))
}

/************** Boll（Welford 版本） **************/

type BollW struct {
	rw *rollingWelford
}

type BollWOut struct {
	Mid, Upper, Lower, Std decimal.Decimal
	Ready                  bool
}

func NewBollW(n int) *BollW {
	return &BollW{rw: newRollingWelford(n)}
}

// Push 輸入收盤價 close 與 k 倍標準差（典型 2.0）
// 回傳：上下軌、均線、標準差、是否就緒（前 n 根為 false）
func (b *BollW) Push(close, k decimal.Decimal) BollWOut {
	b.rw.push(close)
	if !b.rw.ready {
		return BollWOut{Ready: false}
	}
	mid := b.rw.mid()
	std := b.rw.std()

	upper := mid.Add(k.Mul(std))
	lower := mid.Sub(k.Mul(std))
	return BollWOut{
		Mid:   mid,
		Upper: upper,
		Lower: lower,
		Std:   std,
		Ready: true,
	}
}