package cayley

import (
	"fmt"
	"math"
	"math/rand"
	"testing"
)

type WindowedSigner struct {
	GB       *GenerativeBasis
	StdWords map[Mat2][]int8
	W        int
	StdPF    *PathFinder
}

func BuildWindowedSigner(gb *GenerativeBasis, W int) *WindowedSigner {
	stdPF := StandardGens(gb.GS.P).BFS(-1)
	ws := &WindowedSigner{GB: gb, W: W, StdPF: stdPF,
		StdWords: make(map[Mat2][]int8, int(math.Pow(4, float64(W))))}
	for v, d := range stdPF.Dist {
		if d > W { continue }
		path, ok := gb.GPF.PathTo(v)
		if ok { ws.StdWords[v] = path }
	}
	return ws
}

func (ws *WindowedSigner) Sign(target Mat2) ([]int8, bool) {
	P := ws.GB.GS.P
	eucPath, ok := ws.StdPF.PathTo(target)
	if !ok { return nil, false }

	var sig []int8
	for i := 0; i < len(eucPath); i += ws.W {
		end := i + ws.W
		if end > len(eucPath) { end = len(eucPath) }

		// Block contribution: product of block steps FROM IDENTITY.
		blockContribution := ID()
		for j := i; j < end; j++ {
			var gen Mat2
			switch eucPath[j] {
			case 0: gen = modMat2(StdG0, P)
			case 1: gen = modMat2(StdG1, P)
			case 2: gen = modMat2(StdG0I, P)
			case 3: gen = modMat2(StdG1I, P)
			default: return nil, false
			}
			blockContribution = blockContribution.Mul(gen, P)
		}

		word, ok := ws.StdWords[blockContribution]
		if !ok { return nil, false }
		sig = append(sig, word...)
	}
	if ws.GB.GS.Walk(ID(), sig).Eq(target) { return sig, true }
	return nil, false
}

func TestWindowedDebug(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") }

	ws := BuildWindowedSigner(gb, 1)
	fmt.Printf("W=1 stdWords count: %d\n", len(ws.StdWords))

	target := randomSL2Fast(P, rng)
	eucPath, ok := ws.StdPF.PathTo(target)
	if !ok { t.Fatal("target not reachable") }
	fmt.Printf("Euclidean path: %d steps\n", len(eucPath))

	sig, ok := ws.Sign(target)
	if !ok { t.Fatal("sign failed") }
	fmt.Printf("Signature: %d steps\n", len(sig))
	fmt.Printf("Verify: %v\n", gs.Walk(ID(), sig).Eq(target))
}

func TestWindowedBlowup(t *testing.T) {
	P := int64(251)
	if testing.Short() { t.Skip("P=251 BFS slow") }
	rng := rand.New(rand.NewSource(42))
	gs := randomGens(P, 4, rng)
	gb := BuildGenerativeBasis(gs)
	if gb == nil { t.Fatal("no basis") }

	for _, W := range []int{1, 2, 4} {
		ws := BuildWindowedSigner(gb, W)
		var sigT, optT int; n := 0
		for i := 0; i < 20; i++ {
			target := randomSL2Fast(P, rng)
			opt, ok := gb.GPF.PathTo(target)
			if !ok { continue }
			sig, ok := ws.Sign(target)
			if !ok { continue }
			sigT += len(sig); optT += len(opt); n++
		}
		if n > 0 {
			t.Logf("W=%d blocks=%d n=%d: sig=%.1f opt=%.1f ratio=%.2f",
				W, len(ws.StdWords), n, float64(sigT)/float64(n),
				float64(optT)/float64(n), float64(sigT)/float64(optT))
		}
	}
}