package main

// math.go: mathematical routines

import (
	"math"
)

func powInt(x, y int) int {
	return int(math.Pow(float64(x), float64(y)))
}

func egcd(a, b bigint) (g, x, y bigint) {
	if a.Empty() {
		return b, NewBigint(0), NewBigint(1)
	} else {
		g, y, x = egcd(b.Mod(a), a)
		return g, x.Sub(b.Div(a).Mul(y)), y
	}
}

func modinv(a, p bigint) bigint {
	if a.LtInt(0) {
		a = a.Mod(p)
	}

	g, x, _ := egcd(a, p)
	if g.NeInt(1) {
		panic("modular inverse does not exist")
	} else {
		return x.Mod(p)
	}
}

func leftmostBit(x bigint) bigint {
	if x.LteInt(0) {
		panic("x must be greater than 0")
	}

	res := NewBigint(1)
	for res.Lte(x) {
		res = res.MulInt(2)
	}

	return res.DivInt(2)
}

func naf(mult bigint) []bigint {
	ret := make([]bigint, 0)

	for mult.GtInt(0) {
		if mult.ModInt(2).GtInt(0) {
			nd := mult.ModInt(4)
			if nd.GteInt(2) {
				nd = nd.SubInt(4)
			}
			ret = append(ret, nd)
			mult = mult.Sub(nd)
		} else {
			ret = append(ret, NewBigint(0))
		}

		mult = mult.DivInt(2)
	}

	return ret
}

func legendreSymbol(a, p bigint) bigint {
	l := a.ModExp(p.SubInt(1).DivInt(2), p)
	if l.Eq(p.SubInt(1)) {
		return NewBigint(-1)
	} else {
		return l
	}
}

func primeModSqrt(a, p bigint) (bigint, bigint) {
	a = a.Mod(p)

	if a.EqInt(0) {
		return NewBigint(0), NewEmptyBigint()
	}

	if p.EqInt(2) {
		return a, NewEmptyBigint()
	}

	if legendreSymbol(a, p).NeInt(1) {
		return NewEmptyBigint(), NewEmptyBigint()
	}

	if p.ModInt(4).EqInt(3) {
		x := a.ModExp(p.AddInt(1).DivInt(4), p)
		return x, p.Sub(x)
	}

	q, s := p.SubInt(1), 0
	for q.ModInt(2).EqInt(0) {
		s++
		q = q.DivInt(2)
	}

	z := NewBigint(1)
	for legendreSymbol(z, p).NeInt(-1) {
		z = z.AddInt(1)
	}

	c := z.ModExp(q, p)
	x := a.ModExp(q.AddInt(1).DivInt(2), p)
	t := a.ModExp(q, p)
	m := s

	for t.NeInt(1) {
		i, e := 0, NewBigint(2)
		if m > 0 {
			for i = 1; i < m; i++ {
				if t.ModExp(e, p).EqInt(1) {
					break
				}
				e = e.MulInt(2)
			}
		}

		b := c.ModExp(NewBigint(powInt(2, m-i-1)), p)
		x = x.Mul(b).Mod(p)
		t = t.Mul(b).Mul(b).Mod(p)
		c = b.Mul(b).Mod(p)
		m = i
	}

	return x, p.Sub(x)
}