package cayley

import (
	"math/rand"
	"testing"
)

// compressBacktrack removes adjacent inverse generator pairs.
func compressBacktrack(gs *GeneratorSet, path []int8) []int8 {
	inv := make(map[int8]int8)
	for i := int8(0); i < int8(len(gs.Gens)); i++ {
		invMat := gs.Gens[i].Inv(gs.P)
		for j := int8(0); j < int8(len(gs.Gens)); j++ {
			if gs.Gens[j].Eq(invMat) { inv[i] = j; break }
		}
	}
	stack := make([]int8, 0, len(path))
	for _, step := range path {
		if len(stack) > 0 && inv[stack[len(stack)-1]] == step {
			stack = stack[:len(stack)-1]
		} else {
			stack = append(stack, step)
		}
	}
	return stack
}

// compressGreedy optimizes a valid path by greedy step choice.
func compressGreedy(gs *GeneratorSet, path []int8, target Mat2) []int8 {
	pos := ID()
	result := make([]int8, 0, len(path))
	remaining := len(path)

	for !pos.Eq(target) && remaining >= 0 {
		best := int8(-1)
		bestDist := int64(1<<63 - 1)
		for i := int8(0); i < int8(len(gs.Gens)); i++ {
			next := pos.Mul(gs.Gens[i], gs.P)
			dist := matDist(next, target, gs.P)
			if dist < bestDist { bestDist = dist; best = i }
		}
		if best < 0 { break }
		next := pos.Mul(gs.Gens[best], gs.P)
		if next.Eq(pos) { break }
		result = append(result, best)
		pos = next
		remaining--
	}
	if pos.Eq(target) { return result }
	return path // fallback to original
}

func matDist(a, b Mat2, P int64) int64 {
	da := mod(a.A-b.A, P); db := mod(a.B-b.B, P)
	dc := mod(a.C-b.C, P); dd := mod(a.D-b.D, P)
	return int64(da)*int64(da) + int64(db)*int64(db) + int64(dc)*int64(dc) + int64(dd)*int64(dd)
}

// compressWindow BFS-compresses path within windows of size w.
func compressWindow(gs *GeneratorSet, path []int8, target Mat2, w int) []int8 {
	if w <= 1 || len(path) <= w { return path }
	pos := ID()
	result := make([]int8, 0, len(path))
	b := len(gs.Gens)

	for i := 0; i < len(path); {
		end := i + w
		if end > len(path) { end = len(path) }
		endPos := pos
		for j := i; j < end; j++ { endPos = endPos.Mul(gs.Gens[path[j]], gs.P) }

		best := path[i:end]
		bestLen := end - i

		// Enumerate all walks ≤w between pos and endPos.
		type node struct {
			m Mat2
			p []int8
		}
		queue := []node{{pos, nil}}
		seen := map[Mat2]bool{pos: true}

		for len(queue) > 0 && len(queue) < 50000 {
			n := queue[0]; queue = queue[1:]
			if len(n.p) >= bestLen { continue }
			for k := 0; k < b; k++ {
				next := n.m.Mul(gs.Gens[k], gs.P)
				if next.Eq(endPos) {
					np := make([]int8, len(n.p)+1)
					copy(np, n.p); np[len(n.p)] = int8(k)
					best = np; bestLen = len(np)
					goto done
				}
				if len(n.p)+1 >= bestLen-1 { continue }
				if !seen[next] {
					seen[next] = true
					np := make([]int8, len(n.p)+1)
					copy(np, n.p); np[len(n.p)] = int8(k)
					queue = append(queue, node{next, np})
				}
			}
		}
	done:
		result = append(result, best...)
		pos = endPos
		i = end
	}
	if gs.Walk(ID(), result).Eq(target) { return result }
	return path
}

// --- Tests ---

func makeLongPath(target Mat2, gb *GenerativeBasis, P int64, rng *rand.Rand) []int8 {
	euc, err := EuclideanDecomposition(target, P)
	if err != nil { return nil }
	return gb.Sign(euc)
}

func TestCompressBacktrack(t *testing.T) {
	P := int64(31); rng := rand.New(rand.NewSource(42))
	gs := randomGens(P, 4, rng)
	gb := BuildGenerativeBasis(gs)
	if gb == nil { t.Fatal("no basis") }
	pf := gb.GPF

	var origT, newT int; n := 0
	for i := 0; i < 50; i++ {
		target := randomSL2Fast(P, rng)
		path := makeLongPath(target, gb, P, rng)
		if path == nil || len(path) > 400 { continue }
		opt := compressBacktrack(gs, path)
		if !gs.Walk(ID(), opt).Eq(target) { continue }
		origT += len(path); newT += len(opt); n++
	}
	if n == 0 { t.Log("no samples"); return }
	o := float64(origT)/float64(n); nn := float64(newT)/float64(n)
	om := pf.DistStatsMean()
	t.Logf("BACKTRACK: orig=%.0f new=%.0f opt=%.0f ratio=%.2f red=%.0f%% n=%d",
		o, nn, om, nn/o, (1-nn/o)*100, n)
}

func TestCompressGreedy(t *testing.T) {
	P := int64(31); rng := rand.New(rand.NewSource(42))
	gs := randomGens(P, 4, rng)
	gb := BuildGenerativeBasis(gs)
	if gb == nil { t.Fatal("no basis") }
	pf := gb.GPF

	var origT, newT int; n := 0
	for i := 0; i < 50; i++ {
		target := randomSL2Fast(P, rng)
		path := makeLongPath(target, gb, P, rng)
		if path == nil || len(path) > 400 { continue }
		opt := compressGreedy(gs, path, target)
		if !gs.Walk(ID(), opt).Eq(target) { continue }
		origT += len(path); newT += len(opt); n++
	}
	if n == 0 { t.Log("no samples"); return }
	o := float64(origT)/float64(n); nn := float64(newT)/float64(n)
	om := pf.DistStatsMean()
	t.Logf("GREEDY:    orig=%.0f new=%.0f opt=%.0f ratio=%.2f red=%.0f%% n=%d",
		o, nn, om, nn/o, (1-nn/o)*100, n)
}

func TestCompressWindow(t *testing.T) {
	P := int64(31); rng := rand.New(rand.NewSource(42))
	gs := randomGens(P, 4, rng)
	gb := BuildGenerativeBasis(gs)
	if gb == nil { t.Fatal("no basis") }
	pf := gb.GPF

	for _, w := range []int{3, 4, 5} {
		var origT, newT int; n := 0
		for i := 0; i < 30; i++ {
			target := randomSL2Fast(P, rng)
			path := makeLongPath(target, gb, P, rng)
			if path == nil || len(path) > 400 { continue }
			opt := compressWindow(gs, path, target, w)
			if !gs.Walk(ID(), opt).Eq(target) { continue }
			origT += len(path); newT += len(opt); n++
		}
		if n == 0 { continue }
		o := float64(origT)/float64(n); nn := float64(newT)/float64(n)
		om := pf.DistStatsMean()
		t.Logf("WINDOW w=%d: orig=%.0f new=%.0f opt=%.0f ratio=%.2f red=%.0f%% n=%d",
			w, o, nn, om, nn/o, (1-nn/o)*100, n)
	}
}

func TestCompressCompound(t *testing.T) {
	P := int64(31); rng := rand.New(rand.NewSource(42))
	gs := randomGens(P, 4, rng)
	gb := BuildGenerativeBasis(gs)
	if gb == nil { t.Fatal("no basis") }
	pf := gb.GPF

	nSamples := 30
	var origT, btT, grT, winT int; n := 0
	for i := 0; i < nSamples; i++ {
		target := randomSL2Fast(P, rng)
		path := makeLongPath(target, gb, P, rng)
		if path == nil || len(path) > 400 { continue }

		// Compound: backtrack → greedy → window(4).
		p1 := compressBacktrack(gs, path)
		p2 := compressGreedy(gs, p1, target)
		p3 := compressWindow(gs, p2, target, 4)
		if !gs.Walk(ID(), p3).Eq(target) { continue }

		origT += len(path); btT += len(p1); grT += len(p2); winT += len(p3); n++
	}
	if n == 0 { return }
	o := float64(origT)/float64(n)
	t.Logf("COMPOUND P=%d n=%d:", P, n)
	t.Logf("  original:         %.0f", o)
	t.Logf("  +backtrack:       %.0f (%.0f%%)", float64(btT)/float64(n), float64(btT)/float64(origT)*100)
	t.Logf("  +greedy:          %.0f (%.0f%%)", float64(grT)/float64(n), float64(grT)/float64(origT)*100)
	t.Logf("  +window(4):       %.0f (%.0f%%)", float64(winT)/float64(n), float64(winT)/float64(origT)*100)
	t.Logf("  BFS optimal:      %.0f", pf.DistStatsMean())
}