package ringsis_{{.ModulusDegree}}_{{.LogTwoBound}}

import (
	"testing"

	"github.com/consensys/gnark-crypto/ecc/bls12-377/fr/fft"
	"github.com/consensys/linea-monorepo/prover/maths/field"
	"github.com/stretchr/testify/assert"
)

{{- $bitPerField := 256}}
{{- $limbPerField := div $bitPerField .LogTwoBound}}
{{- $fieldPerPoly := div .ModulusDegree $limbPerField}}
{{- $numMask := pow 2 $fieldPerPoly}}

func TestPartialFFT(t *testing.T) {

	var (
		domain   = fft.NewDomain({{.ModulusDegree}})
		twiddles = PrecomputeTwiddlesCoset(domain.Generator, domain.FrMultiplicativeGen)
	)

	for mask := 0; mask < {{$numMask}}; mask++ {

		var (
			a = vec123456()
			b = vec123456()
		)

		zeroizeWithMask(a, mask)
		zeroizeWithMask(b, mask)

		domain.FFT(a, fft.DIF, fft.OnCoset())
		partialFFT[mask](b, twiddles)
		assert.Equal(t, a, b)
	}

}

func vec123456() []field.Element {
	vec := make([]field.Element, {{.ModulusDegree}})
	for i := range vec {
		vec[i].SetInt64(int64(i))
	}
	return vec
}

func zeroizeWithMask(v []field.Element, mask int) {
	for i := 0; i < {{$fieldPerPoly}}; i++ {
		if (mask>>i)&1 == 1 {
			continue
		}

		for j := 0; j < {{$limbPerField}}; j++ {
			v[{{$limbPerField}}*i+j].SetZero()
		}
	}
}
