<commit_before>/* Copyright (C) 2008-2010 University of Massachusetts Amherst,
   Department of Computer Science.
   This file is part of "FACTORIE" (Factor graphs, Imperative, Extensible)
   http://factorie.cs.umass.edu, http://code.google.com/p/factorie/
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License. */

package cc.factorie

import cc.factorie.la._
import scala.collection.mutable.{ArrayBuffer, HashMap, LinkedHashMap, LinkedHashSet}
import scala.collection.{Set}
import scala.collection.mutable.HashSet
import scala.collection.mutable.Queue

/** A factory object creating BPFactors and BPVariables, each of which contain methods for calculating messages. */
trait BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor
}

object BPSumProductRing extends BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1 = new BPVariable1(v)
  /** Construct and return a new BPFactor for "factor", creating/obtaining new BPVariables as necessary from "summary".
      Any of the factor neighbors present in "varying" will get a BPVariable and BPEdge; others will be considered constant.
      If "varying" is null, then any DiscreteVars are considered varying. */
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = factor match {
    case factor:Factor1[DiscreteVar] =>  
      new BPFactor1Factor1(factor, new BPEdge(summary.bpVariable(factor._1))) with BPFactorTreeSumProduct
    case factor:Factor2[DiscreteVar,DiscreteVar] => 
      if (varying == null || (varying.contains(factor._1) && varying.contains(factor._2))) new BPFactor2Factor2(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2SumProduct
      else if (varying.contains(factor._1)) new BPFactor1Factor2first(factor.asInstanceOf[Factor2[DiscreteVar,DiscreteTensorVar]], new BPEdge(summary.bpVariable(factor._1))) with BPFactorTreeSumProduct
      else new BPFactor1Factor2second(factor.asInstanceOf[Factor2[DiscreteTensorVar,DiscreteVar]], new BPEdge(summary.bpVariable(factor._2))) with BPFactorTreeSumProduct
    case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] =>
      new BPFactor2Factor3(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2SumProduct
  }
  def newBPFactorOld(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = {
    val factorVarying = factor.variables.filter(_ match {case v: DiscreteVar => (varying eq null) || varying.contains(v); case _ => false}).asInstanceOf[Seq[DiscreteVar]]
    val edges = factorVarying.map(v => new BPEdge(summary.bpVariable(v))) //_bpVariables.getOrElseUpdate(v, new BPVariable1(v, ring)))
    edges.size match {
      case 1 => factor match {
        case factor:Factor1[DiscreteVar] => new BPFactor1Factor1(factor, edges(0)) with BPFactorTreeSumProduct
        case factor:Factor2[DiscreteVar,DiscreteTensorVar] => new BPFactor1Factor2first(factor, edges(0)) with BPFactorTreeSumProduct
      }
      case 2 => factor match {
        case factor:Factor2[DiscreteVar,DiscreteVar] => new BPFactor2Factor2(factor, edges(0), edges(1)) with BPFactor2SumProduct
        case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] => new BPFactor2Factor3(factor, edges(0), edges(1)) with BPFactor2SumProduct
      }
    }
  }
}

object BPMaxProductRing extends BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1 = new BPVariable1(v)
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = factor match {
    case factor:Factor1[DiscreteVar] =>  
      new BPFactor1Factor1(factor, new BPEdge(summary.bpVariable(factor._1))) with BPFactorMaxProduct
    case factor:Factor2[DiscreteVar,DiscreteVar] => 
      if (factor._2.isInstanceOf[DiscreteVar] && null != varying && varying.contains(factor._2)) new BPFactor2Factor2(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2MaxProduct
      else new BPFactor1Factor2first(factor.asInstanceOf[Factor2[DiscreteVar,DiscreteTensorVar]], new BPEdge(summary.bpVariable(factor._1))) with BPFactorMaxProduct
    case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] =>
      new BPFactor2Factor3(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2MaxProduct
  }
  def newBPFactorOld(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = {
    val factorVarying = factor.variables.filter(_ match {case v: DiscreteVar => (varying eq null) || varying.contains(v); case _ => false}).asInstanceOf[Seq[DiscreteVar]]
    val edges = factorVarying.map(v => new BPEdge(summary.bpVariable(v))) //_bpVariables.getOrElseUpdate(v, new BPVariable1(v, ring)))
    edges.size match {
      case 1 => factor match {
        case factor:Factor1[DiscreteVar] => new BPFactor1Factor1(factor, edges(0)) with BPFactorMaxProduct
        case factor:Factor2[DiscreteVar,DiscreteTensorVar] => new BPFactor1Factor2first(factor, edges(0)) with BPFactorMaxProduct
      }
      case 2 => factor match {
        case factor:Factor2[DiscreteVar,DiscreteVar] => new BPFactor2Factor2(factor, edges(0), edges(1)) with BPFactor2MaxProduct
        case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] => new BPFactor2Factor3(factor, edges(0), edges(1)) with BPFactor2MaxProduct
      }
    }
  }
}

// TODO
// object BPSumProductBeamRing extends BPRing
// object BPSumProductNonLogRing extends BPRing // Not in log space, avoiding maths.logSum

/** A dumb container for messages factor->variable and variable->factor */
class BPEdge(val bpVariable: BPVariable1) {
  // TODO Eventually we should not require that this is a BPVariable1, but work for general BPVariable
  bpVariable.addEdge(this)
  var bpFactor: BPFactor = null
  var factorNeighborIndex: Int = -1 // TODO Is this necessary?
  // Note:  For Bethe cluster graphs with BPVariable1, these messages will be Tensor1, but for other cluster graphs they could have higher dimensionality
  var messageFromVariable: Tensor = new UniformTensor1(bpVariable.variable.domain.size, 0.0) // null // bpVariable.scores.blankCopy
  var messageFromFactor: Tensor = new UniformTensor1(bpVariable.variable.domain.size, 0.0) // null // bpVariable.scores.blankCopy
  def variable = bpVariable.variable
  def factor = bpFactor.factor
}


trait BPVariable {
  def edges: Seq[BPEdge]
  def calculateOutgoing(e:BPEdge): Tensor
  def updateOutgoing(e:BPEdge): Unit = e.messageFromVariable = calculateOutgoing(e)
  def updateOutgoing(): Unit = edges.foreach(updateOutgoing(_))
}
class BPVariable1(val variable: DiscreteVar) extends DiscreteMarginal1(variable, null) with BPVariable {
  private var _edges: List[BPEdge] = Nil
  def addEdge(e:BPEdge): Unit = _edges = e :: _edges
  final def edges: List[BPEdge] = _edges
  def calculateOutgoing(e:BPEdge): Tensor = {
    edges.size match {
      case 0 => throw new Error("BPVariable1 with no edges")
      case 1 => { require(edges.head == e); new UniformTensor1(variable.domain.size, 0.0) }
      case 2 => if (edges.head == e) edges.last.messageFromFactor else if (edges.last == e) edges.head.messageFromFactor else throw new Error
      case _ => Tensor.sum(edges.filter(_ != e).map(_.messageFromFactor)) 
    }
  }
  def calculateBelief: Tensor1 = Tensor.sum(edges.map(_.messageFromFactor)).asInstanceOf[Tensor1] // TODO Make this more efficient for cases of 1, 2, 3 edges.
  def calculateMarginal: Tensor1 = { val result = calculateBelief; result.expNormalize(); result.asInstanceOf[Tensor1] }
  override def proportions: Proportions1 = new NormalizedTensorProportions1(calculateMarginal, false)  // TODO Think about avoiding re-calc every time
  override def value1: DiscreteVar#Value = _1.domain.dimensionDomain(calculateBelief.maxIndex).asInstanceOf[DiscreteVar#Value] // TODO Ug.  This casting is rather sad.  // To avoid normalization compute time
  override def globalize(implicit d:DiffList): Unit = variable match { case v:MutableDiscreteVar[_] => v.set(calculateBelief.maxIndex)(d) }  // To avoid normalization compute time
}
// TODO class BPVariable{2,3,4} would be used for cluster graphs


trait BPFactor extends DiscreteMarginal {
  def factor: Factor
  def edges: Seq[BPEdge]
  /** Re-calculate the message from this factor to edge "e" and set e.messageFromFactor to the result. */
  def updateOutgoing(e: BPEdge): Unit
  def updateOutgoing(): Unit = edges.foreach(updateOutgoing(_))
  def scores: Tensor // All local scores across all dimensions of varying neighbors; does not use messages from variables
  /** Unnormalized log scores over values of varying neighbors */
  def calculateBeliefsTensor: Tensor
  /** The logSum of all entries in the beliefs tensor */
  def calculateLogZ: Double
  /** Normalized probabilities over values of varying neighbors */
  def calculateMarginalTensor: Tensor = calculateBeliefsTensor.expNormalized
  /** Normalized probabilities over values of only the varying neighbors, in the form of a Proportions */
  override def proportions: Proportions // Must be overridden to return "new NormalizedTensorProportions{1,2,3,4}(calculateMarginalTensor, false)"
  /** Returns a Tensor representing the marginal distribution over the values of all the neighbors of the underlying Factor. */
  def marginalTensorValues: Tensor = throw new Error("Not yet implemented")
  /** Returns a Tensor representing the marginal distribution over the statistics of all the neighbors of the underlying Factor. */
  def marginalTensorStatistics: Tensor = throw new Error("Not yet implemented")
}

trait BPFactorTreeSumProduct extends BPFactor {
  override def calculateMarginalTensor: Tensor = {
    val v = calculateBeliefsTensor
    v.expNormalize()
    v
  }
  
  def calculateLogZ: Double = calculateBeliefsTensor match {
    case t:DenseTensor => { var z = Double.NegativeInfinity; val l = t.length; var i = 0; while (i < l) { z = maths.sumLogProb(z, t(i)); i += 1 }; z }
    case t:SparseIndexedTensor => { var z = Double.NegativeInfinity; t.foreachActiveElement((i,v) => { z = maths.sumLogProb(z, v) }); z }
  }
}

trait BPFactorLoopySumProduct extends BPFactorTreeSumProduct {
  override def calculateLogZ: Double = calculateMarginalTensor match {
    case t:DenseTensor => {
      var z = 0.0
      val l = t.length
      var i = 0
      while (i < l) {
        z += t(i) * (math.log(t(i)) + scores(i))
        i += 1
      }
      z
    }
    case t:SparseIndexedTensor => {
      var z = Double.NegativeInfinity
      t.foreachActiveElement((i,v) => {
        z = v * (math.log(z) + scores(i))
      })
      z
    }
  }
}

trait BPFactorMaxProduct extends BPFactor {
  override def calculateMarginalTensor: Tensor1 = calculateBeliefsTensor.maxNormalize().asInstanceOf[Tensor1]
  override def calculateLogZ: Double = calculateBeliefsTensor match {
    case t:DenseTensor => { var z = Double.NegativeInfinity; val l = t.length; var i = 0; while (i < l) { z = math.max(z, t(i)); i += 1 }; z }
    case t:SparseIndexedTensor => { var z = Double.NegativeInfinity; t.foreachActiveElement((i,v) => { z = math.max(z, v) }); z }
  }
}

// An abstract class for BPFactors that has 1 varying neighbor.  They may have additional constant neighbors.
abstract class BPFactor1(val edge1: BPEdge) extends DiscreteMarginal1(edge1.bpVariable.variable, null) with BPFactor {
  override def scores: Tensor1
  def hasLimitedDiscreteValues1: Boolean
  def limitedDiscreteValues1: SparseBinaryTensor1
  edge1.bpFactor = this; edge1.factorNeighborIndex = 0
  val edges = Seq(edge1)
  def updateOutgoing(e: BPEdge): Unit = e match { case this.edge1 => updateOutgoing1() }
  override def updateOutgoing(): Unit = updateOutgoing1()
  def updateOutgoing1(): Unit = edge1.messageFromFactor = calculateOutgoing1
  // TODO See about caching this when possible<commit_msg>Explicitly naming the bethe entropy<commit_after>/* Copyright (C) 2008-2010 University of Massachusetts Amherst,
   Department of Computer Science.
   This file is part of "FACTORIE" (Factor graphs, Imperative, Extensible)
   http://factorie.cs.umass.edu, http://code.google.com/p/factorie/
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License. */

package cc.factorie

import cc.factorie.la._
import scala.collection.mutable.{ArrayBuffer, HashMap, LinkedHashMap, LinkedHashSet}
import scala.collection.{Set}
import scala.collection.mutable.HashSet
import scala.collection.mutable.Queue

/** A factory object creating BPFactors and BPVariables, each of which contain methods for calculating messages. */
trait BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor
}

object BPSumProductRing extends BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1 = new BPVariable1(v)
  /** Construct and return a new BPFactor for "factor", creating/obtaining new BPVariables as necessary from "summary".
      Any of the factor neighbors present in "varying" will get a BPVariable and BPEdge; others will be considered constant.
      If "varying" is null, then any DiscreteVars are considered varying. */
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = factor match {
    case factor:Factor1[DiscreteVar] =>  
      new BPFactor1Factor1(factor, new BPEdge(summary.bpVariable(factor._1))) with BPFactorTreeSumProduct
    case factor:Factor2[DiscreteVar,DiscreteVar] => 
      if (varying == null || (varying.contains(factor._1) && varying.contains(factor._2))) new BPFactor2Factor2(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2SumProduct
      else if (varying.contains(factor._1)) new BPFactor1Factor2first(factor.asInstanceOf[Factor2[DiscreteVar,DiscreteTensorVar]], new BPEdge(summary.bpVariable(factor._1))) with BPFactorTreeSumProduct
      else new BPFactor1Factor2second(factor.asInstanceOf[Factor2[DiscreteTensorVar,DiscreteVar]], new BPEdge(summary.bpVariable(factor._2))) with BPFactorTreeSumProduct
    case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] =>
      new BPFactor2Factor3(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2SumProduct
  }
  def newBPFactorOld(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = {
    val factorVarying = factor.variables.filter(_ match {case v: DiscreteVar => (varying eq null) || varying.contains(v); case _ => false}).asInstanceOf[Seq[DiscreteVar]]
    val edges = factorVarying.map(v => new BPEdge(summary.bpVariable(v))) //_bpVariables.getOrElseUpdate(v, new BPVariable1(v, ring)))
    edges.size match {
      case 1 => factor match {
        case factor:Factor1[DiscreteVar] => new BPFactor1Factor1(factor, edges(0)) with BPFactorTreeSumProduct
        case factor:Factor2[DiscreteVar,DiscreteTensorVar] => new BPFactor1Factor2first(factor, edges(0)) with BPFactorTreeSumProduct
      }
      case 2 => factor match {
        case factor:Factor2[DiscreteVar,DiscreteVar] => new BPFactor2Factor2(factor, edges(0), edges(1)) with BPFactor2SumProduct
        case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] => new BPFactor2Factor3(factor, edges(0), edges(1)) with BPFactor2SumProduct
      }
    }
  }
}

object BPMaxProductRing extends BPRing {
  def newBPVariable(v:DiscreteVar): BPVariable1 = new BPVariable1(v)
  def newBPFactor(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = factor match {
    case factor:Factor1[DiscreteVar] =>  
      new BPFactor1Factor1(factor, new BPEdge(summary.bpVariable(factor._1))) with BPFactorMaxProduct
    case factor:Factor2[DiscreteVar,DiscreteVar] => 
      if (factor._2.isInstanceOf[DiscreteVar] && null != varying && varying.contains(factor._2)) new BPFactor2Factor2(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2MaxProduct
      else new BPFactor1Factor2first(factor.asInstanceOf[Factor2[DiscreteVar,DiscreteTensorVar]], new BPEdge(summary.bpVariable(factor._1))) with BPFactorMaxProduct
    case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] =>
      new BPFactor2Factor3(factor, new BPEdge(summary.bpVariable(factor._1)), new BPEdge(summary.bpVariable(factor._2))) with BPFactor2MaxProduct
  }
  def newBPFactorOld(factor:Factor, varying:Set[DiscreteVar], summary:BPSummary): BPFactor = {
    val factorVarying = factor.variables.filter(_ match {case v: DiscreteVar => (varying eq null) || varying.contains(v); case _ => false}).asInstanceOf[Seq[DiscreteVar]]
    val edges = factorVarying.map(v => new BPEdge(summary.bpVariable(v))) //_bpVariables.getOrElseUpdate(v, new BPVariable1(v, ring)))
    edges.size match {
      case 1 => factor match {
        case factor:Factor1[DiscreteVar] => new BPFactor1Factor1(factor, edges(0)) with BPFactorMaxProduct
        case factor:Factor2[DiscreteVar,DiscreteTensorVar] => new BPFactor1Factor2first(factor, edges(0)) with BPFactorMaxProduct
      }
      case 2 => factor match {
        case factor:Factor2[DiscreteVar,DiscreteVar] => new BPFactor2Factor2(factor, edges(0), edges(1)) with BPFactor2MaxProduct
        case factor:Factor3[DiscreteVar,DiscreteVar,DiscreteTensorVar] => new BPFactor2Factor3(factor, edges(0), edges(1)) with BPFactor2MaxProduct
      }
    }
  }
}

// TODO
// object BPSumProductBeamRing extends BPRing
// object BPSumProductNonLogRing extends BPRing // Not in log space, avoiding maths.logSum

/** A dumb container for messages factor->variable and variable->factor */
class BPEdge(val bpVariable: BPVariable1) {
  // TODO Eventually we should not require that this is a BPVariable1, but work for general BPVariable
  bpVariable.addEdge(this)
  var bpFactor: BPFactor = null
  var factorNeighborIndex: Int = -1 // TODO Is this necessary?
  // Note:  For Bethe cluster graphs with BPVariable1, these messages will be Tensor1, but for other cluster graphs they could have higher dimensionality
  var messageFromVariable: Tensor = new UniformTensor1(bpVariable.variable.domain.size, 0.0) // null // bpVariable.scores.blankCopy
  var messageFromFactor: Tensor = new UniformTensor1(bpVariable.variable.domain.size, 0.0) // null // bpVariable.scores.blankCopy
  def variable = bpVariable.variable
  def factor = bpFactor.factor
}


trait BPVariable {
  def edges: Seq[BPEdge]
  def calculateOutgoing(e:BPEdge): Tensor
  def updateOutgoing(e:BPEdge): Unit = e.messageFromVariable = calculateOutgoing(e)
  def updateOutgoing(): Unit = edges.foreach(updateOutgoing(_))
}
class BPVariable1(val variable: DiscreteVar) extends DiscreteMarginal1(variable, null) with BPVariable {
  private var _edges: List[BPEdge] = Nil
  def addEdge(e:BPEdge): Unit = _edges = e :: _edges
  final def edges: List[BPEdge] = _edges
  def calculateOutgoing(e:BPEdge): Tensor = {
    edges.size match {
      case 0 => throw new Error("BPVariable1 with no edges")
      case 1 => { require(edges.head == e); new UniformTensor1(variable.domain.size, 0.0) }
      case 2 => if (edges.head == e) edges.last.messageFromFactor else if (edges.last == e) edges.head.messageFromFactor else throw new Error
      case _ => Tensor.sum(edges.filter(_ != e).map(_.messageFromFactor)) 
    }
  }
  def calculateBelief: Tensor1 = Tensor.sum(edges.map(_.messageFromFactor)).asInstanceOf[Tensor1] // TODO Make this more efficient for cases of 1, 2, 3 edges.
  def calculateMarginal: Tensor1 = { val result = calculateBelief; result.expNormalize(); result.asInstanceOf[Tensor1] }
  override def proportions: Proportions1 = new NormalizedTensorProportions1(calculateMarginal, false)  // TODO Think about avoiding re-calc every time
  override def value1: DiscreteVar#Value = _1.domain.dimensionDomain(calculateBelief.maxIndex).asInstanceOf[DiscreteVar#Value] // TODO Ug.  This casting is rather sad.  // To avoid normalization compute time
  override def globalize(implicit d:DiffList): Unit = variable match { case v:MutableDiscreteVar[_] => v.set(calculateBelief.maxIndex)(d) }  // To avoid normalization compute time
}
// TODO class BPVariable{2,3,4} would be used for cluster graphs


trait BPFactor extends DiscreteMarginal {
  def factor: Factor
  def edges: Seq[BPEdge]
  /** Re-calculate the message from this factor to edge "e" and set e.messageFromFactor to the result. */
  def updateOutgoing(e: BPEdge): Unit
  def updateOutgoing(): Unit = edges.foreach(updateOutgoing(_))
  def scores: Tensor // All local scores across all dimensions of varying neighbors; does not use messages from variables
  /** Unnormalized log scores over values of varying neighbors */
  def calculateBeliefsTensor: Tensor
  /** The logSum of all entries in the beliefs tensor */
  def calculateLogZ: Double
  /** Normalized probabilities over values of varying neighbors */
  def calculateMarginalTensor: Tensor = calculateBeliefsTensor.expNormalized
  /** Normalized probabilities over values of only the varying neighbors, in the form of a Proportions */
  override def proportions: Proportions // Must be overridden to return "new NormalizedTensorProportions{1,2,3,4}(calculateMarginalTensor, false)"
  override def betheObjective: Double = calculateMarginalTensor match {
    case t:DenseTensor => {
      var z = 0.0
      val l = t.length
      var i = 0
      while (i < l) {
        z += t(i) * (math.log(t(i)) + scores(i))
        i += 1
      }
      z
    }
    case t:SparseIndexedTensor => {
      var z = Double.NegativeInfinity
      t.foreachActiveElement((i,v) => {
        z = v * (math.log(z) + scores(i))
      })
      z
    }
  }
  /** Returns a Tensor representing the marginal distribution over the values of all the neighbors of the underlying Factor. */
  def marginalTensorValues: Tensor = throw new Error("Not yet implemented")
  /** Returns a Tensor representing the marginal distribution over the statistics of all the neighbors of the underlying Factor. */
  def marginalTensorStatistics: Tensor = throw new Error("Not yet implemented")
}

trait BPFactorTreeSumProduct extends BPFactor {
  override def calculateMarginalTensor: Tensor = {
    val v = calculateBeliefsTensor
    v.expNormalize()
    v
  }
  
  def calculateLogZ: Double = calculateBeliefsTensor match {
    case t:DenseTensor => { var z = Double.NegativeInfinity; val l = t.length; var i = 0; while (i < l) { z = maths.sumLogProb(z, t(i)); i += 1 }; z }
    case t:SparseIndexedTensor => { var z = Double.NegativeInfinity; t.foreachActiveElement((i,v) => { z = maths.sumLogProb(z, v) }); z }
  }
}

trait BPFactorMaxProduct extends BPFactor {
  override def calculateMarginalTensor: Tensor1 = calculateBeliefsTensor.maxNormalize().asInstanceOf[Tensor1]
  override def calculateLogZ: Double = calculateBeliefsTensor match {
    case t:DenseTensor => { var z = Double.NegativeInfinity; val l = t.length; var i = 0; while (i < l) { z = math.max(z, t(i)); i += 1 }; z }
    case t:SparseIndexedTensor => { var z = Double.NegativeInfinity; t.foreachActiveElement((i,v) => { z = math.max(z, v) }); z }
  }
}

// An abstract class for BPFactors that has 1 varying neighbor.  They may have additional constant neighbors.
abstract class BPFactor1(val edge1: BPEdge) extends DiscreteMarginal1(edge1.bpVariable.variable, null) with BPFactor {
  override def scores: Tensor1
  def hasLimitedDiscreteValues1: Boolean
  def limitedDiscreteValues1: SparseBinaryTensor1
  edge1.bpFactor = this; edge1.factorNeighborIndex = 0
  val edges = Seq(edge1)
  def updateOutgoing(e: BPEdge): Unit = e match { case this.edge1 => updateOutgoing1() }
  override def updateOutgoing(): Unit = updateOutgoing1()
  def updateOutgoing1(): Unit = edge1.messageFromFactor = calculateOutgoing1
  // TODO See about caching this when possible
{'fraction_non_alphanumeric': 0.09858758198865236, 'fraction_numerical': 0.014405858919158182, 'mean_word_length': 4.684354986276304, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 85, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2346759', 'n_tokens_mistral': 7626, 'n_tokens_neox': 7475, 'n_words': 2135}