/**
 * @file dDNNF.cpp
 * @brief d-DNNF circuit operations and evaluation algorithms.
 *
 * Implements all methods of @c dDNNF declared in @c dDNNF.h:
 * - @c vars(): set of reachable input (IN) gates.
 * - @c makeSmooth(): smooth the circuit so every OR gate's children
 *   mention the same variable set.
 * - @c makeGatesBinary(): binarise n-ary AND/OR gates.
 * - @c simplify(): constant propagation.
 * - @c condition() / @c conditionAndSimplify(): fix one variable.
 * - @c probabilityEvaluation(): exact probability in linear time.
 * - @c shapley() / @c banzhaf(): power index computation.
 * - @c topological_order(): DFS topological sort.
 *
 * The private helpers @c shapley_delta() and @c shapley_alpha() implement
 * the polynomial-time Shapley-value algorithm for d-DNNFs.
 */
#include "dDNNF.h"
#include "Circuit.hpp"

#include <unordered_map>
#include <set>
#include <stack>
#include <variant>
#include <cassert>
#include <algorithm>
#include <functional>
#include <numeric>
#include <sstream>
#include <iomanip>

std::unordered_set<gate_t> dDNNF::vars(gate_t root) const
{
  // No recursion, so using a stack to handle all children; order of
  // calls is here irrelevant
  std::stack<gate_t> to_process;
  std::unordered_set<gate_t> processed;
  to_process.push(root);

  std::unordered_set<gate_t> result;
  while(!to_process.empty())
  {
    auto g = to_process.top();
    to_process.pop();

    if(processed.find(g)==processed.end()) {
      if(getGateType(g)==BooleanGate::IN)
        result.insert(g);
      else {
        for(auto c: getWires(g))
          to_process.push(c);
      }
      processed.insert(g);
    }
  }

  return result;
}

void dDNNF::makeSmooth()
{
  gate_t original_gates_size{gates.size()};
  // gates.size() might change in the loop, but newly added gates do not
  // need to be iterated upon

  for(gate_t g = gate_t{0}; g<original_gates_size; ++g) {
    if(getGateType(g)!=BooleanGate::OR || getWires(g).size()<=1)
      continue;

    std::unordered_set<gate_t> all;
    std::vector<std::unordered_set<gate_t> > varss;

    std::for_each(getWires(g).begin(), getWires(g).end(),
                  [&](gate_t x) {
      varss.push_back(vars(x));
    });
    std::for_each(varss.begin(), varss.end(),
                  [&](const auto &s) {
      all.insert(s.begin(),s.end());
    });

    for(auto v: all) {
      bool modified = false;
      for(size_t i=0; i<varss.size(); ++i) {
        if(varss[i].find(v)==varss[i].end()) {
          if(!modified) {
            if(getGateType(getWires(g)[i])!=BooleanGate::AND) {
              gate_t and_gate = setGate(BooleanGate::AND);
              addWire(and_gate, getWires(g)[i]);
              getWires(g)[i]=and_gate;
            }
            modified = true;
          }

          gate_t dummy_or_gate = setGate(BooleanGate::OR);
          gate_t dummy_not_gate = setGate(BooleanGate::NOT);
          addWire(dummy_or_gate, v);
          addWire(dummy_not_gate, v);
          addWire(dummy_or_gate, dummy_not_gate);
          addWire(getWires(g)[i],dummy_or_gate);
        }
      }
    }
  }
}

void dDNNF::makeGatesBinary(BooleanGate type)
{
  for(gate_t g{0}; g<gates.size(); ++g) {
    if(getGateType(g)!=type || getWires(g).size()<=2)
      continue;

    if(getWires(g).size()==3) {
      const gate_t child = setGate(type);
      auto &w = getWires(g);
      for(size_t i=1; i<3; ++i) {
        addWire(child, w[i]);
      }
      w.resize(1);
      addWire(g, child);
    } else {
      const gate_t child1 = setGate(type);
      const gate_t child2 = setGate(type);

      auto &w = getWires(g);
      const auto k = w.size();

      for(unsigned i=0; i<k; ++i)
        if(i<k/2)
          addWire(child1, w[i]);
        else
          addWire(child2, w[i]);
      w.clear();
      addWire(g, child1);
      addWire(g, child2);
    }
  }
}

double dDNNF::probabilityEvaluation() const
{
  if (gates.size() == 0)
    return 0.;

  // Unfortunately, dDNNFs can be quite deep so we need to simulate
  // recursion with a heap-based stack, to avoid exhausting the actual
  // memory stack
  using RecursionParams = struct {
    gate_t g;
    size_t children_processed;
    double partial_value;
  };
  using RecursionResult = double;
  std::stack<std::variant<RecursionParams,RecursionResult> > stack;
  stack.emplace(RecursionParams{root,0,0.});

  while(1) {
    double child_value{0.};

    if(stack.top().index()==1) {   // RecursionResult
      child_value=std::get<1>(stack.top());
      stack.pop();
    }

    auto [g, children_processed, partial_value]=std::get<0>(stack.top());
    stack.pop();

    auto it = probability_cache.find(g);

    if(it!=probability_cache.end()) {
      if(stack.empty())
        return it->second;
      else
        stack.emplace(it->second);
    } else {
      if(children_processed==0) {
        switch(getGateType(g)) {
        case BooleanGate::IN:
          partial_value = getProb(g);
          break;
        case BooleanGate::NOT:
          // Placeholder; the complement is taken when the child's value
          // comes back (NOT gates descend into their child like AND/OR,
          // so a NOT over an internal gate -- e.g. the De Morgan
          // structures interpretAsDD builds for uncertified ORs -- is
          // evaluated correctly, not read as 1 - prob-array-default).
          partial_value = 0;
          break;
        case BooleanGate::AND:
          partial_value = 1;
          break;
        case BooleanGate::OR:
          partial_value = 0;
          break;
        default:
          assert(false);
        }
      } else {
        if(getGateType(g) == BooleanGate::AND) {
          partial_value *= child_value;
        } else if(getGateType(g) == BooleanGate::NOT) {
          partial_value = 1-child_value;
        } else {   // BooleanGate::OR
          partial_value += child_value;
        }
      }

      if(children_processed<getWires(g).size()) {
        stack.emplace(RecursionParams{g,children_processed+1,partial_value});
        stack.emplace(RecursionParams{getWires(g)[children_processed],0,0.});
      } else {
        double result = partial_value;
        probability_cache[g]=result;

        if(stack.empty())
          return result;
        else
          stack.emplace(result);
      }
    }
  }

  // We return from within the while loop, when the stack is empty
  assert(false);
}

double dDNNF::banzhaf_internal() const {
  std::unordered_map<gate_t, double> result;
  std::unordered_map<gate_t, double> prod_one_plus_p;

  // Stack to simulate recursion: contains a pair (node, b) where b
  // indicates whether this is the beginning (false) or ending (true) of
  // the processing of a node
  std::stack<std::pair<gate_t, bool> > stack;
  stack.emplace(std::make_pair(root, false));

  while(!stack.empty())
  {
    auto [node, b] = stack.top();
    stack.pop();

    if(result.find(node)!=result.end()) {
      // Already processed, skip
      continue;
    }

    switch(getGateType(node)) {
    case BooleanGate::IN:
      result[node] = getProb(node);
      prod_one_plus_p[node] = 1+getProb(node);
      break;

    case BooleanGate::NOT:
      if(!b) {
        stack.push(std::make_pair(node, true));
        stack.push(std::make_pair(getWires(node)[0], false));
      } else {
        auto child = getWires(node)[0];
        result[node] = prod_one_plus_p[child] - result[child];
        prod_one_plus_p[node] = prod_one_plus_p[child];
      }
      break;

    case BooleanGate::OR:
      if(!b) {
        if(getWires(node).size()==0) { // Has to be an OR False gate
          result[node] = 0.;
          prod_one_plus_p[node] = 1.;
        } else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        result[node] =
          std::accumulate(getWires(node).begin(), getWires(node).end(), 0.,
                          [&](auto r, auto g) {
          return r + result[g];
        });
        prod_one_plus_p[node] = prod_one_plus_p[getWires(node)[0]];
      }
      break;

    case BooleanGate::AND:
      if(!b) {
        if(getWires(node).size()==0) { // Has to be an AND True gate
          result[node] = 1.;
          prod_one_plus_p[node] = 1.;
        } else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        result[node] =
          std::accumulate(getWires(node).begin(), getWires(node).end(), 1.,
                          [&](auto r, auto g) {
          return r * result[g];
        });
        prod_one_plus_p[node] =
          std::accumulate(getWires(node).begin(), getWires(node).end(), 1.,
                          [&](auto r, auto g) {
          return r * prod_one_plus_p[g];
        });
      }
      break;

    case BooleanGate::MULIN:
    case BooleanGate::MULVAR:
    case BooleanGate::UNDETERMINED:
      throw CircuitException("Incorrect gate type");
      break;
    }
  }

  return result[root];
}

std::unordered_map<gate_t, std::vector<double> > dDNNF::shapley_delta() const {
  std::unordered_map<gate_t, std::vector<double> > result;

  if(!isProbabilistic())
    return result;

  // Stack to simulate recursion: contains a pair (node, b) where b
  // indicates whether this is the beginning (false) or ending (true) of
  // the processing of a node
  std::stack<std::pair<gate_t, bool> > stack;
  stack.emplace(std::make_pair(root, false));

  while(!stack.empty())
  {
    auto [node, b] = stack.top();
    stack.pop();

    if(result.find(node)!=result.end()) {
      // Already processed, skip
      continue;
    }

    switch(getGateType(node)) {
    case BooleanGate::IN:
      result[node] = {1-getProb(node), getProb(node)};
      break;

    case BooleanGate::NOT:
    case BooleanGate::OR:
      if(!b) {
        if(getWires(node).size()==0) // Has to be an OR False gate
          result[node] = {1};
        else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        result[node] = result[getWires(node)[0]];
      }
      break;

    case BooleanGate::AND:
    {
      if(!b) {
        if(getWires(node).size()==0) // Has to be an AND True gate
          result[node] = {1};
        else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        if(getWires(node).size()==1)
          result[node] = result[getWires(node)[0]];
        else {
          assert(getWires(node).size()==2); // AND has been made binary
          const auto &r1 = result[getWires(node)[0]];
          const auto &r2 = result[getWires(node)[1]];
          const auto n1=r1.size()-1;
          const auto n2=r2.size()-1;
          for(size_t k=0; k<=n1+n2; ++k) {
            double r = 0.;
            for(size_t k1=std::max(0,static_cast<int>(k-n2)); k1<=std::min(k,n1); ++k1) {
              r+=r1[k1]*r2[k-k1];
            }
            result[node].push_back(r);
          }
        }
      }

      break;
    }

    case BooleanGate::MULIN:
    case BooleanGate::MULVAR:
    case BooleanGate::UNDETERMINED:
      throw CircuitException("Incorrect gate type");
      break;
    }
  }

  return result;
}

/**
 * @brief Compute the binomial coefficient C(n, k).
 * @param n  Total number of elements.
 * @param k  Number of elements to choose; must satisfy k ≤ n.
 * @return   The binomial coefficient n-choose-k.
 */
static long long comb(unsigned n, unsigned k)
{
  assert(k<=n);

  if(k == 0)
    return 1;
  else if(k > n/2)
    return comb(n,n-k);
  else return n * comb(n-1,k-1) / k;
}

std::vector<std::vector<double> > dDNNF::shapley_alpha() const {
  std::unordered_map<gate_t, std::vector<double> > delta {shapley_delta()};
  std::unordered_map<gate_t, std::vector<std::vector<double> > > result;

  // Stack to simulate recursion: contains a pair (node, b) where b
  // indicates whether this is the beginning (false) or ending (true) of
  // the processing of a node
  std::stack<std::pair<gate_t, bool> > stack;
  stack.emplace(std::make_pair(root, false));

  while(!stack.empty())
  {
    auto [node, b] = stack.top();
    stack.pop();

    if(result.find(node)!=result.end()) {
      // Already processed, skip
      continue;
    }

    switch(getGateType(node)) {
    case BooleanGate::IN:
      result[node] = {{0},{0,getProb(node)}};
      break;

    case BooleanGate::NOT:
      if(!b) {
        stack.push(std::make_pair(node, true));
        stack.push(std::make_pair(getWires(node)[0], false));
      } else {
        result[node] = result[getWires(node)[0]];
        auto k0=isProbabilistic()?0:result[node].size()-1;
        for(unsigned k=k0; k<result[node].size(); ++k)
          for(unsigned l=0; l<=k; ++l) {
            result[node][k][l] *= -1;
            result[node][k][l] += comb(k,l)*(isProbabilistic()?delta[node][k]:1);
          }
      }
      break;

    case BooleanGate::OR:
      if(!b) {
        if(getWires(node).size()==0) // Has to be an OR False gate
          result[node] = {{0.}};
        else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        result[node] = result[getWires(node)[0]];
        for(size_t i=1; i<getWires(node).size(); ++i) {
          const auto &r = result[getWires(node)[i]];
          auto k0=isProbabilistic()?0:r.size()-1;
          for(unsigned k=k0; k<r.size(); ++k)
            for(unsigned l=0; l<r[k].size(); ++l)
              result[node][k][l]+=r[k][l];
        }
      }
      break;

    case BooleanGate::AND:
      if(!b) {
        if(getWires(node).size()==0) // Has to be an AND True gate
          result[node] = {{1.}};
        else {
          stack.push(std::make_pair(node, true));
          for(auto c: getWires(node))
            stack.push(std::make_pair(c, false));
        }
      } else {
        if(getWires(node).size()==1)
          result[node] = result[getWires(node)[0]];
        else {
          assert(getWires(node).size()==2); // AND has been made binary
          const auto &r1 = result[getWires(node)[0]];
          const auto &r2 = result[getWires(node)[1]];
          const auto n1=r1.size()-1;
          const auto n2=r2.size()-1;
          result[node].resize(n1+n2+1);
          auto k0=isProbabilistic()?0:n1+n2;
          for(size_t k=k0; k<=n1+n2; ++k) {
            result[node][k].resize(k+1);
            for(size_t l=0; l<=k; ++l) {
              for(size_t k1=std::max(0,static_cast<int>(k-n2)); k1<=std::min(k,n1); ++k1)
                for(size_t l1=std::max(0,static_cast<int>(l-k+k1)); l1<=std::min(k1,l); ++l1)
                  result[node][k][l] += r1[k1][l1] * r2[k-k1][l-l1];
            }
          }
        }
      }
      break;

    case BooleanGate::MULIN:
    case BooleanGate::MULVAR:
    case BooleanGate::UNDETERMINED:
      throw CircuitException("Incorrect gate type");
      break;
    }
  }

  return result[root];
}

double dDNNF::shapley(gate_t var) const {
  auto cond_pos = condition(var, true);
  auto cond_neg = condition(var, false);

  auto alpha_pos=cond_pos.shapley_alpha();
  auto alpha_neg=cond_neg.shapley_alpha();

  double result=0.;

  double k0=isProbabilistic()?0:alpha_pos.size()-1;
  for(size_t k=k0; k<alpha_pos.size(); ++k)
    for(size_t l=0; l<=k; ++l) {
      double pos = alpha_pos[k][l];
      double neg = alpha_neg[k][l];
      result += (pos-neg)/comb(k,l)/(k+1);
    }

  result *= getProb(var);

  // Avoid rounding errors that make expected Shapley value outside of [-1,1]
  if(result>1.)
    result=1.;
  else if(result<-1.)
    result=-1.;

  return result;
}

double dDNNF::banzhaf(gate_t var) const {
  auto cond_pos = condition(var, true);
  auto cond_neg = condition(var, false);

  auto env_pos=cond_pos.banzhaf_internal();
  auto env_neg=cond_neg.banzhaf_internal();

  return getProb(var) * (env_pos-env_neg);
}

dDNNF dDNNF::condition(gate_t var, bool value) const {
  assert(getGateType(var)==BooleanGate::IN);

  dDNNF result=*this;

  result.setGateType(var, value ? BooleanGate::AND : BooleanGate::OR);
  result.probability_cache[var] = value?1.:0.;
  result.inputs.erase(var);
  auto it = id2uuid.find(var);
  if(it!=id2uuid.end()) {
    result.uuid2id.erase(it->second);
    result.id2uuid.erase(var);
  }

  return result;
}

std::vector<gate_t> dDNNF::topological_order(const std::vector<std::vector<gate_t> > &reversedWires) const
{
  std::vector<gate_t> result;

  std::stack<gate_t> nodesToProcess;
  std::vector<size_t> inDegree(wires.size());

  for(size_t g=0; g<wires.size(); ++g)
    if(!(inDegree[g] = wires[g].size()))
      nodesToProcess.push(gate_t{g});

  while(!nodesToProcess.empty()) {
    auto g = nodesToProcess.top();
    nodesToProcess.pop();
    result.push_back(g);
    for(auto p: reversedWires[static_cast<size_t>(g)])
      if(!(--inDegree[static_cast<size_t>(p)]))
        nodesToProcess.push(p);
  }

  return result;
}

void dDNNF::simplify() {
  std::vector<std::vector<gate_t> > reversedWires(gates.size());
  for(size_t i=0; i<wires.size(); ++i)
    for(auto g: wires[i])
      reversedWires[static_cast<size_t>(g)].push_back(gate_t{i});

  for(auto node: topological_order(reversedWires)) {
    auto &w = wires[static_cast<size_t>(node)];

    switch(getGateType(node)) {
    case BooleanGate::IN:
      break;

    case BooleanGate::AND:
    case BooleanGate::OR:
      // First, drop same-type empty children (AND of TRUE = TRUE; OR
      // of FALSE = FALSE) and short-circuit on opposite-type empty
      // children (AND of FALSE = FALSE; OR of TRUE = TRUE). May leave
      // `w` with 0 or 1 entry, which the second block then collapses.
      if(w.size()>1) {
        bool shorted = false;
        for(auto c=w.begin(); c!=w.end();) {
          if(getGateType(*c)==getGateType(node) && getWires(*c).size()==0) {
            c = w.erase(c);
          } else if(getGateType(*c)==(getGateType(node)==BooleanGate::AND?BooleanGate::OR:BooleanGate::AND)
                    && getWires(*c).size()==0) {
            setGateType(node, getGateType(*c));
            probability_cache[node] = getGateType(*c)==BooleanGate::AND?1.:0.;
            w.clear();
            shorted = true;
            break;
          } else {
            ++c;
          }
        }
        if(shorted) break;
      }
      if(w.size()==0) {
        probability_cache[node]=(getGateType(node)==BooleanGate::AND?1.:0.);
      } else if(w.size()==1) {
        if(node==getRoot()) {
          root=w[0];
        } else {
          for(auto p: reversedWires[static_cast<size_t>(node)])
            std::replace(wires[static_cast<size_t>(p)].begin(), wires[static_cast<size_t>(p)].end(), node, w[0]);
        }
        w.clear();
      }
      break;

    case BooleanGate::NOT:
      if(getGateType(w[0])==BooleanGate::AND && getWires(w[0]).size()==0) {
        setGateType(node, BooleanGate::OR);
        probability_cache[node]=0.;
        w.clear();
      } else if(getGateType(w[0])==BooleanGate::OR && getWires(w[0]).size()==0) {
        setGateType(node, BooleanGate::AND);
        probability_cache[node]=1.;
        w.clear();
      }
      break;

    case BooleanGate::MULIN:
    case BooleanGate::MULVAR:
    case BooleanGate::UNDETERMINED:
      throw CircuitException("Incorrect gate type");
      break;
    }
  }

  std::vector<bool> used(gates.size());
  std::stack<gate_t> to_process;
  to_process.push(root);

  while(!to_process.empty()) {
    auto g = to_process.top();
    to_process.pop();
    used[static_cast<size_t>(g)]=true;
    for(auto c: wires[static_cast<size_t>(g)])
      if(!used[static_cast<size_t>(c)])
        to_process.push(c);
  }

  size_t newi = 0;
  std::vector<gate_t> relabel(gates.size());
  for(size_t i=0; i<gates.size(); ++i)
  {
    if(!used[i]) {
      inputs.erase(gate_t{i});
      probability_cache.erase(gate_t{i});
      auto it = id2uuid.find(gate_t{i});
      if(it!=id2uuid.end()) {
        uuid2id.erase(it->second);
        id2uuid.erase(it);
      }
      continue;
    }

    relabel[i]=gate_t{newi};

    if(i!=newi) {
      gates[newi] = gates[i];
      wires[newi] = wires[i];
      prob[newi]=prob[i];

      auto it1 = probability_cache.find(gate_t{i});
      if(it1!=probability_cache.end()) {
        probability_cache[gate_t{newi}] = it1->second;
        probability_cache.erase(it1);
      }

      auto it2 = id2uuid.find(gate_t{i});
      if(it2!=id2uuid.end()) {
        id2uuid[gate_t{newi}] = it2->second;
        uuid2id[it2->second] = gate_t{newi};
        id2uuid.erase(it2);
      }

      if(root==gate_t{i})
        root=gate_t{newi};

      auto it3 = inputs.find(gate_t{i});
      if(it3!=inputs.end()) {
        inputs.insert(gate_t{newi});
        inputs.erase(it3);
      }
    }

    ++newi;
  }

  gates.resize(newi);
  wires.resize(newi);
  prob.resize(newi);

  for(auto &w: wires)
    for(size_t i=0; i<w.size(); ++i)
      w[i]=relabel[static_cast<size_t>(w[i])];
}

dDNNF::Stats dDNNF::nodeStats() const
{
  Stats s;

  // Reachable gates from the root (same traversal as toDot).
  std::set<gate_t> seen;
  {
    std::stack<gate_t> stk;
    stk.push(root);
    while(!stk.empty()) {
      gate_t g = stk.top();
      stk.pop();
      if(seen.count(g))
        continue;
      seen.insert(g);
      if(getGateType(g) == BooleanGate::UNDETERMINED)
        continue;
      for(auto c : getWires(g))
        stk.push(c);
    }
  }

  for(gate_t g : seen) {
    auto type = getGateType(g);
    if(type == BooleanGate::UNDETERMINED)
      continue;
    ++s.nodes;
    switch(type) {
    case BooleanGate::AND: ++s.and_gates; break;
    case BooleanGate::OR:  ++s.or_gates;  break;
    case BooleanGate::NOT: ++s.not_gates; break;
    case BooleanGate::IN:  ++s.inputs;    break;
    default: break;
    }
    for(gate_t c : getWires(g))
      if(seen.count(c))
        ++s.edges;
    // Smoothness: every OR gate's children must mention the same set of
    // variables.  Cheap enough for an introspection call (vars() is a
    // reachability set per child).
    if(type == BooleanGate::OR && s.smooth) {
      const auto &children = getWires(g);
      for(std::size_t i = 1; i < children.size() && s.smooth; ++i)
        if(vars(children[i]) != vars(children[0]))
          s.smooth = false;
    }
  }

  // Longest path (in gates) from the root over the reachable DAG,
  // memoised: depth(g) = 0 for a leaf, else 1 + max child depth.
  std::unordered_map<gate_t, int, hash_gate_t> depth;
  std::function<int(gate_t)> longest = [&](gate_t g) -> int {
    auto it = depth.find(g);
    if(it != depth.end())
      return it->second;
    int d = 0;
    for(gate_t c : getWires(g))
      if(seen.count(c))
        d = std::max(d, 1 + longest(c));
    depth[g] = d;
    return d;
  };
  s.depth = seen.empty() ? 0 : longest(root);

  return s;
}

std::string dDNNF::toNNF(
  const std::function<int(const std::string &)> &var_of_uuid) const
{
  // Emit nodes in post-order so a child's NNF index is always smaller
  // than its parent's; memoise gate -> assigned NNF node index.
  std::vector<std::string> lines;
  std::unordered_map<gate_t, int, hash_gate_t> idx;
  std::size_t edges = 0;
  int max_var = 0;

  // Variable index of an input gate: the caller-supplied numbering (by
  // original UUID) when available, else the d-DNNF's own gate id + 1.
  auto var_of = [&](gate_t g) -> int {
    if(var_of_uuid) {
      auto u = id2uuid.find(g);
      if(u != id2uuid.end()) {
        int v = var_of_uuid(u->second);
        if(v > 0)
          return v;
      }
    }
    return static_cast<int>(
      static_cast<std::underlying_type<gate_t>::type>(g)) + 1;
  };

  std::function<int(gate_t)> emit = [&](gate_t g) -> int {
    auto found = idx.find(g);
    if(found != idx.end())
      return found->second;

    std::string line;
    switch(getGateType(g)) {
    case BooleanGate::IN: {
      int v = var_of(g);
      max_var = std::max(max_var, v);
      line = "L " + std::to_string(v);
      break;
    }
    case BooleanGate::NOT: {
      // A literal: NOT must sit directly over an input in NNF. The IN
      // child is not emitted as a node of its own here -- the negative
      // literal is the leaf.
      gate_t c = *getWires(g).begin();
      if(getGateType(c) != BooleanGate::IN)
        throw CircuitException(
          "toNNF: NOT over a non-input gate (circuit not in negation normal form)");
      int v = var_of(c);
      max_var = std::max(max_var, v);
      line = "L -" + std::to_string(v);
      break;
    }
    case BooleanGate::AND: {
      std::vector<int> ch;
      for(gate_t c : getWires(g))
        ch.push_back(emit(c));
      edges += ch.size();
      line = "A " + std::to_string(ch.size());
      for(int i : ch)
        line += " " + std::to_string(i);
      break;
    }
    case BooleanGate::OR: {
      std::vector<int> ch;
      for(gate_t c : getWires(g))
        ch.push_back(emit(c));
      edges += ch.size();
      // j (decision variable) is 0: ProvSQL's OR gates do not record
      // which variable the branches split on.
      line = "O 0 " + std::to_string(ch.size());
      for(int i : ch)
        line += " " + std::to_string(i);
      break;
    }
    default:
      throw CircuitException("toNNF: unexpected gate type in d-DNNF");
    }

    int my = static_cast<int>(lines.size());
    lines.push_back(std::move(line));
    idx[g] = my;
    return my;
  };

  emit(root);

  std::ostringstream out;
  out << "nnf " << lines.size() << " " << edges << " " << max_var << "\n";
  for(const auto &l : lines)
    out << l << "\n";
  return out.str();
}

std::string dDNNF::toDot() const
{
  std::ostringstream out;
  out << "digraph dDNNF {\n"
      << "  graph [rankdir=TB];\n"
      << "  node [fontname=\"Helvetica\"];\n";

  // Collect reachable gates from the root in deterministic gate-id
  // order (std::set ordering) so the output is stable across runs.
  std::stack<gate_t> stk;
  std::set<gate_t> seen;
  stk.push(root);
  while(!stk.empty()) {
    gate_t g = stk.top();
    stk.pop();
    if(seen.count(g))
      continue;
    seen.insert(g);
    if(getGateType(g) == BooleanGate::UNDETERMINED)
      continue;
    for(auto c : getWires(g))
      stk.push(c);
  }

  for(auto g : seen) {
    auto type = getGateType(g);
    if(type == BooleanGate::UNDETERMINED)
      continue;

    auto id = static_cast<std::underlying_type<gate_t>::type>(g);
    std::string label;
    std::string shape = "circle";
    std::string fill;
    std::string tooltip;
    switch(type) {
    case BooleanGate::AND:
      label = "∧";  // ∧
      break;
    case BooleanGate::OR:
      label = "∨";  // ∨
      break;
    case BooleanGate::NOT:
      label = "¬";  // ¬
      break;
    case BooleanGate::IN: {
      shape = "box";
      fill = ", style=filled, fillcolor=\"#e6f0fa\"";
      std::string u;
      auto it = id2uuid.find(g);
      if(it != id2uuid.end())
        u = it->second;
      std::string sh = u.size() > 8 ? u.substr(0, 8) : u;
      std::ostringstream lab;
      lab << sh << "\\np=" << std::fixed << std::setprecision(3) << prob[id];
      label = lab.str();
      // Surface the full original-circuit UUID through a `tooltip`
      // attribute so the Studio renderer can look the source row up
      // via /api/leaf/<uuid> (same behaviour as input gates in the
      // provenance circuit). The short prefix on the label is purely
      // for the in-circle text.
      tooltip = u;
      break;
    }
    default:
      continue;
    }

    out << "  n" << id << " [label=\"" << label << "\", shape=" << shape
        << fill;
    if(!tooltip.empty())
      out << ", tooltip=\"" << tooltip << "\"";
    if(g == root)
      out << ", penwidth=2";
    out << "];\n";
  }

  for(auto g : seen) {
    auto type = getGateType(g);
    if(type == BooleanGate::IN || type == BooleanGate::UNDETERMINED)
      continue;
    auto id = static_cast<std::underlying_type<gate_t>::type>(g);
    for(auto c : getWires(g)) {
      if(!seen.count(c))
        continue;
      auto cid = static_cast<std::underlying_type<gate_t>::type>(c);
      out << "  n" << id << " -> n" << cid << ";\n";
    }
  }

  out << "}\n";
  return out.str();
}
