utils.go

package spn

import (
	"fmt"
	"math"

	"github.com/RenatoGeh/gospn/common"
	"github.com/RenatoGeh/gospn/sys"
	"github.com/RenatoGeh/gospn/utils"
)

// Some of the following functions are non-recursive versions of equivalent spn.SPN methods. They
// are done using a Queue or Stack to perform the graph search instead of the recursion call stack.
// When the SPN is dense, running the recursive versions can take exponential time (as we do not
// account for already visited vertices). In these static function versions, all searches are done
// in time linear to the graphs. For this reason, unless the SPN is a tree (or the graph sparse
// enough), the preferred method is using the static function version. When the SPN is a tree, the
// best method is the recursive version, as it takes less memory and same time usage in average
// when compared to the static versions.

// InferenceY returns the value of S(I, Y=y). This convenience function allows for fast computation
// of soft inference values without having to create another VarSet for each valuation of Y.
func InferenceY(S SPN, I VarSet, Y, y int) float64 {
	if len(S.Ch()) == 0 {
		return 0
	}
	J := map[int]int{Y: y}
	O := common.Queue{}
	TopSortTarjan(S, &O)
	V := make(map[SPN]float64)
	for !O.Empty() {
		s := O.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			if varid := s.Sc()[0]; varid == Y {
				V[s] = s.Value(J)
			} else {
				V[s] = s.Value(I)
			}
		case "sum":
			sum := s.(*Sum)
			ch := sum.Ch()
			W := sum.Weights()
			n := len(ch)
			vals := make([]float64, n)
			for i, cs := range ch {
				v := V[cs]
				vals[i] = v + math.Log(W[i])
			}
			V[s] = sum.Compute(vals)
		case "product":
			prod := s.(*Product)
			ch := prod.Ch()
			n := len(ch)
			vals := make([]float64, n)
			for i, ch := range ch {
				vals[i] = V[ch]
			}
			V[s] = prod.Compute(vals)
		}
	}
	return V[S]
}

// Inference simply returns the value of S(I), without storing values for later use.
func Inference(S SPN, I VarSet) float64 {
	if len(S.Ch()) == 0 {
		return 0
	}
	O := common.Queue{}
	TopSortTarjan(S, &O)
	V := make(map[SPN]float64)
	for !O.Empty() {
		s := O.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			V[s] = s.Value(I)
		case "sum":
			sum := s.(*Sum)
			ch := sum.Ch()
			W := sum.Weights()
			n := len(ch)
			vals := make([]float64, n)
			for i, cs := range ch {
				v := V[cs]
				vals[i] = v + math.Log(W[i])
			}
			V[s] = sum.Compute(vals)
		case "product":
			prod := s.(*Product)
			ch := prod.Ch()
			n := len(ch)
			vals := make([]float64, n)
			for i, ch := range ch {
				vals[i] = V[ch]
			}
			V[s] = prod.Compute(vals)
		}
	}
	return V[S]
}

// StoreInference takes an SPN S and stores the values for an instance I on a DP table storage
// at the position designated by the ticket tk. Returns S and the ticket used (if tk < 0,
// StoreInference creates a new ticket).
func StoreInference(S SPN, I VarSet, tk int, storage *Storer) (SPN, int) {
	if len(S.Ch()) == 0 {
		return nil, -1
	}

	if tk < 0 {
		tk = storage.NewTicket()
	}

	// Get topological order.
	O := common.Queue{}
	TopSortTarjan(S, &O)
	sys.Free()

	table, _ := storage.Table(tk)
	for !O.Empty() {
		s := O.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			table.StoreSingle(s, s.Value(I))
		case "sum":
			sum := s.(*Sum)
			ch := sum.Ch()
			W := sum.Weights()
			n := len(ch)
			vals := make([]float64, n)
			for i, cs := range ch {
				v, e := table.Single(cs)
				if !e {
					fmt.Println("Error: The SPN graph has sums or products as leaves.")
					sys.Free()
					return nil, -1
				}
				vals[i] = v + math.Log(W[i])
			}
			if n == 0 {
				table.StoreSingle(s, utils.LogZero)
			} else {
				table.StoreSingle(s, sum.Compute(vals))
			}
		case "product":
			prod := s.(*Product)
			ch := prod.Ch()
			n := len(ch)
			vals := make([]float64, n)
			for i, cs := range ch {
				vals[i], _ = table.Single(cs)
			}
			table.StoreSingle(s, prod.Compute(vals))
		}
	}
	sys.Free()
	return S, tk
}

// StoreMAP takes an SPN S and stores the MAP values for an instance I on a DP table storage
// at the position designated by the ticket tk. Returns S and the ticket used (if tk < 0,
// StoreMAP creates a new ticket).
func StoreMAP(S SPN, I VarSet, tk int, storage *Storer) (SPN, int, VarSet) {
	if len(S.Ch()) == 0 {
		return nil, -1, nil
	}

	if tk < 0 {
		tk = storage.NewTicket()
	}

	// Get topological order.
	T := common.Queue{}
	TopSortTarjan(S, &T)

	tab, _ := storage.Table(tk)
	// Find max values.
	for !T.Empty() {
		s := T.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			m := s.Max(I)
			tab.StoreSingle(s, m)
		case "sum":
			sum := s.(*Sum)
			W := sum.Weights()
			ch := s.Ch()
			mv := math.Inf(-1)
			for i, c := range ch {
				v, _ := tab.Single(c)
				u := math.Log(W[i]) + v
				if u > mv {
					mv = u
				}
			}
			tab.StoreSingle(s, mv)
		case "product":
			ch := s.Ch()
			var v float64
			for _, c := range ch {
				cv, _ := tab.Single(c)
				v += cv
			}
			tab.StoreSingle(s, v)
		}
	}

	Q := common.Queue{}
	V := make(map[SPN]bool)
	Q.Enqueue(S)
	V[S] = true
	M := make(VarSet)

	// Find MAP states.
	for !Q.Empty() {
		s := Q.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			N, _ := s.ArgMax(I)
			for k, v := range N {
				M[k] = v
			}
		case "sum":
			sum := s.(*Sum)
			W := sum.Weights()
			ch := s.Ch()
			m := math.Inf(-1)
			var mv []SPN
			for i, c := range ch {
				v, _ := tab.Single(c)
				u := math.Log(W[i]) + v
				if u > m {
					mv, m = []SPN{c}, u
				} else if u == m {
					mv = append(mv, c)
				}
			}
			// Randomly break ties.
			mvc := mv[sys.RandIntn(len(mv))]
			if mvc != nil && !V[mvc] {
				Q.Enqueue(mvc)
				V[mvc] = true
			}
		case "product":
			ch := s.Ch()
			for _, c := range ch {
				if !V[c] {
					Q.Enqueue(c)
					V[c] = true
				}
			}
		}
	}

	return S, tk, M
}

func norm(v []float64) {
	var min float64
	for _, u := range v {
		if u < min {
			min = u
		}
	}
	if min < 0 {
		for i := range v {
			v[i] += min
		}
	}

	var norm float64
	for i := range v {
		norm += v[i]
	}
	for i := range v {
		v[i] /= norm
	}
}

// NormalizeSPN recursively normalizes the SPN S.
func NormalizeSPN(S SPN) SPN {
	Q := common.Queue{}
	V := make(map[SPN]bool)

	Q.Enqueue(S)
	V[S] = true

	for !Q.Empty() {
		s := Q.Dequeue().(SPN)
		if s.Type() == "sum" {
			z := s.(*Sum)
			W := z.Weights()
			norm(W)
		}
		ch := s.Ch()
		for _, c := range ch {
			if !V[c] {
				Q.Enqueue(c)
				V[c] = true
			}
		}
	}

	return S
}

// ComputeHeight computes the height of a certain SPN S.
func ComputeHeight(S SPN) int {
	T := common.Stack{}
	V := make(map[SPN]int)

	T.Push(S)
	V[S] = 0

	var h int
	for !T.Empty() {
		s := T.Pop().(SPN)
		if s.Type() == "leaf" && V[s] > h {
			h = V[s]
		}
		ch := s.Ch()
		for _, c := range ch {
			if _, e := V[c]; !e {
				T.Push(c)
				V[c] = V[s] + 1
			}
		}
	}

	return h
}

// ComputeScope computes the scope of a certain SPN S.
func ComputeScope(S SPN) []int {
	T := common.Queue{}
	TopSortTarjan(S, &T)

	for !T.Empty() {
		s := T.Dequeue().(SPN)
		ch := s.Ch()
		if len(ch) > 0 {
			msc := make(map[int]bool)
			for _, c := range ch {
				csc := c.rawSc()
				for _, v := range csc {
					msc[v] = true
				}
			}
			var sc []int
			for k, _ := range msc {
				sc = append(sc, k)
			}
			s.setRawSc(sc)
		}
	}

	return S.rawSc()
}

// Complete returns whether the SPN is complete.
func Complete(S SPN) bool {
	ComputeScope(S)
	Q := common.Queue{}
	V := make(map[SPN]bool)

	Q.Enqueue(S)
	V[S] = true

	for !Q.Empty() {
		s := Q.Dequeue().(SPN)
		ch := s.Ch()
		if s.Type() == "sum" {
			sc := s.rawSc()
			v := make(map[int]int)
			for _, u := range sc {
				v[u]++
			}
			for _, c := range ch {
				csc := c.rawSc()
				// Invariant: ComputeScope guarantees that there will be no duplicates.
				if len(csc) != len(sc) {
					sys.Printf("len(csc)=%d != len(sc)=%d\n", len(csc), len(sc))
					sys.Printf("%v\n%v\n", csc, sc)
					return false
				}
				for _, u := range csc {
					_, e := v[u]
					if !e {
						sys.Printf("v[%d] does not exist\n", u)
						return false
					}
					v[u]++
				}
			}
			k := len(ch) + 1
			for _, u := range v {
				if u != k {
					sys.Printf("u=%d != k=%d\n", u, k)
					return false
				}
			}
		}
		for _, c := range ch {
			if !V[S] && c.Type() != "leaf" {
				Q.Enqueue(c)
				V[c] = true
			}
		}
	}

	return true
}

// Decomposable returns whether the SPN is decomposable.
func Decomposable(S SPN) bool {
	ComputeScope(S)
	Q := common.Queue{}
	V := make(map[SPN]bool)

	Q.Enqueue(S)
	V[S] = true

	for !Q.Empty() {
		s := Q.Dequeue().(SPN)
		ch := s.Ch()
		if s.Type() == "product" {
			sc := s.rawSc()
			v := make(map[int]int)
			for _, u := range sc {
				v[u]++
			}
			n := len(sc)
			for _, c := range ch {
				csc := c.rawSc()
				// Invariant: ComputeScope guarantees that there will be no duplicates.
				for _, u := range csc {
					_, e := v[u]
					if !e {
						return false
					}
					v[u]++
				}
				n -= len(csc)
			}
			for _, u := range v {
				if u != 2 {
					return false
				}
			}
			if n != 0 {
				return false
			}
		}
		for _, c := range ch {
			if !V[S] && c.Type() != "leaf" {
				Q.Enqueue(c)
				V[c] = true
			}
		}
	}

	return true
}

// TraceMAP returns the max child index of each sum node in a map. We assume decomposability and
// completeness. When this condition is not met, one arbitrary MAP state is chosen. When the SPN is
// both decomposable and complete, it is easy to see that the induced MAP trace of the SPN's graph
// is a tree, and thus no two paths from the root to a leaf intersect. For the negative case, there
// can be two paths that do intersect, and thus we could have randomly chosen different max
// children in case of ties. In this case, TraceMAP chooses the first child it finds to meet the
// criteria.
func TraceMAP(S SPN, I VarSet) map[SPN]int {
	tab := make(map[SPN]float64)
	P := make(map[SPN]int)

	// Get topological order.
	T := common.Queue{}
	TopSortTarjan(S, &T)

	// Find max values.
	for !T.Empty() {
		s := T.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "leaf":
			m := s.Max(I)
			tab[s] = m
		case "sum":
			sum := s.(*Sum)
			W := sum.Weights()
			ch := s.Ch()
			mv := math.Inf(-1)
			for i, c := range ch {
				v := tab[c]
				u := math.Log(W[i]) + v
				if u > mv {
					mv = u
				}
			}
			tab[s] = mv
		case "product":
			ch := s.Ch()
			var v float64
			for _, c := range ch {
				cv := tab[c]
				v += cv
			}
			tab[s] = v
		}
	}

	Q := common.Queue{}
	V := make(map[SPN]bool)
	Q.Enqueue(S)
	V[S] = true

	for !Q.Empty() {
		s := Q.Dequeue().(SPN)
		switch t := s.Type(); t {
		case "sum":
			sum := s.(*Sum)
			W := sum.Weights()
			ch := s.Ch()
			m := math.Inf(-1)
			var mi []int
			for i, c := range ch {
				v := tab[c]
				u := math.Log(W[i]) + v
				if u > m {
					mi, m = []int{i}, u
				} else if u == m {
					mi = append(mi, i)
				}
			}
			// Randomly break ties.
			i := sys.RandIntn(len(mi))
			mvc := ch[mi[i]]
			if mvc != nil && !V[mvc] {
				Q.Enqueue(mvc)
				V[mvc] = true
				P[s] = mi[i]
			}
		case "product":
			ch := s.Ch()
			for _, c := range ch {
				if !V[c] {
					Q.Enqueue(c)
					V[c] = true
				}
			}
		}
	}

	return P
}