package gnarl

// Matrix types for the Gnarl signature scheme using Montgomery field elements.
//
// Operates over the Gnarl prime P
// (216 bits). The non-split torus has order N = P+1 = 6Q, and the Schnorr
// subgroup has prime order Q = (P+1)/6.
//
// mat4 - full 2x2 matrix [[A, B], [C, D]] over Z_P.
// tmat - torus-constrained 2x2 matrix where B = C. Stores only (A, B, D).

import "math/big"

// mat4 is a full 2x2 matrix over Z_P in Montgomery form.
type mat4 struct {
	a, b, c, d fe
}

// tmat is a torus-constrained 2x2 matrix where B = C, in Montgomery form.
type tmat struct {
	a, b, d fe
}

func m4Eye() (r mat4) {
	return mat4{a: feOne, b: feZero, c: feZero, d: feOne}
}

func tmEye() (r tmat) {
	return tmat{a: feOne, b: feZero, d: feOne}
}

// m4Mul computes r = x * y for full 2x2 matrices. 8 field multiplications.
func m4Mul(r, x, y *mat4) {
	var ra, rb, rc, rd, t1, t2 fe
	montMul(&t1, &x.a, &y.a)
	montMul(&t2, &x.b, &y.c)
	feAdd(&ra, &t1, &t2)

	montMul(&t1, &x.a, &y.b)
	montMul(&t2, &x.b, &y.d)
	feAdd(&rb, &t1, &t2)

	montMul(&t1, &x.c, &y.a)
	montMul(&t2, &x.d, &y.c)
	feAdd(&rc, &t1, &t2)

	montMul(&t1, &x.c, &y.b)
	montMul(&t2, &x.d, &y.d)
	feAdd(&rd, &t1, &t2)

	r.a = ra
	r.b = rb
	r.c = rc
	r.d = rd
}

// tmMul computes r = x * y for torus matrices (B=C constraint preserved).
// 5 field multiplications.
func tmMul(r, x, y *tmat) {
	var aa, bb, dd, ab, bd, ra, rb, rd fe

	montMul(&aa, &x.a, &y.a)
	montMul(&bb, &x.b, &y.b)
	montMul(&dd, &x.d, &y.d)
	montMul(&ab, &x.a, &y.b)
	montMul(&bd, &x.b, &y.d)

	feAdd(&ra, &aa, &bb)
	feAdd(&rb, &ab, &bd)
	feAdd(&rd, &bb, &dd)

	r.a = ra
	r.b = rb
	r.d = rd
}

// tmSquare computes r = x^2 for a torus matrix.
func tmSquare(r, x *tmat) {
	var a2, b2, d2, apd, ra, rb, rd fe

	montSquare(&a2, &x.a)
	montSquare(&b2, &x.b)
	montSquare(&d2, &x.d)
	feAdd(&apd, &x.a, &x.d)
	montMul(&rb, &x.b, &apd)

	feAdd(&ra, &a2, &b2)
	feAdd(&rd, &b2, &d2)

	r.a = ra
	r.b = rb
	r.d = rd
}

func tmIsIdentity(m *tmat) (result bool) {
	return feEqual(&m.a, &feOne) == 1 &&
		feIsZero(&m.b) == 1 &&
		feEqual(&m.d, &feOne) == 1
}

func tmEqual(a, b *tmat) (result bool) {
	return feEqual(&a.a, &b.a) == 1 &&
		feEqual(&a.b, &b.b) == 1 &&
		feEqual(&a.d, &b.d) == 1
}

func tmInv(r, m *tmat) {
	var ra, rb, rd fe
	feSet(&ra, &m.d)
	feNeg(&rb, &m.b)
	feSet(&rd, &m.a)
	r.a = ra
	r.b = rb
	r.d = rd
}

func tmToMat4(r *mat4, m *tmat) {
	feSet(&r.a, &m.a)
	feSet(&r.b, &m.b)
	feSet(&r.c, &m.b)
	feSet(&r.d, &m.d)
}

func tmFromMat4(r *tmat, m *mat4) {
	feSet(&r.a, &m.a)
	feSet(&r.b, &m.b)
	feSet(&r.d, &m.d)
}

func m4Det(r *fe, m *mat4) {
	var ad, bc fe
	montMul(&ad, &m.a, &m.d)
	montMul(&bc, &m.b, &m.c)
	feSub(r, &ad, &bc)
}

func m4Trace(r *fe, m *mat4) {
	feAdd(r, &m.a, &m.d)
}

func m4Equal(a, b *mat4) (result bool) {
	return feEqual(&a.a, &b.a) == 1 &&
		feEqual(&a.b, &b.b) == 1 &&
		feEqual(&a.c, &b.c) == 1 &&
		feEqual(&a.d, &b.d) == 1
}

func m4IsIdentity(m *mat4) (result bool) {
	return feEqual(&m.a, &feOne) == 1 &&
		feIsZero(&m.b) == 1 &&
		feIsZero(&m.c) == 1 &&
		feEqual(&m.d, &feOne) == 1
}

func m4Inv(r, m *mat4) {
	var ra, rb, rc, rd fe
	feSet(&ra, &m.d)
	feNeg(&rb, &m.b)
	feNeg(&rc, &m.c)
	feSet(&rd, &m.a)
	r.a = ra
	r.b = rb
	r.c = rc
	r.d = rd
}

func m4FromSmall(a, b, c, d int64) (r mat4) {
	feFromSmall(&r.a, a)
	feFromSmall(&r.b, b)
	feFromSmall(&r.c, c)
	feFromSmall(&r.d, d)
	return r
}

func tmToBytes(buf []byte, m *tmat) {
	feToBytes27(buf[0:27], &m.a)
	feToBytes27(buf[27:54], &m.b)
	feToBytes27(buf[54:81], &m.d)
}

func tmFromBytes(m *tmat, buf []byte) (ok bool) {
	if int32(len(buf)) < 81 {
		return false
	}
	if !feFromBytes27(&m.a, buf[0:27]) {
		return false
	}
	if !feFromBytes27(&m.b, buf[27:54]) {
		return false
	}
	if !feFromBytes27(&m.d, buf[54:81]) {
		return false
	}
	return true
}

// m4PowBig computes r = base^exp for a full matrix with an arbitrary big.Int exponent.
func m4PowBig(r *mat4, base *mat4, exp *big.Int) {
	*r = m4Eye()
	var b mat4
	b = *base

	for i := int32(0); i < int32(exp.BitLen()); i++ {
		if exp.Bit(int32(i)) == 1 {
			m4Mul(r, r, &b)
		}
		m4Mul(&b, &b, &b)
	}
}

// bigToFe converts a big.Int to a Montgomery field element.
func bigToFe(r *fe, v *big.Int) {
	pBigLocal := &big.Int{}
	pBigLocal.SetUint64(pLimbs[3])
	pBigLocal.Lsh(pBigLocal, 64)
	pBigLocal.Or(pBigLocal, (&big.Int{}).SetUint64(pLimbs[2]))
	pBigLocal.Lsh(pBigLocal, 64)
	pBigLocal.Or(pBigLocal, (&big.Int{}).SetUint64(pLimbs[1]))
	pBigLocal.Lsh(pBigLocal, 64)
	pBigLocal.Or(pBigLocal, (&big.Int{}).SetUint64(pLimbs[0]))

	norm := (&big.Int{}).Mod(v, pBigLocal)
	var buf [27]byte
	normBytes := norm.Bytes()
	copy(buf[27-int32(len(normBytes)):], normBytes)
	r[3] = uint64(buf[0])<<16 | uint64(buf[1])<<8 | uint64(buf[2])
	r[2] = uint64(buf[3])<<56 | uint64(buf[4])<<48 | uint64(buf[5])<<40 | uint64(buf[6])<<32 |
		uint64(buf[7])<<24 | uint64(buf[8])<<16 | uint64(buf[9])<<8 | uint64(buf[10])
	r[1] = uint64(buf[11])<<56 | uint64(buf[12])<<48 | uint64(buf[13])<<40 | uint64(buf[14])<<32 |
		uint64(buf[15])<<24 | uint64(buf[16])<<16 | uint64(buf[17])<<8 | uint64(buf[18])
	r[0] = uint64(buf[19])<<56 | uint64(buf[20])<<48 | uint64(buf[21])<<40 | uint64(buf[22])<<32 |
		uint64(buf[23])<<24 | uint64(buf[24])<<16 | uint64(buf[25])<<8 | uint64(buf[26])
	feToMont(r, r)
}

func feToBig(a *fe) (result *big.Int) {
	var buf [27]byte
	feToBytes27(buf[:], a)
	return (&big.Int{}).SetBytes(buf[:])
}

func bigToScalar(r *scalar, v *big.Int) {
	norm := (&big.Int{}).Mod(v, qBig)
	var buf [27]byte
	normBytes := norm.Bytes()
	copy(buf[27-int32(len(normBytes)):], normBytes)
	scFromBytes27(r, buf[:])
}

func scalarToBig(s *scalar) (result *big.Int) {
	var buf [27]byte
	scToBytes27(buf[:], s)
	return (&big.Int{}).SetBytes(buf[:])
}