package bn256

import (
	"crypto/rand"
	"fmt"
	"io"
	"math/big"
	"testing"
	"time"

	"xdx.jelly/xgcl/gmath"
)

var twoTo256 = new(big.Int).Lsh(big.NewInt(1), 256)

func gfPFromBigInt(A *big.Int) *gfP {
	a := &gfP{}
	buf := make([]byte, 32)
	A.FillBytes(buf)
	a.Unmarshal(buf)
	montEncode(a, a)
	return a
}

func randomInt(r io.Reader, max *big.Int) *big.Int {
	a, _ := rand.Int(r, max)
	return a
}

func (e *gfP) random(r io.Reader) *gfP {
	e.Set(gfPFromBigInt(randomInt(r, P)))
	return e
}

func (e *gfP2) random(r io.Reader) *gfP2 {
	e.x.random(r)
	e.y.random(r)
	return e
}

func (f *gfP6) random(r io.Reader) *gfP6 {
	f.x.random(r)
	f.z.random(r)
	f.y.random(r)
	return f
}

func (f *gfP12) random(r io.Reader) *gfP12 {
	f.x.random(r)
	f.y.random(r)
	return f
}

func testOp(t *testing.T, arg1, arg2 *big.Int, op1 func(*big.Int, *big.Int, *big.Int), op2 func(*gfP, *gfP, *gfP), equal func(*gfP, *gfP) bool) {
	var a, b *gfP

	a = gfPFromBigInt(arg1)
	if arg2 != nil {
		b = gfPFromBigInt(arg2)
	}

	C := new(big.Int)
	op1(C, arg1, arg2)
	c1 := gfPFromBigInt(C)

	c2 := &gfP{}
	op2(c2, a, b)
	if !equal(c1, c2) {
		fmt.Println("failed:")
		fmt.Println("\ta =", a)
		fmt.Println("\tb =", b)
		fmt.Println("\tc1 =", c1)
		fmt.Println("\tc2 =", c2)
		t.Fatal()
	}
}

func TestGFpSqrt(t *testing.T) {
	for i := 0; i < 100000; i++ {
		a := randomInt(rand.Reader, twoTo256)
		a.Mul(a, a).Mod(a, N)
		testOp(t,
			big.NewInt(1),
			nil,
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.ModSqrt(a, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				c.Sqrt(a)
			},
			func(a *gfP, b *gfP) bool {
				c := &gfP{}
				gfpAdd(c, a, b)
				return *a == *b || *c == gfP{}
			},
		)
	}
}

func TestGFpNeg(t *testing.T) {
	a := &gfP{0xe56f9b27e351457c, 0x21f2934b1a7aeedb, 0xd603ab4ff58ec745, 0xb640000002a3a6f1}
	gfpNeg(a, a)
	if *a != (gfP{1, 0, 0, 0}) {
		t.Fatal()
	}

	for i := 0; i < 1000000; i++ {
		a := randomInt(rand.Reader, P)
		a.Mul(a, a).Mod(a, N)
		testOp(t,
			randomInt(rand.Reader, P),
			nil,
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Neg(a).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpNeg(c, a)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)
	}

}

func TestGFpAdd(t *testing.T) {
	for i := 0; i < 1000000; i++ {
		testOp(t,
			randomInt(rand.Reader, twoTo256),
			randomInt(rand.Reader, twoTo256),
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Add(a, b).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpAdd(c, a, b)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)
	}
}

func TestGFpSub(t *testing.T) {
	for i := 0; i < 1000000; i++ {
		a := randomInt(rand.Reader, twoTo256)
		b := randomInt(rand.Reader, twoTo256)
		testOp(t,
			a,
			b.Sub(a, big.NewInt(1)),
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Sub(a, b).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpSub(c, a, b)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)
	}
}

func TestGFpHalf(t *testing.T) {
	for i := 0; i < 100000; i++ {
		testOp(t,
			randomInt(rand.Reader, P),
			nil,
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Set(a)
				if c.Bit(0) == 1 {
					c.Add(a, P)
				}
				c.Rsh(c, 1)
			},
			func(c *gfP, a *gfP, b *gfP) {
				c.half(a)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)
	}

}

// FIXME: amd64平台不支持 a and b > p
func TestGFpMul(t *testing.T) {
	for i := 0; i < 100000; i++ {
		testOp(t,
			randomInt(rand.Reader, P),
			randomInt(rand.Reader, P),
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Mul(a, b).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpMul(c, a, b)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)

		testOp(t,
			randomInt(rand.Reader, twoTo256),
			randomInt(rand.Reader, twoTo256),
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Mul(a, b).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpMul(c, a, b)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)

		testOp(t,
			new(big.Int).Sub(twoTo256, big.NewInt(int64(i+1))),
			new(big.Int).Sub(twoTo256, big.NewInt(int64(i+2))),
			func(c *big.Int, a *big.Int, b *big.Int) {
				c.Mul(a, b).Mod(c, P)
			},
			func(c *gfP, a *gfP, b *gfP) {
				gfpMul(c, a, b)
			},
			func(a *gfP, b *gfP) bool {
				return *a == *b
			},
		)
	}
}

func TestSpeedGFp(t *testing.T) {
	a := newGFp(123)
	b := newGFp(123)
	c := newGFp(1)

	buf := make([]byte, 32)

	a.Marshal(buf)
	A := new(big.Int).SetBytes(buf)
	b.Marshal(buf)
	B := new(big.Int).SetBytes(buf)
	c.Marshal(buf)
	C := new(big.Int).SetBytes(buf)

	// fmt.Println(a, b, c)
	begin := time.Now()
	total := 10000000
	for i := 0; i < total; i++ {
		gfpMul(c, a, b)
	}
	elaspe := time.Since(begin)
	fmt.Println("time: ", elaspe.Milliseconds(), "ms")
	fmt.Println(int(float64(total) / float64(elaspe.Milliseconds()) * 1000))

	begin = time.Now()
	for i := 0; i < total; i++ {
		C.Mul(A, B)
	}
	elaspe = time.Since(begin)
	fmt.Println("time: ", elaspe.Milliseconds(), "ms")
	fmt.Println(int(float64(total) / float64(elaspe.Milliseconds()) * 1000))

}

func (e *gfP2) Exp(a *gfP2, power *big.Int) *gfP2 {
	sum := (&gfP2{}).SetOne()
	t := &gfP2{}

	for i := power.BitLen() - 1; i >= 0; i-- {
		t.Square(sum)
		if power.Bit(i) != 0 {
			sum.Mul(t, a)
		} else {
			sum.Set(t)
		}
	}

	e.Set(sum)
	return e
}

// Exp set c = a^power and return c
func (f *gfP6) Exp(a *gfP6, power *big.Int) *gfP6 {
	sum := (&gfP6{}).SetOne()
	t := &gfP6{}

	for i := power.BitLen() - 1; i >= 0; i-- {
		t.Square(sum)
		if power.Bit(i) != 0 {
			sum.Mul(t, a)
		} else {
			sum.Set(t)
		}
	}

	f.Set(sum)
	return f
}
func TestGFp(t *testing.T) {
	one := newGFp(1)
	//one := &gfP{1}
	//two := &gfP{2}
	gfpMul(one, r2, one)
	gfpMul(one, &gfP{1}, one)
	montDecode(one, one)
	if *one != (gfP{1}) {
		t.Fail()
		fmt.Println(one)
		return
	}

	// a := (&gfP{}).random(rand.Reader)
	a := &gfP{0x0a1c7970e5df544d, 0xe74504e9a96b56cc, 0xcda02d92d4d62924, 0x7d2bc576fdf597d1}
	b := &gfP{}
	c := &gfP{}
	fmt.Println("Test Exponent")
	b.exp(a, pMinus1)
	//b.exp(a,[4]uint64{2,0,0,0})
	//fmt.Println("a:", b)
	//gfpMul(c,a,a)
	//fmt.Println("c:",c)
	//gfpSub(b,c,b)
	// montDecode(b, b)
	if *b != *newGFp(1) {
		t.Fail()
		fmt.Println(b)
		return
	}
	fmt.Println("pass")
	var rN1 = &gfP{0x0a1c7970e5df544d, 0xe74504e9a96b56cc, 0xcda02d92d4d62924, 0x7d2bc576fdf597d1}
	fmt.Println("Test Rn")
	gfpMul(b, rN1, r3)
	montDecode(b, b)
	if *b != (gfP{1}) {
		t.Fail()
		fmt.Println(b)
		return
	}
	fmt.Println("pass")

	fmt.Println("Test Invert")
	a = newGFp(1)
	c.Invert(a)
	fmt.Println(a, c)
	gfpMul(c, c, a)
	montDecode(c, c)
	if *c != (gfP{1}) {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")
}

func TestGfP2(t *testing.T) {

	p2Minus1 := new(big.Int)
	p2Minus1.Mul(P, P)
	p2Minus1.Sub(p2Minus1, gmath.BigInt1)

	a := (&gfP2{}).random(rand.Reader)
	b := &gfP2{}
	c := &gfP2{}

	fmt.Println("Test gfP2 mul and square")
	fmt.Println("a:  ", a)
	b.Mul(a, a)
	fmt.Println("a^2:", b)

	c.Square(a)
	fmt.Println("a*a:", c)
	if !c.Sub(c, b).IsZero() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")

	fmt.Println("Test exponential")
	c.Exp(a, p2Minus1)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")

	fmt.Println("Test Invert")
	c.Invert(a)
	c.Mul(c, a)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")
}

func TestGFp2Sqrt(t *testing.T) {
	for i := 0; i < 100000; i++ {
		a := (&gfP2{}).random(rand.Reader)
		a.Mul(a, a)
		b := *a
		if !a.Sqrt(a) {
			t.Fatal()
		}
		a.Square(a)
		if *a != b {
			t.Fatal()
		}
		c := &gfP2{}
		if !c.Sqrt(a) {
			t.Fatal()
		}
		a.Square(c)
		if *a != b {
			t.Fatal()
		}
	}
}

func TestGfP6(t *testing.T) {
	p6Minus1 := new(big.Int)
	p6Minus1.Exp(P, gmath.BigInt6, nil)
	p6Minus1.Sub(p6Minus1, gmath.BigInt1)

	a := (&gfP6{}).random(rand.Reader)
	b := &gfP6{}
	c := &gfP6{}
	fmt.Println(a)

	fmt.Println("Test gfP6 mul and square")
	fmt.Println("a:  ", a)
	b.Mul(a, a)
	fmt.Println("a^2:", b)

	c.Square(a)
	fmt.Println("a*a:", c)
	if !c.Sub(c, b).IsZero() {
		t.Fail()
		fmt.Println(c)
		return
	}

	fmt.Println("Test exponential")

	c.Exp(a, p6Minus1)
	//c.Exp(a,gmath.BigInt1)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")

	fmt.Println("Test Invert")
	c.Invert(a)
	c.Mul(c, a)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")
}

func TestGfP12(t *testing.T) {
	p12Minus1 := new(big.Int)
	p12Minus1.Exp(P, gmath.BigInt12, nil)
	p12Minus1.Sub(p12Minus1, gmath.BigInt1)

	a := (&gfP12{}).random(rand.Reader)
	b := &gfP12{}
	c := &gfP12{}
	fmt.Println(a)

	fmt.Println("Test gfP12 mul and square")
	fmt.Println("a:  ", a)
	b.Mul(a, a)
	fmt.Println("a^2:", b)

	c.Square(a)
	fmt.Println("a*a:", c)
	if !c.Sub(c, b).IsZero() {
		t.Fail()
		fmt.Println(c)
		return
	}

	fmt.Println("Test exponential")

	c.Exp(a, p12Minus1)
	//c.Exp(a,gmath.BigInt1)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")

	fmt.Println("Test Invert")
	c.Invert(a)
	c.Mul(c, a)
	if !c.IsOne() {
		t.Fail()
		fmt.Println(c)
		return
	}
	fmt.Println("pass")

}

func TestFrobenius6(t *testing.T) {
	p2 := new(big.Int).Mul(P, P)
	a := (&gfP6{}).random(rand.Reader)
	b := &gfP6{}
	c := &gfP6{}
	b.Frobenius(a)
	c.Exp(a, P)
	if *b != *c {
		fmt.Println("a:", a)
		fmt.Println("b:", b)
		fmt.Println("c:", c)
		t.Fatal()
	}

	b.FrobeniusP2(a)
	c.Exp(a, p2)
	if *b != *c {
		t.Fatal()
	}

}

func TestFrobenius12(t *testing.T) {
	p2 := new(big.Int).Mul(P, P)
	a := (&gfP12{}).random(rand.Reader)
	b := &gfP12{}
	c := &gfP12{}

	fmt.Println("Test Frobenius on GFP12")
	b.Frobenius(a)
	c.Exp(a, P)
	if *b != *c {
		fmt.Println("a:", a)
		fmt.Println("b:", b)
		fmt.Println("c:", c)
		t.Fail()
		return
	}
	fmt.Println("pass")

	fmt.Println("Test Frobenius2 on GFP12")
	b.FrobeniusP2(a)
	c.Exp(a, p2)
	if *b != *c {
		t.Fail()
		return
	}
	fmt.Println("pass")
}

func present(p *gfP) string {

	return fmt.Sprintf("gfP{0x%x,0x%x,0x%x,0x%x}", p[0], p[1], p[2], p[3])
}

// ExampleGFP12Gen generate gfP12Gen
func TestGFP12Gen(t *testing.T) {
	e := Pair(&G1{*curveGen}, &G2{*twistGen}).p
	fmt.Printf(`
	var gfP12Gen *gfP12 = &gfP12{
		x: gfP6{
			x: gfP2{
				x: %s,
				y: %s,
			},
			y: gfP2{
				x: %s,
				y: %s,
			},
			z: gfP2{
				x: %s,
				y: %s,
			},
		},
		y: gfP6{
			x: gfP2{
				x: %s,
				y: %s,
			},
			y: gfP2{
				x: %s,
				y: %s,
			},
			z: gfP2{
				x: %s,
				y: %s,
			},
		},
	}
	`, present(&e.x.x.x),
		present(&e.x.x.y),
		present(&e.x.y.x),
		present(&e.x.y.y),
		present(&e.x.z.x),
		present(&e.x.z.y),
		present(&e.y.x.x),
		present(&e.y.x.y),
		present(&e.y.y.x),
		present(&e.y.y.y),
		present(&e.y.z.x),
		present(&e.y.z.y))
}