package graph

// AdjEdge 带权邻接边（导出：index 包 LoadRWRGraph 需引用该类型）
type AdjEdge struct {
	To     int64
	Weight float64
}

// RWR Random-Walk-with-Restart（个性化 PageRank）：
// 从种子集合出发在无向带权图上做 power iteration，restart 概率 alpha 回到种子。
// 返回归一化到 [0,1] 的游走质量。确定性算法，无随机数。
func RWR(seedIDs []int64, adj map[int64][]AdjEdge, alpha float64) map[int64]float64 {
	if len(seedIDs) == 0 {
		return map[int64]float64{}
	}
	// 收集参与节点：种子 + 邻接可达
	nodes := make(map[int64]bool)
	for _, s := range seedIDs {
		nodes[s] = true
	}
	for from := range adj {
		nodes[from] = true
		for _, e := range adj[from] {
			nodes[e.To] = true
		}
	}
	n := len(nodes)
	if n == 0 {
		return map[int64]float64{}
	}
	idx := make(map[int64]int, n)
	ids := make([]int64, 0, n)
	for id := range nodes {
		idx[id] = len(ids)
		ids = append(ids, id)
	}

	// 重启向量：种子均匀
	r := make([]float64, n)
	for _, s := range seedIDs {
		if i, ok := idx[s]; ok {
			r[i] = 1.0 / float64(len(seedIDs))
		}
	}

	const maxIter = 50
	const eps = 1e-6
	newR := make([]float64, n)
	for iter := 0; iter < maxIter; iter++ {
		for i := range newR {
			newR[i] = 0
		}
		// 游走传播：out[i] = sum(w_ij)
		for i, id := range ids {
			var total float64
			for _, e := range adj[id] {
				total += e.Weight
			}
			if total == 0 {
				continue
			}
			for _, e := range adj[id] {
				j, ok := idx[e.To]
				if !ok {
					continue
				}
				newR[j] += (1 - alpha) * r[i] * e.Weight / total
			}
		}
		// restart
		for i := range newR {
			newR[i] += alpha * r[i]
		}
		// 收敛判断
		diff := 0.0
		for i := range newR {
			d := newR[i] - r[i]
			if d < 0 {
				d = -d
			}
			diff += d
		}
		r, newR = newR, r
		if diff < eps {
			break
		}
	}
	// 归一化 [0,1]
	maxV := 0.0
	for _, v := range r {
		if v > maxV {
			maxV = v
		}
	}
	out := make(map[int64]float64, n)
	if maxV == 0 {
		for _, id := range ids {
			out[id] = 0
		}
		return out
	}
	for i, id := range ids {
		out[id] = r[i] / maxV
	}
	return out
}