package blang.engines.internals.factories;

import org.eclipse.xtext.xbase.lib.Pair;

import bayonet.distributions.ExhaustiveDebugRandom;
import blang.engines.internals.PosteriorInferenceEngine;
import blang.inits.Arg;
import blang.inits.DefaultValue;
import blang.inits.GlobalArg;
import blang.inits.experiments.ExperimentResults;
import blang.io.BlangTidySerializer;
import blang.runtime.Runner;
import blang.runtime.SampledModel;
import blang.runtime.internals.objectgraph.GraphAnalysis;
import blang.types.ExtensionUtils;

public class Exact implements PosteriorInferenceEngine
{
  @GlobalArg ExperimentResults results;
  
  @Arg(description = "Stop and warn the user if more than this number of random draws is used in a simulation.")
         @DefaultValue("15")
  public int maxDepth = 15;
  
  @Arg(description = "Stop and warn the user if more than this number of traces is used in a simulation.")
          @DefaultValue("100_000")
  public int maxTraces = 100_000;
  
  @Arg(description = "If all generate blocks have the property that for each realization there is at most "
      + "one execution trace generating it, then we can check that the logf and randomness used in the "
      + "generate block match.")
                               @DefaultValue("false")
  public boolean checkLawsGenerateAgreement = false;
  
  @SuppressWarnings("unchecked")
  @Override
  public void performInference() 
  {
    BlangTidySerializer tidySerializer = new BlangTidySerializer(results.child(Runner.SAMPLES_FOLDER));
    ExhaustiveDebugRandom exhaustive = new ExhaustiveDebugRandom();
    int i = 0;
    while (exhaustive.hasNext())
    {
      model.forwardSample(exhaustive, true);
      
      // check this is not going to take forever
      if (exhaustive.lastDepth() > maxDepth || i > maxTraces)
        throw new RuntimeException("The number of traces or trace depth has been exceeded. "
            + "You should probably switch to an approximate inference engine. "
            + "You can also increase these limits via command line arguments if you are sure. "
            + "currentDepth=" + exhaustive.lastDepth() + " nTraces=" + i);
      
      double logWeightFromModel = model.logDensity(1.0);
      
      if (checkLawsGenerateAgreement) 
      {
        double logWeightFromGeneration = Math.log(exhaustive.lastProbability()) + model.logDensity(1.0) - model.logDensity(0.0);
        if (!ExtensionUtils.isClose(logWeightFromModel, logWeightFromGeneration))
          throw new RuntimeException("generate(rand){..} block not faithful with its laws{..} block. \n"
              + "Common mistake: forgetting to take the log of the answer in logf() { .. } constructs. \n"
              + "Diverging values: " + logWeightFromModel + " vs " + logWeightFromGeneration);
      }
      
      model.getSampleWriter(tidySerializer).write(Pair.of("sample", i++), Pair.of("logWeight", logWeightFromModel)); 
    }
  }
  
  SampledModel model;

  @Override
  public void setSampledModel(SampledModel model) 
  {
    this.model = model;
  }

  @Override
  public void check(GraphAnalysis analysis) 
  {
  }
}

@@ -131,16 +131,52 @@
 Random;%0A
+import bayonet.math.NumericalUtils;%0A
 import b
@@ -1551,16 +1551,66 @@
 LDER));%0A
+    double logNormalization = logNormalization();%0A
     Exha
@@ -2247,239 +2247,172 @@
     
-double logWeightFromModel = model.logDensity(1.0);%0A      %0A      if (checkLawsGenerateAgreement) %0A      %7B%0A        double logWeightFromGeneration = Math.log(exhaustive.lastProbability()) + model.logDensity(1.0) - model.logDensity(0.0
+if (checkLawsGenerateAgreement) %0A      %7B%0A        double logPrior = model.logDensity(0.0);%0A        double logPriorRealization = Math.log(exhaustive.lastProbability()
 );%0A 
@@ -2453,43 +2453,29 @@
 (log
-WeightFromModel, logWeightFromGener
+Prior, logPriorRealiz
 atio
@@ -2736,53 +2736,39 @@
  log
-WeightFromModel + %22 vs %22 + logWeightFromGener
+Prior + %22 vs %22 + logPriorRealiz
 atio
@@ -2876,52 +2876,871 @@
 %22log
-Weight%22, logWeightFromModel)); %0A    %7D
+Probability%22, logWeight(exhaustive) - logNormalization)); %0A    %7D%0A  %7D%0A  %0A  private double logWeight(ExhaustiveDebugRandom exhaustive)%0A  %7B%0A    double logLikelihood = model.logDensity(1.0) - model.logDensity(0.0);%0A    /*%0A     * Prior is not necessarily just model.logDensity(0.0), because the generate code %0A     * might potentially have several exec traces generating the same prior realization. %0A     */%0A    double priorRealization = Math.log(exhaustive.lastProbability()); %0A    return priorRealization + logLikelihood;%0A  %7D%0A  %0A  private double logNormalization() %0A  %7B%0A    double logSum = Double.NEGATIVE_INFINITY;%0A    ExhaustiveDebugRandom exhaustive = new ExhaustiveDebugRandom();%0A    while (exhaustive.hasNext())%0A    %7B%0A      model.forwardSample(exhaustive, true);%0A      logSum = NumericalUtils.logAdd(logSum, logWeight(exhaustive));%0A    %7D%0A    return logSum;
 %0A  %7D%0A
-  
 %0A  S

Fix bug in Exact when more than one realization maps to same generated,