{{if .IsG2 -}}
//go:build g2

{{end -}}
package {{.PackageName}}

import (
	"github.com/stretchr/testify/assert"
	"testing"

	"github.com/ingonyama-zk/icicle/wrappers/golang/core"
	cr "github.com/ingonyama-zk/icicle/wrappers/golang/cuda_runtime"

	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark-crypto/ecc/{{.GnarkImport}}"
	"github.com/consensys/gnark-crypto/ecc/{{.GnarkImport}}/fp"
	"github.com/consensys/gnark-crypto/ecc/{{.GnarkImport}}/fr"
)
{{$isBW6 := eq .Curve "bw6_761"}}{{$isG1 := not .IsG2}}{{if or $isBW6 $isG1}}
func projectiveToGnarkAffine{{if and $isBW6 .IsG2}}G2{{end}}(p {{if and $isBW6 .IsG2}}G2{{end}}Projective) {{toPackage .GnarkImport}}.{{if and $isBW6 .IsG2}}G2{{else}}G1{{end}}Affine {
	px, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.X).ToBytesLittleEndian()))
	py, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.Y).ToBytesLittleEndian()))
	pz, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.Z).ToBytesLittleEndian()))

	zInv := new(fp.Element)
	x := new(fp.Element)
	y := new(fp.Element)

	zInv.Inverse(&pz)

	x.Mul(&px, zInv)
	y.Mul(&py, zInv)

	return {{toPackage .GnarkImport}}.{{if and $isBW6 .IsG2}}G2{{else}}G1{{end}}Affine{X: *x, Y: *y}
}
{{end}}
{{- $isNotBW6 := ne .Curve "bw6_761"}}{{if and $isNotBW6 .IsG2 }}
func projectiveToGnarkAffineG2(p G2Projective) {{toPackage .GnarkImport}}.G2Affine {
	pxBytes := p.X.ToBytesLittleEndian()
	pxA0, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pxBytes[:fp.Bytes]))
	pxA1, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pxBytes[fp.Bytes:]))
	x := {{toPackage .GnarkImport}}.E2{
		A0: pxA0,
		A1: pxA1,
	}

	pyBytes := p.Y.ToBytesLittleEndian()
	pyA0, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pyBytes[:fp.Bytes]))
	pyA1, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pyBytes[fp.Bytes:]))
	y := {{toPackage .GnarkImport}}.E2{
		A0: pyA0,
		A1: pyA1,
	}

	pzBytes := p.Z.ToBytesLittleEndian()
	pzA0, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pzBytes[:fp.Bytes]))
	pzA1, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)(pzBytes[fp.Bytes:]))
	z := {{toPackage .GnarkImport}}.E2{
		A0: pzA0,
		A1: pzA1,
	}

	var zSquared {{toPackage .GnarkImport}}.E2
	zSquared.Mul(&z, &z)

	var X {{toPackage .GnarkImport}}.E2
	X.Mul(&x, &z)

	var Y {{toPackage .GnarkImport}}.E2
	Y.Mul(&y, &zSquared)

	g2Jac := {{toPackage .GnarkImport}}.G2Jac{
		X: X,
		Y: Y,
		Z: z,
	}

	var g2Affine {{toPackage .GnarkImport}}.G2Affine
	return *g2Affine.FromJacobian(&g2Jac)
}
{{end}}
func testAgainstGnarkCryptoMsm{{if .IsG2}}G2{{end}}(scalars core.HostSlice[ScalarField], points core.HostSlice[{{if .IsG2}}G2{{end}}Affine], out {{if .IsG2}}G2{{end}}Projective) bool {
	scalarsFr := make([]fr.Element, len(scalars))
	for i, v := range scalars {
		slice64, _ := fr.LittleEndian.Element((*[fr.Bytes]byte)(v.ToBytesLittleEndian()))
		scalarsFr[i] = slice64
	}

	pointsFp := make([]{{toPackage .GnarkImport}}.{{if .IsG2}}G2{{else}}G1{{end}}Affine, len(points))
	for i, v := range points {
		pointsFp[i] = projectiveToGnarkAffine{{if .IsG2}}G2{{end}}(v.ToProjective())
	}
	var msmRes {{toPackage .GnarkImport}}.{{if .IsG2}}G2{{else}}G1{{end}}Jac
	msmRes.MultiExp(pointsFp, scalarsFr, ecc.MultiExpConfig{})

	var icicleResAsJac {{toPackage .GnarkImport}}.{{if .IsG2}}G2{{else}}G1{{end}}Jac
	proj := projectiveToGnarkAffine{{if .IsG2}}G2{{end}}(out)
	icicleResAsJac.FromAffine(&proj)

	return msmRes.Equal(&icicleResAsJac)
}

func TestMSM{{if .IsG2}}G2{{end}}(t *testing.T) {
	cfg := GetDefaultMSMConfig()
	for _, power := range []int{2, 3, 4, 5, 6, 7, 8, 10, 18} {
		size := 1 << power

		scalars := GenerateScalars(size)
		points := {{if .IsG2}}G2{{end}}GenerateAffinePoints(size)

		stream, _ := cr.CreateStream()
		var p {{if .IsG2}}G2{{end}}Projective
		var out core.DeviceSlice
		_, e := out.MallocAsync(p.Size(), p.Size(), stream)
		assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
		cfg.Ctx.Stream = &stream
		e = {{if .IsG2}}G2{{end}}Msm(scalars, points, &cfg, out)
		assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
		outHost := make(core.HostSlice[{{if .IsG2}}G2{{end}}Projective], 1)
		outHost.CopyFromDevice(&out)
		out.Free()

		// Check with gnark-crypto
		assert.True(t, testAgainstGnarkCryptoMsm{{if .IsG2}}G2{{end}}(scalars, points, outHost[0]))
	}
}

func TestMSM{{if .IsG2}}G2{{end}}Batch(t *testing.T) {
	cfg := GetDefaultMSMConfig()
	for _, power := range []int{10, 16} {
		for _, batchSize := range []int{1, 3, 16} {
			size := 1 << power
			totalSize := size * batchSize
			scalars := GenerateScalars(totalSize)
			points := {{if .IsG2}}G2{{end}}GenerateAffinePoints(totalSize)

			var p {{if .IsG2}}G2{{end}}Projective
			var out core.DeviceSlice
			_, e := out.Malloc(batchSize*p.Size(), p.Size())
			assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")

			e = {{if .IsG2}}G2{{end}}Msm(scalars, points, &cfg, out)
			assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
			outHost := make(core.HostSlice[{{if .IsG2}}G2{{end}}Projective], batchSize)
			outHost.CopyFromDevice(&out)
			out.Free()

			// Check with gnark-crypto
			for i := 0; i < batchSize; i++ {
				scalarsSlice := scalars[i*size : (i+1)*size]
				pointsSlice := points[i*size : (i+1)*size]
				out := outHost[i]
				assert.True(t, testAgainstGnarkCryptoMsm{{if .IsG2}}G2{{end}}(scalarsSlice, pointsSlice, out))
			}
		}
	}
}

func TestMSM{{if .IsG2}}G2{{end}}SkewedDistribution(t *testing.T) {
	cfg := GetDefaultMSMConfig()
	for _, power := range []int{2, 3, 4, 5, 6, 7, 8, 10, 18} {
		size := 1 << power

		scalars := GenerateScalars(size)
		for i := size / 4; i < size; i++ {
			scalars[i].One()
		}
		points := {{if .IsG2}}G2{{end}}GenerateAffinePoints(size)
		for i := 0; i < size/4; i++ {
			points[i].Zero()
		}

		var p {{if .IsG2}}G2{{end}}Projective
		var out core.DeviceSlice
		_, e := out.Malloc(p.Size(), p.Size())
		assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")

		e = {{if .IsG2}}G2{{end}}Msm(scalars, points, &cfg, out)
		assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
		outHost := make(core.HostSlice[{{if .IsG2}}G2{{end}}Projective], 1)
		outHost.CopyFromDevice(&out)
		out.Free()

		// Check with gnark-crypto
		assert.True(t, testAgainstGnarkCryptoMsm{{if .IsG2}}G2{{end}}(scalars, points, outHost[0]))
	}
}
