package cayley

import (
	"fmt"
	"sort"
)

// PathFinder stores precomputed shortest paths from root to all vertices
// in the Cayley graph of SL(2, Z_P) with a given generator set.
type PathFinder struct {
	GS     *GeneratorSet
	P      int64
	Dist   map[Mat2]int    // distance from root (ID) to each vertex
	Parent map[Mat2]Mat2   // parent in BFS tree
	Edge   map[Mat2]int8   // generator index from parent to this vertex
	BFSMax int              // maximum distance found
}

// BFS computes shortest paths from ID to all reachable vertices via BFS.
// Returns the number of vertices reached and the maximum distance.
func (gs *GeneratorSet) BFS(maxDist int) *PathFinder {
	pf := &PathFinder{
		GS:     gs,
		P:      gs.P,
		Dist:   make(map[Mat2]int),
		Parent: make(map[Mat2]Mat2),
		Edge:   make(map[Mat2]int8),
	}

	root := ID()
	pf.Dist[root] = 0
	queue := []Mat2{root}
	pf.BFSMax = 0

	for len(queue) > 0 {
		cur := queue[0]
		queue = queue[1:]
		d := pf.Dist[cur]

		if maxDist >= 0 && d >= maxDist {
			continue
		}

		for i, g := range gs.Gens {
			next := cur.Mul(g, gs.P)
			if _, ok := pf.Dist[next]; !ok {
				pf.Dist[next] = d + 1
				pf.Parent[next] = cur
				pf.Edge[next] = int8(i)
				queue = append(queue, next)
				if d+1 > pf.BFSMax {
					pf.BFSMax = d + 1
				}
			}
		}
	}
	return pf
}

// PathTo returns the shortest path from root to target (reversed — root to target).
func (pf *PathFinder) PathTo(target Mat2) ([]int8, bool) {
	if _, ok := pf.Dist[target]; !ok {
		return nil, false
	}
	var path []int8
	cur := target
	for !cur.Eq(ID()) {
		parent := pf.Parent[cur]
		// Find which generator takes parent → cur.
		edge := pf.Edge[cur]
		path = append([]int8{edge}, path...)
		cur = parent
	}
	return path, true
}

// Distance returns the shortest path distance from root to target.
func (pf *PathFinder) Distance(target Mat2) int {
	return pf.Dist[target]
}

// Reachable returns the number of vertices reached by BFS.
func (pf *PathFinder) Reachable() int {
	return len(pf.Dist)
}

// MaxDist returns the maximum distance found (diameter lower bound).
func (pf *PathFinder) MaxDist() int { return pf.BFSMax }

// DistHist returns a histogram of distances.
func (pf *PathFinder) DistHist() map[int]int {
	h := make(map[int]int)
	for _, d := range pf.Dist {
		h[d]++
	}
	return h
}

// DistStats returns mean, median, max distance.
func (pf *PathFinder) DistStats() (mean float64, median float64, maxDist int) {
	dists := make([]int, 0, len(pf.Dist))
	for _, d := range pf.Dist {
		dists = append(dists, d)
		if d > maxDist {
			maxDist = d
		}
	}
	sort.Ints(dists)
	n := len(dists)
	var sum int
	for _, d := range dists {
		sum += d
	}
	mean = float64(sum) / float64(n)
	median = float64(dists[n/2])
	return
}

// Report prints BFS statistics.
func (pf *PathFinder) Report() string {
	reachable := pf.Reachable()
	maxD := pf.MaxDist()
	mean, median, _ := pf.DistStats()
	hist := pf.DistHist()
	return fmt.Sprintf("BFS: P=%d reached=%d diam≥%d mean=%.1f median=%.0f hist=%v",
		pf.P, reachable, maxD, mean, median, hist)
}

// DistStatsMean returns the mean distance.
func (pf *PathFinder) DistStatsMean() float64 {
	var sum int
	for _, d := range pf.Dist {
		sum += d
	}
	return float64(sum) / float64(len(pf.Dist))
}