1package gocircomprover
  2
  3import (
  4	"fmt"
  5	"math"
  6	"math/big"
  7)
  8
  9type rootsT struct {
 10	roots [][]*big.Int
 11	w     []*big.Int
 12}
 13
 14func newRootsT() rootsT {
 15	var roots rootsT
 16
 17	rem := new(big.Int).Sub(R, big.NewInt(1))
 18	s := 0
 19	for rem.Bit(0) == 0 { // rem.Bit==0 when even
 20		s++
 21		rem = new(big.Int).Rsh(rem, 1)
 22	}
 23	roots.w = make([]*big.Int, s+1)
 24	roots.w[s] = FExp(big.NewInt(5), rem)
 25
 26	n := s - 1
 27	for n >= 0 {
 28		roots.w[n] = FMul(roots.w[n+1], roots.w[n+1])
 29		n--
 30	}
 31	roots.roots = make([][]*big.Int, 50)
 32
 33	roots.setRoots(15)
 34	return roots
 35}
 36
 37func (roots rootsT) setRoots(n int) {
 38	// var roots []bool
 39	for i := n; i >= 0 && nil == roots.roots[i]; i-- { // TODO tmp i<=len(r)
 40		r := big.NewInt(1)
 41		nroots := 1 << i
 42		var rootsi []*big.Int
 43		for j := 0; j < nroots; j++ {
 44			rootsi = append(rootsi, r)
 45			r = FMul(r, roots.w[i])
 46		}
 47		// fmt.Println("rootsi", rootsi)
 48		roots.roots[i] = rootsi
 49	}
 50}
 51
 52func fft(roots rootsT, pall []*big.Int, bits, offset, step int) []*big.Int {
 53	n := 1 << bits
 54	if n == 1 {
 55		return []*big.Int{pall[offset]}
 56	} else if n == 2 {
 57		return []*big.Int{
 58			FAdd(pall[offset], pall[offset+step]), // TODO tmp
 59			FSub(pall[offset], pall[offset+step]),
 60		}
 61	}
 62
 63	ndiv2 := n >> 1
 64	p1 := fft(roots, pall, bits-1, offset, step*2)
 65	p2 := fft(roots, pall, bits-1, offset+step, step*2)
 66
 67	// var out []*big.Int
 68	out := make([]*big.Int, n)
 69	for i := 0; i < ndiv2; i++ {
 70		fmt.Println(i, len(roots.roots))
 71		out[i] = FAdd(p1[i], FMul(roots.roots[bits][i], p2[i]))
 72		out[i+ndiv2] = FSub(p1[i], FMul(roots.roots[bits][i], p2[i]))
 73	}
 74	return out
 75}
 76
 77func ifft(p []*big.Int) []*big.Int {
 78	if len(p) <= 1 {
 79		return p
 80	}
 81	bits := math.Log2(float64(len(p)-1)) + 1
 82	roots := newRootsT()
 83	roots.setRoots(int(bits))
 84	m := 1 << int(bits)
 85	ep := extend(p, m)
 86	res := fft(roots, ep, int(bits), 0, 1)
 87
 88	twoinvm := FInv(FMul(big.NewInt(1), big.NewInt(int64(m))))
 89
 90	var resn []*big.Int
 91	for i := 0; i < m; i++ {
 92		resn = append(resn, FMul(res[(m-i)%m], twoinvm))
 93	}
 94
 95	return resn
 96}
 97
 98func extend(p []*big.Int, e int) []*big.Int {
 99	if e == len(p) {
100		return p
101	}
102	z := arrayOfZeroes(e - len(p))
103	return append(p, z...)
104}