package fpe

import (
	"math"
	"math/big"
	"math/bits"
)

type Radix struct {
	r int

	e int // the maximum number such that r^e is a uint64

	rBig *big.Int // radix as big

	reBig *big.Int // radix^e as big.Int, the same as rPowers[-1]

	rPowers []big.Int // [r, r^2, ..., r^e]
}

func NewRadix(r int) *Radix {
	radix := &Radix{}
	radix.Set(r)
	return radix
}

func (r *Radix) Set(v int) {
	if v < 2 || v >= math.MaxUint16 {
		panic("radix overflow")
	}
	r.r = v
	r.e = 64 / bits.Len32(uint32(v))
	r.rPowers = make([]big.Int, r.e)
	r.rPowers[0].SetUint64(uint64(v))
	for i, vi := 1, uint64(v); i < r.e; i++ {
		vi *= uint64(v)
		r.rPowers[i].SetUint64(vi)
	}
	r.rBig = &r.rPowers[0]
	r.reBig = &r.rPowers[r.e-1]
}

func (r *Radix) Radix() int {
	return r.r
}

func (r *Radix) Num(X []Char) *big.Int {
	if false {
		radix := big.NewInt(int64(r.r))
		n := new(big.Int)
		for _, d := range X {
			n.Mul(n, radix)
			n.Add(n, big.NewInt(int64(d)))
		}
		return n
	} else {
		radix := uint64(r.r)
		l := len(X) % r.e

		tail := uint64(0)
		for i := 0; i < l; i++ {
			tail *= uint64(radix)
			tail += uint64(X[i])
		}

		n := new(big.Int).SetUint64(tail)
		for i := l; i < len(X); i += r.e {
			d := uint64(0)
			for j := 0; j < r.e; j++ {
				d = d*uint64(radix) + uint64(X[i+j])
			}

			n.Mul(n, r.reBig)
			n.Add(n, new(big.Int).SetUint64(d))
		}
		return n
	}
}

// x should < radix^(lenX)
func (r *Radix) Str(X []Numeral, x *big.Int) {
	if false {
		n := len(X)
		radix := big.NewInt(int64(r.r))
		d := new(big.Int)
		xx := new(big.Int).Set(x) // make a copy
		for i := 0; i < n; i++ {
			xx, d = xx.DivMod(xx, radix, d)
			X[n-i-1] = Numeral(d.Int64())
		}
		// if xx is not zero, then the input x must too big
		if xx.Sign() != 0 {
			panic("x should less then radix^m")
		}
	} else {
		xx := new(big.Int).Set(x) // make a copy
		lx := len(X)
		radix := uint64(r.r)
		d := new(big.Int)

		// the tail part
		m := lx % r.e
		if m > 0 {
			xx, d = xx.DivMod(xx, &r.rPowers[m-1], d)
			n := d.Uint64()
			for i := lx - 1; i >= lx-m; i-- {
				X[i] = Numeral(n % radix)
				n /= radix
			}
		}

		for i := lx - m; i > 0; i -= r.e {
			xx, d = xx.DivMod(xx, r.reBig, d)
			n := d.Uint64()
			for j := i - 1; j >= i-r.e; j-- {
				X[j] = uint16(n % radix)
				n /= radix
			}
		}
		// if xx is not zero, then the input x must too big
		if xx.Sign() != 0 {
			panic("x should less then radix^m")
		}
	}
}