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 2×2 matrix [[A, B], [C, D]] over Z_P.
// tmat — torus-constrained 2×2 matrix where B = C. Stores only (A, B, D).

import "math/big"

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

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

// m4Eye returns the identity matrix.
func m4Eye() mat4 {
	return mat4{a: feOne, b: feZero, c: feZero, d: feOne}
}

// tmEye returns the torus identity matrix.
func tmEye() tmat {
	return tmat{a: feOne, b: feZero, d: feOne}
}

// m4Mul computes r = x * y for full 2×2 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
}

// tmIsIdentity returns true if m is the identity torus matrix.
func tmIsIdentity(m *tmat) bool {
	return feEqual(&m.a, &feOne) == 1 &&
		feIsZero(&m.b) == 1 &&
		feEqual(&m.d, &feOne) == 1
}

// tmEqual returns true if a and b represent the same torus matrix.
func tmEqual(a, b *tmat) bool {
	return feEqual(&a.a, &b.a) == 1 &&
		feEqual(&a.b, &b.b) == 1 &&
		feEqual(&a.d, &b.d) == 1
}

// tmInv computes r = m^{-1} for a torus matrix.
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
}

// tmToMat4 converts a torus matrix to a full matrix (setting C = B).
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)
}

// tmFromMat4 converts a full matrix to a torus matrix (assumes B == C).
func tmFromMat4(r *tmat, m *mat4) {
	feSet(&r.a, &m.a)
	feSet(&r.b, &m.b)
	feSet(&r.d, &m.d)
}

// m4Det computes det(m) = a*d - b*c in Montgomery form.
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)
}

// m4Trace computes trace(m) = a + d in Montgomery form.
func m4Trace(r *fe, m *mat4) {
	feAdd(r, &m.a, &m.d)
}

// m4Equal returns true if two full matrices are equal.
func m4Equal(a, b *mat4) 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
}

// m4IsIdentity returns true if m is the identity matrix.
func m4IsIdentity(m *mat4) bool {
	return feEqual(&m.a, &feOne) == 1 &&
		feIsZero(&m.b) == 1 &&
		feIsZero(&m.c) == 1 &&
		feEqual(&m.d, &feOne) == 1
}

// m4Inv computes the inverse of an SL(2) matrix: [[d, -b], [-c, a]].
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
}

// m4FromSmall creates a mat4 from small integer values.
func m4FromSmall(a, b, c, d int64) mat4 {
	var r mat4
	feFromSmall(&r.a, a)
	feFromSmall(&r.b, b)
	feFromSmall(&r.c, c)
	feFromSmall(&r.d, d)
	return r
}

// tmToBytes serializes a torus matrix as 81 bytes (A, B, D as 3×27 bytes big-endian).
func tmToBytes(buf []byte, m *tmat) {
	feToBytes27(buf[0:27], &m.a)
	feToBytes27(buf[27:54], &m.b)
	feToBytes27(buf[54:81], &m.d)
}

// tmFromBytes deserializes a torus matrix from 81 bytes.
func tmFromBytes(m *tmat, buf []byte) bool {
	if 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 := 0; i < exp.BitLen(); i++ {
		if exp.Bit(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) {
	pBig := new(big.Int)
	var pbuf [27]byte
	feToBytes27(pbuf[:], &feZero)
	// Actually we need P as big.Int. Reconstruct from pLimbs.
	pBig.SetUint64(pLimbs[3])
	pBig.Lsh(pBig, 64)
	pBig.Or(pBig, new(big.Int).SetUint64(pLimbs[2]))
	pBig.Lsh(pBig, 64)
	pBig.Or(pBig, new(big.Int).SetUint64(pLimbs[1]))
	pBig.Lsh(pBig, 64)
	pBig.Or(pBig, new(big.Int).SetUint64(pLimbs[0]))

	norm := new(big.Int).Mod(v, pBig)
	var buf [27]byte
	normBytes := norm.Bytes()
	copy(buf[27-len(normBytes):], normBytes)
	// Decode big-endian into little-endian limbs.
	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)
}

// feToBig converts a Montgomery field element to a big.Int.
func feToBig(a *fe) *big.Int {
	var buf [27]byte
	feToBytes27(buf[:], a)
	return new(big.Int).SetBytes(buf[:])
}

// bigToScalar converts a big.Int to a scalar mod Q.
func bigToScalar(r *scalar, v *big.Int) {
	norm := new(big.Int).Mod(v, qBig)
	var buf [27]byte
	normBytes := norm.Bytes()
	copy(buf[27-len(normBytes):], normBytes)
	scFromBytes27(r, buf[:])
}

// scalarToBig converts a scalar to a big.Int.
func scalarToBig(s *scalar) *big.Int {
	var buf [27]byte
	scToBytes27(buf[:], s)
	return new(big.Int).SetBytes(buf[:])
}