package edu.stanford.nlp.sempre;

import com.google.common.base.Joiner;
import fig.basic.*;
import fig.exec.Execution;
import fig.prob.SampleUtils;

import java.util.*;

/**
 * @author joberant
 * Parser and learning using re-inforcement learning, that is maximizing future rewards.
 * During learning we sample from the agenda that contains all possible next actions
 * Then we update the gradient so that actions that were taken and led to good reward
 * will get better weights than other actions
 *
 * At test time we simply find the highest probability parse so that is done using regular parsing
 *
 * We implement this using a chart and an agenda that holds the possible future actions
 *
 * NOTE: during search we search with the agenda and then do re-ranking
 * features are separated by a "search" prefix.
 * For this when we compute a derivation score on the agenda we prefix the derivation
 * features with "search", and in addition, when updating the expected counts we add the
 * "search" prefix. This is brittle... but at the end we score with re-ranking features
 * and then the score of the derivation is \theta \times \phi as usual
 */
public class ReinforcementParser extends Parser {
  public static class Options {
    @Option(gloss = "Whether to do coarse pruning")
    public boolean efficientCoarsePrune = true;
    @Option(gloss = "Whether to do importance sampling")
    public double multiplicativeBonus = 1000d;
    @Option (gloss = "Number of samples")
    public int numOfSamplesPerExample = 1;
    @Option (gloss = "Whether to update gradient only for correct moves")
    public boolean updateGradientForCorrectMovesOnly = true;
    @Option (gloss = "Low probability for which we don't unroll the stream")
    public double lowProb = 0.01;
    @Option (gloss = "Whether to simulate the log liklihood objective")
    public boolean simulateNonRlObjective = false;
    @Option (gloss = "Whether to always unroll (even at test time)")
    public boolean alwaysUnroll = false;
  }
  public static Options opts = new Options();

  // we assume here a binarized grammar
  Map<String, Map<String, List<Rule>>> leftToRightSiblingMap = new HashMap<>();
  Map<String, Map<String, List<Rule>>> rightToLeftSiblingMap = new HashMap<>();
  Map<String, List<Rule>> terminalsToRulesList = new HashMap<>();
  final CoarseParser coarseParser;
  public static final String SEARCH_PREFIX = "search_";
  public final String searchPrefix;

  public ReinforcementParser(Spec spec) {
    super(spec);
    coarseParser = new CoarseParser(grammar);

    // generate maps from left to right and vice versa
    for (Rule rule : grammar.rules) {
      if (rule.rhs.size() > 2)
        throw new RuntimeException("We assume that the grammar is binarized, rule: "  + rule);
      if (rule.rhs.size() == 2) {
        String left = rule.rhs.get(0);
        String right = rule.rhs.get(1);
        addToSiblingMap(left, right, rule, leftToRightSiblingMap);
        addToSiblingMap(right, left, rule, rightToLeftSiblingMap);
      }
      if (rule.isRhsTerminals())
        MapUtils.addToList(terminalsToRulesList, Joiner.on(' ').join(rule.rhs), rule);
    }
    if (Parser.opts.visualizeChartFilling)
      this.chartFillOut = IOUtils.openOutAppendEasy(Execution.getFile("chartfill"));
    searchPrefix = opts.simulateNonRlObjective ? "" : SEARCH_PREFIX;
    LogInfo.logs("ReinforcementParser(): search prefix is %s", searchPrefix);
  }

  private void addToSiblingMap(String keySibling, String valueSibling, Rule rule,
                               Map<String, Map<String, List<Rule>>> siblingToSiblingMap) {
    Map<String, List<Rule>> valueSiblingMap = siblingToSiblingMap.get(keySibling);
    if (valueSiblingMap == null)
      siblingToSiblingMap.put(keySibling, valueSiblingMap = new HashMap<>());
    MapUtils.addToList(valueSiblingMap, valueSibling, rule);
  }

  @Override
  public ParserState newParserState(Params params, Example ex, boolean computeExpectedCounts) {
    if (computeExpectedCounts) { // if we learn - use sampling, otherwise, use max
      // if we simulate non RL we just take the max and not do sampling
      if (opts.simulateNonRlObjective) {
        return new ReinforcementParserState.StateBuilder()
                .parser(this)
                .params(params)
                .example(ex)
                .samplingStrategy("max")
                .computeExpectedCounts(true)
                .createState();
      } else {
        return (new ReinforcementParserState.StateBuilder()
                .parser(this)
                .params(params)
                .example(ex))
                .samplingStrategy("proposal")
                .computeExpectedCounts(true)
                .createState();
      }
    }
    return (new ReinforcementParserState.StateBuilder()
            .parser(this)
            .params(params)
            .example(ex))
            .samplingStrategy("max")
            .computeExpectedCounts(false)
            .createState();
  }

  @Override
  public Params getSearchParams(Params params) {
    return params.copyParamsByPrefix(searchPrefix);
  }
}

// Parsing by sampling items from the agenda
// Defined by: (a) sampling strategy (from agenda, from true derivation, take the max)
// (b) whether an update needs to be done
final class ReinforcementParserState extends AbstractReinforcementParserState {

  public enum NecessaryDeriv { NECESSARY_DERIV, UNNECESSARY_DERIV, UNKNOWN }
  private static final double LOG_SMALL_PROB = Math.log(ReinforcementParser.opts.lowProb);

  private final ParserAgenda<PrioritizedDerivationStream> agenda;

  private int completeDerivationsPushed = 0;
  private int firstCorrectItem = -1;
  private String samplingStrategy;
  private Sampler sampler;
  List<Derivation> correctDerivations = new ArrayList<>();
  private Map<String, Double> stateSequenceExpectedCounts = new HashMap<>();
  Random randGen = new Random(1);
  // backpointers for remembering what derivations on the stream were popped before others
  private Map<Long, Pair<ArrayList<Derivation>, Integer>> backpointerList;
  private int numItemsSampled = 0;

  public static class StateBuilder {
    private ReinforcementParser parser;
    private Params params;
    private Example example;
    private ParserState coarseState;
    private String samplingStrategy = null;
    private boolean computeExpectedCounts;

    public StateBuilder parser(ReinforcementParser parser) { this.parser = parser; return this; }
    public StateBuilder params(Params params) { this.params = params; return this; }
    public StateBuilder example(Example example) { this.example = example; return this; }
    public StateBuilder samplingStrategy(String samplingStrategy) { this.samplingStrategy = samplingStrategy; return this; }
    public StateBuilder computeExpectedCounts(boolean computeExpectedCounts) { this.computeExpectedCounts = computeExpectedCounts; return this; }
    public ReinforcementParserState createState() {
      return new ReinforcementParserState(this.parser, this.params, this.example, this.computeExpectedCounts,
              this.samplingStrategy);
    }
  }

  // note that the sampler has a pointer to the fields of the state where they were created
  private ReinforcementParserState(ReinforcementParser parser, Params params, Example ex, boolean computeExpectedCounts,
                                   String samplingStrategy) {
    super(parser, params, ex, computeExpectedCounts);
    this.samplingStrategy = samplingStrategy;
    backpointerList = new HashMap<>();
    agenda = samplingStrategy.equals("max") ? new QueueParserAgenda() : new ListParserAgenda();
  }

  private void clearState() {
    agenda.clear();
    clearChart();
    completeDerivationsPushed = 0;
    firstCorrectItem = -1;
    correctDerivations.clear();
    stateSequenceExpectedCounts.clear();
    backpointerList.clear();
    numItemsSampled = 0;
  }



  protected void addToAgenda(DerivationStream derivationStream) {
    addToAgenda(derivationStream, 0d);
  }

  private void addToAgenda(DerivationStream derivationStream, double probSum) {

    if (!derivationStream.hasNext()) return;

    //if it's less than one we can just add it even if we unroll everything (optimization)
    if(!ReinforcementParser.opts.alwaysUnroll || derivationStream.estimatedSize() <= 1) {
      Derivation deriv = derivationStream.peek();
      featurizeAndScoreDerivation(deriv);
      addToAgendaWithScore(derivationStream, deriv.score, probSum);
      if (completeDerivationsPushed % 100 == 0) // sort the agenda
        agenda.sort();
    }
    else {
      while(derivationStream.hasNext()) {
        Derivation deriv = derivationStream.next();
        featurizeAndScoreDerivation(deriv);
        DerivationStream newStream = SingleDerivationStream.constant(deriv);
        addToAgendaWithScore(newStream, deriv.score, probSum);
        if (completeDerivationsPushed % 100 == 0) // sort the agenda
          agenda.sort();
      }
    }
  }

  // we need to override the method because parameters are prefixed with "search_"
  // this means that the score will not be the dot product and features and weights
  protected void featurizeAndScoreDerivation(Derivation deriv) {
    if (deriv.isFeaturizedAndScored()) return;

    // Compute features
    parser.extractor.extractLocal(ex, deriv);

    // Compute score by adding |SEARCH_PREFIX| prefix and adding children scores
    FeatureVector searchFV = deriv.addPrefixLocalFeatureVector(parser.searchPrefix);
    deriv.score = searchFV.dotProduct(params);
    if (deriv.children != null)
      for (Derivation child : deriv.children)
        deriv.score += child.score;

    if (parser.verbose(3))
      LogInfo.logs(
              "featurizeAndScore(score=%s) %s %s: %s [rule: %s]",
              Fmt.D(deriv.score), deriv.cat, ex.spanString(deriv.start, deriv.end), deriv, deriv.rule);
    numOfFeaturizedDerivs++;
  }

  private void addToAgendaWithScore(DerivationStream derivationStream, double derivScore, double probSum) {
    if (derivScore == Double.NEGATIVE_INFINITY) return; // no need to add bad derivations - shouldn't happen

    Derivation deriv = derivationStream.peek(); // Score a DerivationStream based on the first item in the stream.
    double priority = derivScore - (completeDerivationsPushed++ * EPSILON);
    agenda.add(new PrioritizedDerivationStream(derivationStream, priority, probSum), priority); // when adding to agenda probsum is 0

    if (parser.verbose(3)) {
      LogInfo.logs("ReinforcementParser: adding to agenda: size=%s, priority=%s, deriv=%s(%s,%s), formula=%s,|pushed|=%s",
              agenda.size(), priority, deriv.cat, deriv.start, deriv.end, deriv.formula, completeDerivationsPushed);
    }
  }

  public boolean continueParsing() {
    if (agenda.size() == 0) {
      LogInfo.log("Agenda is empty");
      return false;
    }

    return chart[0][numTokens].get(Rule.rootCat) == null ||
            chart[0][numTokens].get(Rule.rootCat).size() < getBeamSize();
  }

  public void infer() {
    if (numTokens == 0)
      return;

    ReinforcementParserState oracleState = null;
    expectedCounts = new HashMap<>();
    if (computeExpectedCounts && !ReinforcementParser.opts.simulateNonRlObjective) { // when updating params we first find a correct derivation to set the oracle sampler
      // TODO(jonathan): move to ReinforcementParser, not ParserState
      LogInfo.begin_track("Finding oracle derivation");
      oracleState = new StateBuilder()
              .parser(this.parser)
              .params(this.params)
              .example(this.ex)
              .samplingStrategy("agenda")
              .computeExpectedCounts(false).createState(); // update params is false preventing an infinite loop
      oracleState.infer();
      LogInfo.end_track();
      if (oracleState.correctDerivations.isEmpty()) {
        LogInfo.logs("No oracle derivation found");
        return;
      }
    }
    createSampler(oracleState); // we can only create the sampler after we have the oracle derivation

    LogInfo.begin_track("Coarse parsing");
    coarseParserState = null;
    if (ReinforcementParser.opts.efficientCoarsePrune)
      coarseParserState = coarseParser.getCoarsePrunedChart(ex);
    LogInfo.end_track();

    // draw a sample to compute gradient and expected reward
    LogInfo.begin_track("ReinforcementParserState.inferBySampling");
    sampleHistoryAndInfer();
    LogInfo.end_track();

    // Compute gradient
    setPredDerivations();
    if (parser.verbose(3))
      LogInfo.logs("Expected reward = %s", objectiveValue);
    visualizeChart();
  }

  public void execute() {};

  private void sampleHistoryAndInfer() {

    // add to chart the token and phrase parts
    for (Derivation deriv : gatherTokenAndPhraseDerivations())
      addToAgenda(SingleDerivationStream.constant(deriv));
    // add to agenda unaries where RHS is just terminals
    for (DerivationStream derivStream : gatherRhsTerminalsDerivations())
      addToAgenda(derivStream);

    ensureExecuted();

    while (continueParsing()) {

      unrollHighProbStreams();
      Pair<PrioritizedDerivationStream, Double> pdsAndProbability = sampler.sample();

      DerivationStream sampledDerivations = pdsAndProbability.getFirst().derivStream;
      Derivation sampledDerivation = sampledDerivations.next();
      updateBackpointers(sampledDerivations, sampledDerivation); // to be able to get all correct actions
      numItemsSampled++;

      assert sampledDerivation.isFeaturizedAndScored() : "top derivation is not featurized and scored: " + sampledDerivation;
      assert Math.abs(sampledDerivation.score - pdsAndProbability.getFirst().priority) < 1e-4 :
              sampledDerivation.score + " != " + pdsAndProbability.getFirst().priority;

      if (parser.verbose(2)) {
        LogInfo.begin_track("Item %d (|agenda|=%d), priority %s: |item|=%s -> %s %s %s [%s], prob=%s",
                numItemsSampled, agenda.size() + 1, Fmt.D(pdsAndProbability.getFirst().priority), sampledDerivations.estimatedSize(),
                sampledDerivation.cat, ex.spanString(sampledDerivation.start, sampledDerivation.end), sampledDerivation,
                sampledDerivation.rule, pdsAndProbability.getSecond());
      }

      // handle root derivations - get compatibility and record number of compatible derivations
      handleRootDerivation(ex, numItemsSampled, sampledDerivation);

      if (computeExpectedCounts) {
        Map<String, Double> counts = new HashMap<>();
        // add the feature vector and subtract for the time it was in the agenda unless has negative probability
        //pretty hacky
        if (pdsAndProbability.getSecond() > -0.0001)
          sampledDerivation.incrementLocalFeatureVector(1 - pdsAndProbability.getFirst().probSum, counts);
        else
          sampledDerivation.incrementLocalFeatureVector(-pdsAndProbability.getFirst().probSum, counts);
        if (parser.verbose(3))
          SempreUtils.logMap(counts, "agenda item gradient");
        ReinforcementUtils.addToDoubleMap(stateSequenceExpectedCounts, counts, parser.searchPrefix); // upate the gradient incrementally
      }
      // only after update of params we can change the chart and the agenda
      if (addToBoundedChart(sampledDerivation)) {
        if (parser.verbose(5))
          LogInfo.logs("ReinforcementParserState.infer: adding to chart %s(%s,%s) formula=%s",
                  sampledDerivation.cat, sampledDerivation.start, sampledDerivation.end, sampledDerivation.formula);
        combineWithChartDerivations(sampledDerivation);
      }
      addToAgenda(sampledDerivations);
      if (parser.verbose(2))
        LogInfo.end_track();
    }

    finalizeSearchExpectedCounts(); // gradient for remaining agenda items
    rerankRootDerivations(); // last action
    if (computeExpectedCounts) {
      computeGradient();
    }
  }

  private void unrollHighProbStreams() {

    if (samplingStrategy.equals("max")) return;

    sampler.unroll(); // if multiplicative, then unroll oracle stuff (ignore \beta currently!)

    if (parser.verbose(3))
      LogInfo.begin_track("Unrolling high probability streams");

    double lb=Double.NEGATIVE_INFINITY;
    int numOfHiddenStreams = 0;
    for (PrioritizedDerivationStream pds : agenda) {
      lb = NumUtils.logAdd(lb, pds.getScore());
      if(pds.derivStream.estimatedSize() > 1)
        numOfHiddenStreams++;
    }

    if (parser.verbose(3))
      LogInfo.logs("unrollHighProbStreams(): |agenda|=%s, lb=%s, |hiddenstreams|=%s", agenda.size(), lb, numOfHiddenStreams);

    List<Pair<DerivationStream, Double>> derivsToAdd = new ArrayList<>();
    List<Integer> indicesToRemove = new ArrayList<>();
    for (int i = 0; i < agenda.size(); ++i) {
      PrioritizedDerivationStream pds = agenda.get(i);
      boolean modified = false;
      while (pds.derivStream.hasNext() && pds.derivStream.estimatedSize() > 1 &&
              illegalStream(pds.derivStream, lb, pds.derivStream.estimatedSize(), numOfHiddenStreams)) {
        modified = true;
        Derivation nextDeriv = pds.derivStream.next();
        updateBackpointers(pds.derivStream, nextDeriv);

        DerivationStream derivStream = SingleDerivationStream.constant(nextDeriv);
        if (parser.verbose(3) && derivStream.hasNext()) {
          Derivation deriv  = derivStream.peek();
          LogInfo.logs("unrollIllegalStreams(): add deriv=%s(%s,%s) [%s] score=%s, |stream|=%s",
                  deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.score, pds.derivStream.estimatedSize());
        }
        derivsToAdd.add(Pair.newPair(derivStream, pds.probSum));
        //update lb
        if (pds.derivStream.hasNext()) {
          featurizeAndScoreDerivation(pds.derivStream.peek());
          lb = NumUtils.logAdd(lb, pds.getScore());
        }
        //update num of hidden streams
        if(pds.derivStream.estimatedSize() <= 1)
          numOfHiddenStreams--;
      }
      if (modified) {
        indicesToRemove.add(i);
        derivsToAdd.add(Pair.newPair(pds.derivStream, pds.probSum));
      }
    }

    // remove - need to make sure indices don't change due to removal so go from end to start
    for (int i = indicesToRemove.size() - 1; i >= 0; --i)
      agenda.remove(agenda.get(indicesToRemove.get(i)), indicesToRemove.get(i));

    // add
    for (Pair<DerivationStream, Double> pair : derivsToAdd)
      addToAgenda(pair.getFirst(), pair.getSecond());

    if (parser.verbose(3))
      LogInfo.logs("unrollHighProbStreams(): |agenda|=%s", agenda.size());
    if (parser.verbose(3))
      LogInfo.end_track();
  }

  private boolean illegalStream(DerivationStream derivStream, double logSum, int estimatedSize, int numOfHiddenStreams) {
    Derivation deriv = derivStream.peek();
    double firstItemLogProb = deriv.score - logSum; //log(exp(s(g_1))/L)
    double upperBound = Math.log(estimatedSize) + Math.log(numOfHiddenStreams); //log(M(g)|G'|)

    if (parser.verbose(3))
      LogInfo.logs("IllegalStream(): score=%s, logsum=%s, |stream|=%s, |hiddenstreams|=%s, deriv=%s(%s,%s) %s, sum=%s",
              deriv.score, logSum, estimatedSize, numOfHiddenStreams, deriv.cat, deriv.start, deriv.end, deriv.formula, firstItemLogProb+upperBound);

    return (firstItemLogProb+upperBound) > LOG_SMALL_PROB;
  }

  private boolean isHighProbStream(DerivationStream derivStream, double maxScore, int estimatedSize) {
    Derivation deriv = derivStream.peek();
    double gapFromMax = deriv.score - maxScore;
    double threshold = LOG_SMALL_PROB - Math.log(estimatedSize);
    if (parser.verbose(3))
      LogInfo.logs("isHighProbStream(): gapFromMax=%s, threshold=%s, deriv=%s(%s,%s) %s |stream|=%s", gapFromMax, threshold,
              deriv.cat, deriv.start, deriv.end, deriv.formula, derivStream.estimatedSize());

    return gapFromMax > threshold;
  }

  // recompute score using dot product of features and re-ranking feature weights
  private void rerankRootDerivations() {
    setPredDerivations();
    for (Derivation rootDeriv : predDerivations) {
      double oldScore = rootDeriv.score;
      rootDeriv.computeScore(params);
      if (parser.verbose(3))
        LogInfo.logs("ReinforcementParser.rerankRootDerivations: deriv=%s, old=%s, new=%s", rootDeriv, oldScore, rootDeriv.score);
    }
    Derivation.sortByScore(predDerivations);
  }

  private void updateBackpointers(DerivationStream stream, Derivation sampledDeriv) {
    Pair<ArrayList<Derivation>, Integer> pair = backpointerList.get(sampledDeriv.creationIndex);
    if (!stream.hasNext()) return;

    if (pair == null) {
      ArrayList<Derivation> list = new ArrayList<>();
      list.add(sampledDeriv);
      backpointerList.put(sampledDeriv.creationIndex, pair = Pair.newPair(list, 0));
    }
    List<Derivation> list = pair.getFirst();
    Derivation nextDeriv = stream.peek();
    list.add(nextDeriv);
    backpointerList.put(nextDeriv.creationIndex, Pair.newPair(pair.getFirst(), list.size() - 1));
  }

  private double computeExpectedReward(List<Derivation> predDerivations, double[] probs) {
    double rewardExpectation = 0d;
    for (int i = 0; i < predDerivations.size(); ++i) {
      rewardExpectation += probs[i] * compatibilityToReward(predDerivations.get(i).compatibility);
    }
    return rewardExpectation;
  }

  // q - proposal distribution
  // pi = model distribution
  private void computeGradient() {
    if (predDerivations.isEmpty()) return;

    double[] qDist = sampler.getDerivDistribution(predDerivations); //uniform over correct things when \beta is high
    double[] piDist = ReinforcementUtils.expNormalize(predDerivations);
    // compute E_q(R(d))
    LogInfo.begin_track("Computing gradient");
    double rewardExpectation = computeExpectedReward(predDerivations, qDist);

    // compute E_q(\phi(d) R(d)) and E_pi(\phi(d))
    Map<String, Double> featureExpectation = new HashMap<>(), rewardInfusedFeatureExpectation = new HashMap<>();
    for (int i = 0; i < predDerivations.size(); ++i) {
      Derivation deriv = predDerivations.get(i);
      deriv.incrementAllFeatureVector(piDist[i], featureExpectation);
      deriv.incrementAllFeatureVector(qDist[i] * compatibilityToReward(deriv.compatibility), rewardInfusedFeatureExpectation);
    }
    // final gradient computation
    Map<String, Double> sampleCounts = new HashMap<>();
    if (ReinforcementParser.opts.simulateNonRlObjective) {
      ParserState.computeExpectedCounts(predDerivations, sampleCounts);
    } else {
      sampleCounts = ReinforcementUtils.multiplyDoubleMap(stateSequenceExpectedCounts, rewardExpectation);
      SempreUtils.addToDoubleMap(sampleCounts, rewardInfusedFeatureExpectation);
      ReinforcementUtils.subtractFromDoubleMap(sampleCounts, ReinforcementUtils.multiplyDoubleMap(featureExpectation, rewardExpectation));
    }
    SempreUtils.addToDoubleMap(expectedCounts, sampleCounts);

    double sum = 0d;
    for (String key : sampleCounts.keySet()) {
      double value = sampleCounts.get(key);
      if (parser.verbose(3))
        LogInfo.logs("feature=%s, value=%s", key, value);
      sum += value * value;
    }
    LogInfo.logs("L2 norm: %s", Math.sqrt(sum));
    LogInfo.end_track();
  }

  private void createSampler(ReinforcementParserState oracleState) {
    if ("proposal".equals(samplingStrategy)) {
      if (oracleState == null)
        throw new RuntimeException("missing oracle state");
      this.sampler = new MultiplicativeProposalSampler(oracleState);
    } else if ("max".equals(samplingStrategy)) {
      this.sampler = new MaxSampler();
    } else if ("agenda".equals(samplingStrategy) || samplingStrategy == null) // default
      this.sampler = new AgendaSampler();
  }

  // info for visualizing chart
  private void visualizeChart() {
    if (parser.chartFillOut != null && Parser.opts.visualizeChartFilling) {
      parser.chartFillOut.println(Json.writeValueAsStringHard(new ChartFillingData(ex.id, chartFillingList,
              ex.utterance, ex.numTokens())));
      parser.chartFillOut.flush();
    }
  }

  // go over all agenda items that we did not subtract counts for and finalize
  private void finalizeSearchExpectedCounts() {
    if (ReinforcementParser.opts.simulateNonRlObjective) return;
    if (!computeExpectedCounts) return;
    Map<String, Double> counts = new HashMap<>();
    for (PrioritizedDerivationStream pds : agenda) {
      pds.derivStream.peek().incrementLocalFeatureVector(pds.probSum, counts);
    }
    if (parser.verbose(3)) {
      SempreUtils.logMap(counts, "subtracted");
    }
    ReinforcementUtils.subtractFromDoubleMap(stateSequenceExpectedCounts, counts, parser.searchPrefix);
    if (parser.verbose(3)) {
      SempreUtils.logMap(stateSequenceExpectedCounts, "final search gradient");
    }
  }

  private void handleRootDerivation(Example ex, int numItemsSampled, Derivation sampledDerivation) {
    if (!sampledDerivation.isRoot(ex.numTokens())) return;

    sampledDerivation.ensureExecuted(parser.executor, ex.context);
    if (ex.targetValue != null)
      sampledDerivation.compatibility = parser.valueEvaluator.getCompatibility(ex.targetValue, sampledDerivation.value);
    if (Parser.opts.partialReward ? (sampledDerivation.compatibility > 0) : (sampledDerivation.compatibility == 1)) {
      if (parser.verbose(2))
        LogInfo.logs("Top-level %s: reward = %s", numItemsSampled, sampledDerivation.compatibility);
      // put in position 0 the derivation with best compatibility
      correctDerivations.add(sampledDerivation);
      if (correctDerivations.get(0).compatibility < sampledDerivation.compatibility) {
        Collections.swap(correctDerivations, 0, correctDerivations.size() - 1);
      }

      if (firstCorrectItem == -1)
        firstCorrectItem = numItemsSampled;
    }
  }

  public void setEvaluation() {
    LogInfo.begin_track_printAll("ReinforcementParserParserState.setEvaluation");
    super.setEvaluation();

    if (coarseParserState != null)
      evaluation.add("coarseParseTime", coarseParserState.getCoarseParseTime());
    if (firstCorrectItem != -1)
      evaluation.add("firstCorrectItem", firstCorrectItem);
    LogInfo.end_track();
  }

  // Defines how to sample the next state
  abstract class Sampler {
    // returns a derivation stream with the sample probability
    public abstract Pair<PrioritizedDerivationStream, Double> sample();
    public abstract double[] getDerivDistribution(List<Derivation> rootDerivs);
    public abstract void unroll();

    // go over agenda and update the probability sum for gradient computation before sampling
    public void updateProbSum(double[] modelProbs) {
      for (int i = 0; i < agenda.size(); ++i) {
        PrioritizedDerivationStream pds = agenda.get(i);
        pds.addProb(modelProbs[i]);
        if (parser.verbose(3))
          LogInfo.logs("updateProbSum(): deriv=%s, probSum=%s", pds.derivStream.peek(), pds.probSum);
      }
    }
  }

  // sample from agenda based on agenda scores
  class AgendaSampler extends Sampler {

    @Override

    public Pair<PrioritizedDerivationStream, Double> sample() {

      double[] modelProbs = ReinforcementUtils.expNormalize(agenda);
      if (computeExpectedCounts)// compute probability sum before sampling for gradient computation (easier before sampling)
        updateProbSum(modelProbs);

      int sampledIndex = SampleUtils.sampleMultinomial(randGen, modelProbs);
      PrioritizedDerivationStream pds = agenda.get(sampledIndex);
      double prob = modelProbs[sampledIndex];
      agenda.remove(pds, sampledIndex);
      return Pair.newPair(pds, prob);
    }

    @Override
    public double[] getDerivDistribution(List<Derivation> rootDerivs) {
      return ReinforcementUtils.expNormalize(rootDerivs);
    }

    @Override
    public void unroll() { }
  }

  class MaxSampler extends Sampler {

    @Override
    public Pair<PrioritizedDerivationStream, Double> sample() {
      PrioritizedDerivationStream pds = agenda.pop();
      return Pair.newPair(pds, 1d);
    }

    @Override
    public double[] getDerivDistribution(List<Derivation> rootDerivs) {
      double[] res = new double[rootDerivs.size()];
      Arrays.fill(res, 0d);
      res[0] = 1d;
      return res;
    }

    @Override
    public void unroll() { }
  }

  // sample categories and spans that are in a correct derivation
  class OracleInfo  {

    private List<DerivInfo> necessaryDerivInfos; // info for derivations necessary to generate oracle derivations
    private List<DerivInfo> oracleDerivInfos; // info for oracle derivations
    private Map<Long, Pair<ArrayList<Derivation>, Integer>> backPointers;
    private NecessaryDeriv[] necessaryDerivsCache; // memorizing whether a deriv is necessary or not
    long firstCorrectDerivNumber = -1; // offset for necessaryDerivs

    // we sample from the end to the start
    public OracleInfo(ReinforcementParserState oracleState) {
      if (oracleState == null) throw new RuntimeException("oracle state is null");
      this.necessaryDerivInfos = new ArrayList<>();
      this.oracleDerivInfos = new ArrayList<>();
      if (!oracleState.correctDerivations.isEmpty()) {
        Collections.sort(oracleState.correctDerivations, new CorrectDerivationComparator());
        this.backPointers = oracleState.backpointerList;
        Derivation oracleDeriv = oracleState.correctDerivations.get(0);
        this.firstCorrectDerivNumber = oracleDeriv.creationIndex;

        LogInfo.logs("OracleSampler: deriv=%s, comp=%s", oracleDeriv, oracleDeriv.compatibility);
        populateCorrectDerivations(oracleDeriv);

        if (parser.verbose(2)) {
          LogInfo.begin_track("OracleSampler: necessary infos:");
          for (DerivInfo necessaryInfo : necessaryDerivInfos) LogInfo.log(necessaryInfo);
          LogInfo.end_track();
          LogInfo.begin_track("OracleSampler: oracle infos:");
          for (DerivInfo oracleInfo : oracleDerivInfos) LogInfo.log(oracleInfo);
          LogInfo.end_track();
        }
      }
    }

    private void populateCorrectDerivations(Derivation oracleDeriv) {
      // add derivation info and also all upstream derivations
      if (parser.verbose(4))
        LogInfo.logs("populateCorrectDerivations(): oracle deriv: %s", oracleDeriv);

      Pair<ArrayList<Derivation>, Integer> listAndIndex = this.backPointers.get(oracleDeriv.creationIndex);
      if (listAndIndex != null) {
        for (int i = listAndIndex.getSecond() - 1; i >= 0; i--) {
          Derivation deriv = listAndIndex.getFirst().get(i);
          if (parser.verbose(4))
            LogInfo.logs("populateCorrectDerivations(): necessary deriv: %s", deriv);
          DerivInfo derivInfo = new DerivInfo(deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.rule);
          if (!necessaryDerivInfos.contains(derivInfo))
            necessaryDerivInfos.add(derivInfo);
        }
      }
      DerivInfo derivInfo = new DerivInfo(oracleDeriv.cat, oracleDeriv.start, oracleDeriv.end, oracleDeriv.formula, oracleDeriv.rule);
      if (!oracleDerivInfos.contains(derivInfo)) {
        necessaryDerivInfos.add(derivInfo);
        oracleDerivInfos.add(derivInfo);
      }

      // recurse
      for (Derivation child : oracleDeriv.children) {
        populateCorrectDerivations(child);
      }
    }

    // checking if derivation is necessary using the cache
    protected boolean isNecessaryDeriv(Derivation deriv) {

      if (necessaryDerivInfos.isEmpty()) return false;
      if (necessaryDerivsCache == null) {
        necessaryDerivsCache = new NecessaryDeriv[200000];
        Arrays.fill(necessaryDerivsCache, NecessaryDeriv.UNKNOWN);
      }
      int index = (int) (deriv.creationIndex - firstCorrectDerivNumber);
      if (index < 0) throw new RuntimeException("Negative index - correct index larger than deriv number");
      if (index >= 200000) {
        LogInfo.warnings("isNecessaryDeriv(): index larger than 200000: %s", index);
        return necessaryDerivInfos.contains(new DerivInfo(deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.rule));
      }
      if (necessaryDerivsCache[index] == NecessaryDeriv.UNNECESSARY_DERIV)
        return false;
      if (necessaryDerivsCache[index] == NecessaryDeriv.NECESSARY_DERIV)
        return true;
      // unknown
      boolean res = necessaryDerivInfos.contains(new DerivInfo(deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.rule));
      necessaryDerivsCache[index] = res ? NecessaryDeriv.NECESSARY_DERIV : NecessaryDeriv.UNNECESSARY_DERIV;
      return res;
    }
  }

  class MultiplicativeProposalSampler extends Sampler {

    private double bonus;
    private OracleInfo oracleInfo;

    public MultiplicativeProposalSampler(ReinforcementParserState oracleState) {
      oracleInfo = new OracleInfo(oracleState);
      bonus = ReinforcementParser.opts.multiplicativeBonus;
      LogInfo.logs("Bonus=%s", bonus);
    }

    // We assume that oracle stuff has been unrolled
    @Override
    public Pair<PrioritizedDerivationStream, Double> sample() {

      double[] modelProbs = ReinforcementUtils.expNormalize(agenda);
      double[] samplerProbs = getUnnormalizedAgendaDistribution();
      if (!NumUtils.expNormalize(samplerProbs)) throw new RuntimeException("Normalization failed" + Arrays.toString(samplerProbs));

      int sampledIndex = ReinforcementUtils.sampleIndex(randGen, samplerProbs);
      PrioritizedDerivationStream pds = agenda.get(sampledIndex);
      double prob = samplerProbs[sampledIndex];

      if (parser.verbose(3)) {
        Derivation deriv = pds.derivStream.peek();
        if (oracleInfo.oracleDerivInfos.contains(new DerivInfo(deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.rule)))
          LogInfo.logs("MultiplicativeProposalSampler.sample(): Sampled from correct!, prob=%s", prob);
        else
          LogInfo.logs("MultiplicativeProposalSampler.sample(): Sampled from incorrect!, prob=%s", prob);
      }
      boolean returnProb = true;

      // whether to update only for correct moves or not hack
      if (computeExpectedCounts && ReinforcementParser.opts.updateGradientForCorrectMovesOnly) {
        Derivation deriv = pds.derivStream.peek();
        DerivInfo derivInfo = new DerivInfo(deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.rule);
        if (oracleInfo.oracleDerivInfos.contains(derivInfo))
          updateProbSum(modelProbs);
        else returnProb = false;

        if (parser.verbose(3)) {
          LogInfo.logs("Updating gradient=%s", returnProb);
        }
      } else {
        if (computeExpectedCounts) // compute probability sum before sampling for gradient computation (easier before sampling)
          updateProbSum(modelProbs);
      }

      agenda.remove(pds, sampledIndex);
      return Pair.newPair(pds, returnProb ? prob : -1d);
    }

    //todo : This unrolls all necessary derivations, and ignore \beta
    //assumes probability mass on non-necessary things is very very small
    //might need to be fixed if we anneal \beta to 0
    @Override
    public void unroll() {

      if (parser.verbose(3))
        LogInfo.begin_track("MultiplicativeBonusSampler.unroll()");

      List<Pair<DerivationStream, Double>> derivsToAdd = new ArrayList<>();
      List<Integer> indicesToRemove = new ArrayList<>();

      for (int i = 0; i < agenda.size(); ++i) {
        PrioritizedDerivationStream pds = agenda.get(i);
        boolean modified = false;
        while (pds.derivStream.hasNext() && pds.derivStream.estimatedSize() > 1 && oracleInfo.isNecessaryDeriv(pds.derivStream.peek())) {
          modified = true;
          Derivation nextDeriv = pds.derivStream.next();
          DerivationStream newDerivStream = SingleDerivationStream.constant(nextDeriv);
          if (parser.verbose(3) && newDerivStream.hasNext()) {
            Derivation deriv  = newDerivStream.peek();
            LogInfo.logs("MultiplicativeSampler.unroll(): add necessary deriv=%s(%s,%s) [%s] score=%s, |stream|=%s, creationIndex=%s",
                    deriv.cat, deriv.start, deriv.end, deriv.formula, deriv.score, pds.derivStream.estimatedSize(), deriv.creationIndex);
          }
          derivsToAdd.add(Pair.newPair(newDerivStream, pds.probSum));
          if (pds.derivStream.hasNext())
            featurizeAndScoreDerivation(pds.derivStream.peek());
        }
        if (modified) {
          indicesToRemove.add(i);
          derivsToAdd.add(Pair.newPair(pds.derivStream, pds.probSum));
        }
      }
      // remove - need to make sure indices don't change due to removal so go from end to start
      for (int i = indicesToRemove.size() - 1; i >= 0; --i)
        agenda.remove(agenda.get(indicesToRemove.get(i)), indicesToRemove.get(i));
      // add
      for (Pair<DerivationStream, Double> pair: derivsToAdd)
        addToAgenda(pair.getFirst(), pair.getSecond());
      if (parser.verbose(3))
        LogInfo.end_track();
    }

    private double[] getUnnormalizedAgendaDistribution() {
      double[] probs = new double[agenda.size()];
      for (int i = 0; i < agenda.size(); ++i) {
        Derivation d = agenda.get(i).derivStream.peek();
        probs[i] = d.score;
        // we assume all necessary things have been unrolled already so no need to handle that
        if (oracleInfo.oracleDerivInfos.contains(new DerivInfo(d.cat, d.start, d.end, d.formula, d.rule))) {
          probs[i] += bonus;
        }
      }
      return probs;
    }

    @Override
    public double[] getDerivDistribution(List<Derivation> rootDerivs) {
      double[] res = new double[rootDerivs.size()];
      for (int i = 0; i < rootDerivs.size(); ++i) {
        Derivation rootDeriv = rootDerivs.get(i);
        res[i] = rootDeriv.score + bonus * rootDeriv.compatibility;
      }
      NumUtils.expNormalize(res);
      return res;
    }
  }

  //A heuristic for choosing the oracle derivation. It'd be good to get rid of this or simplify
  public static class CorrectDerivationComparator implements Comparator<Derivation> {
    @Override
    public int compare(Derivation deriv1, Derivation deriv2) {
      if (deriv1.compatibility > deriv2.compatibility) return -1;
      if (deriv1.compatibility < deriv2.compatibility) return +1;

      boolean deriv1Join = containsJoin(deriv1);
      boolean deriv2Join = containsJoin(deriv2);
      if (deriv1Join && !deriv2Join) return -1;
      if (!deriv1Join && deriv2Join) return +1;
      // by score
      if (deriv1.score > deriv2.score) return -1;
      if (deriv1.score < deriv2.score) return +1;
      // Ensure reproducible randomness
      if (deriv1.creationIndex < deriv2.creationIndex) return -1;
      if (deriv1.creationIndex > deriv2.creationIndex) return +1;
      return 0;
    }

    private boolean containsJoin(Derivation d) {
      SemanticFn semanticFn = d.rule.getSem();
      if (semanticFn != null) {
        if (semanticFn instanceof JoinFn)
          return true;
      }
      for (Derivation child : d.children) {
        if (containsJoin(child))
          return true;
      }
      return false;
    }
  }
}

//holds the stream, the priority, and a probability sum to make grdient computation efficient
class PrioritizedDerivationStream implements Comparable<PrioritizedDerivationStream>, HasScore {
  public final DerivationStream derivStream;
  public final double priority;
  public double probSum;

  PrioritizedDerivationStream(DerivationStream derivStream, double priority, double probSum) {
    this.derivStream = derivStream;
    this.priority = priority;
    this.probSum = probSum;
  }

  @Override
  public int compareTo(PrioritizedDerivationStream o) {
    if (this.priority > o.priority) return -1;
    if (this.priority < o.priority) return +1;
    return 0;
  }

  public double getScore() { return derivStream.peek().score; }
  public void addProb(double prob) { probSum += prob; }
}

class DerivInfo {
  public final String cat;
  public final int start;
  public final int end;
  public final Formula formula;
  public final Rule rule;

  DerivInfo(String cat, int start, int end, Formula formula, Rule rule) {
    this.cat = cat;
    this.start = start;
    this.end = end;
    this.formula = formula;
    this.rule = rule;
  }

  @Override
  public boolean equals(Object o) {
    if (this == o) return true;
    if (o == null || getClass() != o.getClass()) return false;

    DerivInfo derivInfo = (DerivInfo) o;

    if (end != derivInfo.end) return false;
    if (start != derivInfo.start) return false;
    if (cat != null ? !cat.equals(derivInfo.cat) : derivInfo.cat != null) return false;
    if (formula != null ? !formula.equals(derivInfo.formula) : derivInfo.formula != null) return false;
    return !(rule != null ? !rule.equals(derivInfo.rule) : derivInfo.rule != null);

  }

  @Override
  public int hashCode() {
    int result = cat != null ? cat.hashCode() : 0;
    result = 31 * result + start;
    result = 31 * result + end;
    result = 31 * result + (formula != null ? formula.hashCode() : 0);
    result = 31 * result + (rule != null ? rule.hashCode() : 0);
    return result;
  }

  public String toString() {
    return cat + "(" + start + "," + end + ") " + formula.toString();
  }
}

