// Copyright 2020 ConsenSys Software Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Code generated by consensys/gnark-crypto DO NOT EDIT

package fft

import (
	"math/big"
	"strconv"
	"testing"

	"github.com/consensys/linea-monorepo/prover/maths/common/poly"
	"github.com/consensys/linea-monorepo/prover/maths/field"
	"github.com/stretchr/testify/require"

	"github.com/leanovate/gopter"
	"github.com/leanovate/gopter/gen"
	"github.com/leanovate/gopter/prop"
)

func TestFFT(t *testing.T) {
	const maxSize = 1 << 10

	nbCosets := 3
	domainWithPrecompute := NewDomain(maxSize).WithCoset()

	parameters := gopter.DefaultTestParameters()
	parameters.MinSuccessfulTests = 5

	properties := gopter.NewProperties(parameters)

	// checks that a random evaluation of a dual function eval(gen**ithpower) is consistent with the FFT result
	testA := func(ithpower int) bool {

		pol := make([]field.Element, maxSize)
		backupPol := make([]field.Element, maxSize)

		for i := 0; i < maxSize; i++ {
			pol[i].SetRandom()
		}
		copy(backupPol, pol)

		domainWithPrecompute.FFT(pol, DIF)
		BitReverse(pol)

		sample := domainWithPrecompute.Generator
		sample.Exp(sample, big.NewInt(int64(ithpower)))

		eval := poly.EvalUnivariate(backupPol, sample)

		return eval.Equal(&pol[ithpower])
	}

	testA(53)

	properties.Property("DIF FFT on cosets should be consistent with dual basis", prop.ForAll(

		// checks that a random evaluation of a dual function eval(gen**ithpower) is consistent with the FFT result
		func(ithpower int) bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			domainWithPrecompute.FFT(pol, DIF, OnCoset())
			BitReverse(pol)

			sample := domainWithPrecompute.Generator
			sample.Exp(sample, big.NewInt(int64(ithpower))).
				Mul(&sample, &domainWithPrecompute.FrMultiplicativeGen)

			eval := poly.EvalUnivariate(backupPol, sample)

			return eval.Equal(&pol[ithpower])

		},
		gen.IntRange(0, maxSize-1),
	))

	properties.Property("DIT FFT should be consistent with dual basis", prop.ForAll(

		// checks that a random evaluation of a dual function eval(gen**ithpower) is consistent with the FFT result
		func(ithpower int) bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			BitReverse(pol)
			domainWithPrecompute.FFT(pol, DIT)

			sample := domainWithPrecompute.Generator
			sample.Exp(sample, big.NewInt(int64(ithpower)))

			eval := poly.EvalUnivariate(backupPol, sample)

			return eval.Equal(&pol[ithpower])

		},
		gen.IntRange(0, maxSize-1),
	))

	properties.Property("bitReverse(DIF FFT(DIT FFT (bitReverse))))==id", prop.ForAll(

		func() bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			BitReverse(pol)
			domainWithPrecompute.FFT(pol, DIT)
			domainWithPrecompute.FFTInverse(pol, DIF)
			BitReverse(pol)

			check := true
			for i := 0; i < len(pol); i++ {
				check = check && pol[i].Equal(&backupPol[i])
			}
			return check
		},
	))

	properties.Property("bitReverse(DIF FFT(DIT FFT (bitReverse))))==id on cosets", prop.ForAll(

		func() bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			check := true

			for i := 1; i <= nbCosets; i++ {

				BitReverse(pol)
				domainWithPrecompute.FFT(pol, DIT, OnCoset())
				domainWithPrecompute.FFTInverse(pol, DIF, OnCoset())
				BitReverse(pol)

				for i := 0; i < len(pol); i++ {
					check = check && pol[i].Equal(&backupPol[i])
				}
			}

			return check
		},
	))

	properties.Property("DIT FFT(DIF FFT)==id", prop.ForAll(

		func() bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			domainWithPrecompute.FFTInverse(pol, DIF)
			domainWithPrecompute.FFT(pol, DIT)

			check := true
			for i := 0; i < len(pol); i++ {
				check = check && (pol[i] == backupPol[i])
			}
			return check
		},
	))

	properties.Property("DIT FFT(DIF FFT)==id on cosets", prop.ForAll(

		func() bool {

			pol := make([]field.Element, maxSize)
			backupPol := make([]field.Element, maxSize)

			for i := 0; i < maxSize; i++ {
				pol[i].SetRandom()
			}
			copy(backupPol, pol)

			domainWithPrecompute.FFTInverse(pol, DIF, OnCoset())
			domainWithPrecompute.FFT(pol, DIT, OnCoset())

			check := true
			for i := 0; i < len(pol); i++ {
				check = check && (pol[i] == backupPol[i])
			}
			return check
		},
	))

	properties.TestingRun(t, gopter.ConsoleReporter(false))

}

func TestOmega(t *testing.T) {
	n := 4
	// Test that it returns indeed a root of unity
	omega := GetOmega(n)
	omega.Exp(omega, big.NewInt(int64(n)))
	require.Equal(t, "1", omega.String())
}

// --------------------------------------------------------------------
// benches
func BenchmarkFFT(b *testing.B) {

	const maxSize = 1 << 20

	pol := make([]field.Element, maxSize)
	pol[0].SetRandom()
	for i := 1; i < maxSize; i++ {
		pol[i] = pol[i-1]
	}

	for i := 8; i < 20; i++ {
		sizeDomain := 1 << i
		b.Run("fft 2**"+strconv.Itoa(i), func(b *testing.B) {
			domain := NewDomain(sizeDomain).WithCoset()
			b.ResetTimer()
			for j := 0; j < b.N; j++ {
				domain.FFT(pol[:sizeDomain], DIT)
			}
		})
		// b.Run("fft 2**"+strconv.Itoa(i)+"(coset)", func(b *testing.B) {
		// 	domain := NewDomain(uint64(sizeDomain))
		// 	b.ResetTimer()
		// 	for j := 0; j < b.N; j++ {
		// 		domain.FFT(pol[:sizeDomain], DIT, true)
		// 	}
		// })
	}

}

func BenchmarkFFTDITCosetReference(b *testing.B) {
	const maxSize = 1 << 20

	pol := make([]field.Element, maxSize)
	pol[0].SetRandom()
	for i := 1; i < maxSize; i++ {
		pol[i] = pol[i-1]
	}

	domain := NewDomain(maxSize).WithCoset()

	b.ResetTimer()
	for j := 0; j < b.N; j++ {
		domain.FFT(pol, DIT, OnCoset())
	}
}

func BenchmarkFFTDIFReference(b *testing.B) {
	const maxSize = 1 << 20

	pol := make([]field.Element, maxSize)
	pol[0].SetRandom()
	for i := 1; i < maxSize; i++ {
		pol[i] = pol[i-1]
	}

	domain := NewDomain(maxSize)

	b.ResetTimer()
	for j := 0; j < b.N; j++ {
		domain.FFT(pol, DIF)
	}
}
