/*
 * Copyright © 2003-2026 Dynare Team
 *
 * This file is part of Dynare.
 *
 * Dynare is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Dynare is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Dynare.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <algorithm>
#include <cassert>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <iterator>
#include <ranges>
#include <string>

#include "DataTree.hh"

using namespace std::string_literals;

bool DataTree::no_commutativity = false;

void
DataTree::initConstants()
{
  Zero = AddNonNegativeConstant("0");
  One = AddNonNegativeConstant("1");
  Two = AddNonNegativeConstant("2");
  Three = AddNonNegativeConstant("3");

  MinusOne = AddUMinus(One);

  NaN = AddNonNegativeConstant("NaN");
  Infinity = AddNonNegativeConstant("Inf");
  MinusInfinity = AddUMinus(Infinity);

  Pi = AddNonNegativeConstant("3.141592653589793");
}

DataTree::DataTree(SymbolTable& symbol_table_arg, NumericalConstants& num_constants_arg,
                   ExternalFunctionsTable& external_functions_table_arg,
                   HeterogeneityTable& heterogeneity_table_arg, DatabaseTable& database_table_arg,
                   bool is_dynamic_arg) :
    symbol_table {symbol_table_arg},
    num_constants {num_constants_arg},
    external_functions_table {external_functions_table_arg},
    heterogeneity_table {heterogeneity_table_arg},
    database_table {database_table_arg},
    is_dynamic {is_dynamic_arg}
{
  initConstants();
}

DataTree::DataTree(const DataTree& d) :
    symbol_table {d.symbol_table},
    num_constants {d.num_constants},
    external_functions_table {d.external_functions_table},
    heterogeneity_table {d.heterogeneity_table},
    database_table {d.database_table},
    is_dynamic {d.is_dynamic},
    local_variables_vector {d.local_variables_vector}
{
  // Constants must be initialized first because they are used in some Add* methods
  initConstants();

  // See commment in DataTree::operator=() for the rationale
  for (int symb_id : d.local_variables_vector)
    local_variables_table[symb_id] = d.local_variables_table.at(symb_id)->clone(*this);

  for (const auto& it : d.node_list)
    it->clone(*this);

  assert(node_list.size() == d.node_list.size());
}

DataTree&
DataTree::operator=(const DataTree& d)
{
  assert(&symbol_table == &d.symbol_table);
  assert(&num_constants == &d.num_constants);
  assert(&external_functions_table == &d.external_functions_table);
  assert(&heterogeneity_table == &d.heterogeneity_table);
  assert(&database_table == &d.database_table);
  assert(is_dynamic == d.is_dynamic);

  num_const_node_map.clear();
  variable_node_map.clear();
  unary_op_node_map.clear();
  binary_op_node_map.clear();
  trinary_op_node_map.clear();
  external_function_node_map.clear();
  var_expectation_node_map.clear();
  pac_expectation_node_map.clear();
  pac_target_nonstationary_node_map.clear();
  first_deriv_external_function_node_map.clear();
  second_deriv_external_function_node_map.clear();

  node_list.clear();

  // Constants must be initialized first because they are used in some Add* methods
  initConstants();

  /* Model local variables must be next, because they can be evaluated in Add*
     methods when the model equations are added. They need to be cloned in
     order of appearance in the model block (hence with
     local_variables_vector), because if there is a model_local_variable statement
     the symbol IDs ordering may not be the right one (see dynare#1782) */
  for (int symb_id : d.local_variables_vector)
    local_variables_table[symb_id] = d.local_variables_table.at(symb_id)->clone(*this);

  for (const auto& it : d.node_list)
    it->clone(*this);

  assert(node_list.size() == d.node_list.size());

  local_variables_vector = d.local_variables_vector;

  return *this;
}

NumConstNode*
DataTree::AddNonNegativeConstant(const string& value)
{
  int id = num_constants.AddNonNegativeConstant(value);

  if (auto it = num_const_node_map.find(id); it != num_const_node_map.end())
    return it->second;

  auto sp = make_unique<NumConstNode>(*this, node_list.size(), id);
  auto p = sp.get();
  node_list.push_back(move(sp));
  num_const_node_map.emplace(id, p);
  return p;
}

VariableNode*
DataTree::AddVariable(int symb_id, int lag)
{
  if (lag != 0 && !is_dynamic)
    {
      cerr << "Leads/lags not authorized in this DataTree" << '\n';
      exit(EXIT_FAILURE);
    }

  if (auto it = variable_node_map.find({symb_id, lag}); it != variable_node_map.end())
    return it->second;

  auto sp = make_unique<VariableNode>(*this, node_list.size(), symb_id, lag);
  auto p = sp.get();
  node_list.push_back(move(sp));
  variable_node_map.try_emplace({symb_id, lag}, p);
  return p;
}

VariableNode*
DataTree::getVariable(int symb_id, int lag) const
{
  auto it = variable_node_map.find({symb_id, lag});
  if (it == variable_node_map.end())
    {
      cerr << "DataTree::getVariable: unknown variable node for symb_id=" << symb_id
           << " and lag=" << lag << '\n';
      exit(EXIT_FAILURE);
    }
  return it->second;
}

NamespaceQualifiedVariableNode*
DataTree::AddNamespaceQualifiedVariable(NamespaceQualifiedVariableNode::NamespaceType Namespace,
                                        int symb_id, int lag, int database_id,
                                        const variant<int, string>& learnt_in_period)
{
  if (auto it = namespace_qualified_variable_node_map.find(
          {Namespace, symb_id, lag, learnt_in_period, database_id});
      it != namespace_qualified_variable_node_map.end())
    return it->second;

  auto sp = make_unique<NamespaceQualifiedVariableNode>(*this, node_list.size(), Namespace, symb_id,
                                                        lag, database_id, learnt_in_period);
  auto p = sp.get();
  node_list.push_back(move(sp));
  namespace_qualified_variable_node_map.try_emplace(
      {Namespace, symb_id, lag, learnt_in_period, database_id}, p);
  return p;
}

set<string>
DataTree::ParamUsedWithLeadLag() const
{
  set<string> r;
  for (const auto& [symb_lag, expr] : variable_node_map)
    if (symbol_table.getType(symb_lag.first) == SymbolType::parameter && symb_lag.second != 0)
      r.insert(symbol_table.getName(symb_lag.first));
  return r;
}

expr_t
DataTree::AddPlus(expr_t iArg1, expr_t iArg2)
{
  if (iArg2 == Zero)
    return iArg1;

  if (iArg1 == Zero)
    return iArg2;

  // Simplify x+(-y) in x-y
  if (auto uarg2 = dynamic_cast<UnaryOpNode*>(iArg2);
      uarg2 && uarg2->op_code == UnaryOpcode::uminus)
    return AddMinus(iArg1, uarg2->arg);

  // Simplify (-x)+y in y-x
  if (auto uarg1 = dynamic_cast<UnaryOpNode*>(iArg1);
      uarg1 && uarg1->op_code == UnaryOpcode::uminus)
    return AddMinus(iArg2, uarg1->arg);

  // Simplify (x-y)+y in x
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::minus && barg1->arg2 == iArg2)
    return barg1->arg1;

  // Simplify y+(x-y) in x
  if (auto barg2 = dynamic_cast<BinaryOpNode*>(iArg2);
      barg2 && barg2->op_code == BinaryOpcode::minus && barg2->arg2 == iArg1)
    return barg2->arg1;

  // To treat commutativity of "+"
  // Nodes iArg1 and iArg2 are sorted by index
  if (iArg1->idx > iArg2->idx && !no_commutativity) // NOLINT(clang-analyzer-core.NullDereference)
    swap(iArg1, iArg2);
  return AddBinaryOp(iArg1, BinaryOpcode::plus, iArg2);
}

expr_t
DataTree::AddMinus(expr_t iArg1, expr_t iArg2)
{
  if (iArg2 == Zero)
    return iArg1;

  if (iArg1 == Zero)
    return AddUMinus(iArg2);

  if (iArg1 == iArg2)
    return Zero;

  // Simplify x-(-y) in x+y
  if (auto uarg2 = dynamic_cast<UnaryOpNode*>(iArg2);
      uarg2 && uarg2->op_code == UnaryOpcode::uminus)
    return AddPlus(iArg1, uarg2->arg);

  // Simplify (x+y)-y and (y+x)-y in x
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::plus)
    {
      if (barg1->arg2 == iArg2)
        return barg1->arg1;
      if (barg1->arg1 == iArg2)
        return barg1->arg2;
    }

  return AddBinaryOp(iArg1, BinaryOpcode::minus, iArg2);
}

expr_t
DataTree::AddUMinus(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  // Simplify -(-x) in x
  if (auto uarg = dynamic_cast<UnaryOpNode*>(iArg1); uarg && uarg->op_code == UnaryOpcode::uminus)
    return uarg->arg;

  return AddUnaryOp(UnaryOpcode::uminus, iArg1);
}

expr_t
DataTree::AddTimes(expr_t iArg1, expr_t iArg2)
{
  if (iArg1 == Zero || iArg2 == Zero)
    return Zero;

  if (iArg1 == One)
    return iArg2;

  if (iArg2 == One)
    return iArg1;

  if (iArg1 == MinusOne)
    return AddUMinus(iArg2);

  if (iArg2 == MinusOne)
    return AddUMinus(iArg1);

  // Simplify (x/y)*y in x
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::divide && barg1->arg2 == iArg2)
    return barg1->arg1;

  // Simplify y*(x/y) in x
  if (auto barg2 = dynamic_cast<BinaryOpNode*>(iArg2);
      barg2 && barg2->op_code == BinaryOpcode::divide && barg2->arg2 == iArg1)
    return barg2->arg1;

  // To treat commutativity of "*"
  // Nodes iArg1 and iArg2 are sorted by index
  if (iArg1->idx > iArg2->idx && !no_commutativity) // NOLINT(clang-analyzer-core.NullDereference)
    swap(iArg1, iArg2);
  return AddBinaryOp(iArg1, BinaryOpcode::times, iArg2);
}

expr_t
DataTree::AddDivide(expr_t iArg1, expr_t iArg2) noexcept(false)
{
  if (iArg2 == One)
    return iArg1;

  // This test should be before the next two, otherwise 0/0 won't be rejected
  if (iArg2 == Zero)
    {
      throw DivisionByZeroException(
          "Division by zero when forming ("s + iArg1->toString() + ")/(" + iArg2->toString()
          + "); denominator simplified to 0 (possibly after substituting a variable set to 0).");
    }

  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == iArg2)
    return One;

  // Simplify x/(1/y) in x*y
  if (auto barg2 = dynamic_cast<BinaryOpNode*>(iArg2);
      barg2 && barg2->op_code == BinaryOpcode::divide && barg2->arg1 == One)
    return AddTimes(iArg1, barg2->arg2);

  // Simplify (x*y)/y and (y*x)/y in x
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::times)
    {
      if (barg1->arg2 == iArg2)
        return barg1->arg1;
      if (barg1->arg1 == iArg2)
        return barg1->arg2;
    }

  return AddBinaryOp(iArg1, BinaryOpcode::divide, iArg2);
}

expr_t
DataTree::AddLess(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::less, iArg2);
}

expr_t
DataTree::AddGreater(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::greater, iArg2);
}

expr_t
DataTree::AddLessEqual(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::lessEqual, iArg2);
}

expr_t
DataTree::AddGreaterEqual(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::greaterEqual, iArg2);
}

expr_t
DataTree::AddEqualEqual(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::equalEqual, iArg2);
}

expr_t
DataTree::AddDifferent(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::different, iArg2);
}

expr_t
DataTree::AddPower(expr_t iArg1, expr_t iArg2)
{
  // This one comes first, because 0⁰=1
  if (iArg2 == Zero)
    return One;

  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == One)
    return One;

  if (iArg2 == One)
    return iArg1;

  return AddBinaryOp(iArg1, BinaryOpcode::power, iArg2);
}

expr_t
DataTree::AddPowerDeriv(expr_t iArg1, expr_t iArg2, int powerDerivOrder)
{
  assert(powerDerivOrder > 0);
  return AddBinaryOp(iArg1, BinaryOpcode::powerDeriv, iArg2, powerDerivOrder);
}

expr_t
DataTree::AddDiff(expr_t iArg1)
{
  if (iArg1->maxLead() > 0)
    // Issue preprocessor#21: always expand diffs with leads
    return AddMinus(iArg1, iArg1->decreaseLeadsLags(1));
  return AddUnaryOp(UnaryOpcode::diff, iArg1);
}

expr_t
DataTree::AddExp(expr_t iArg1)
{
  if (iArg1 == Zero)
    return One;

  return AddUnaryOp(UnaryOpcode::exp, iArg1);
}

expr_t
DataTree::AddLog(expr_t iArg1)
{
  if (iArg1 == One)
    return Zero;

  if (iArg1 == Zero)
    {
      cerr << "ERROR: log(0) not defined!" << '\n';
      exit(EXIT_FAILURE);
    }

  // Simplify log(1/x) in −log(x)
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::divide && barg1->arg1 == One)
    return AddUMinus(AddLog(barg1->arg2));

  return AddUnaryOp(UnaryOpcode::log, iArg1);
}

expr_t
DataTree::AddLog10(expr_t iArg1)
{
  if (iArg1 == One)
    return Zero;

  if (iArg1 == Zero)
    {
      cerr << "ERROR: log10(0) not defined!" << '\n';
      exit(EXIT_FAILURE);
    }

  // Simplify log₁₀(1/x) in −log₁₀(x)
  if (auto barg1 = dynamic_cast<BinaryOpNode*>(iArg1);
      barg1 && barg1->op_code == BinaryOpcode::divide && barg1->arg1 == One)
    return AddUMinus(AddLog10(barg1->arg2));

  return AddUnaryOp(UnaryOpcode::log10, iArg1);
}

expr_t
DataTree::AddCos(expr_t iArg1)
{
  if (iArg1 == Zero)
    return One;

  return AddUnaryOp(UnaryOpcode::cos, iArg1);
}

expr_t
DataTree::AddSin(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::sin, iArg1);
}

expr_t
DataTree::AddTan(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::tan, iArg1);
}

expr_t
DataTree::AddAcos(expr_t iArg1)
{
  if (iArg1 == One)
    return Zero;

  return AddUnaryOp(UnaryOpcode::acos, iArg1);
}

expr_t
DataTree::AddAsin(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::asin, iArg1);
}

expr_t
DataTree::AddAtan(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::atan, iArg1);
}

expr_t
DataTree::AddCosh(expr_t iArg1)
{
  if (iArg1 == Zero)
    return One;

  return AddUnaryOp(UnaryOpcode::cosh, iArg1);
}

expr_t
DataTree::AddSinh(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::sinh, iArg1);
}

expr_t
DataTree::AddTanh(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::tanh, iArg1);
}

expr_t
DataTree::AddAcosh(expr_t iArg1)
{
  if (iArg1 == One)
    return Zero;

  return AddUnaryOp(UnaryOpcode::acosh, iArg1);
}

expr_t
DataTree::AddAsinh(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::asinh, iArg1);
}

expr_t
DataTree::AddAtanh(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::atanh, iArg1);
}

expr_t
DataTree::AddSqrt(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == One)
    return One;

  return AddUnaryOp(UnaryOpcode::sqrt, iArg1);
}

expr_t
DataTree::AddCbrt(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == One)
    return One;

  return AddUnaryOp(UnaryOpcode::cbrt, iArg1);
}

expr_t
DataTree::AddAbs(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == One)
    return One;

  return AddUnaryOp(UnaryOpcode::abs, iArg1);
}

expr_t
DataTree::AddSign(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  if (iArg1 == One)
    return One;

  return AddUnaryOp(UnaryOpcode::sign, iArg1);
}

expr_t
DataTree::AddErf(expr_t iArg1)
{
  if (iArg1 == Zero)
    return Zero;

  return AddUnaryOp(UnaryOpcode::erf, iArg1);
}

expr_t
DataTree::AddErfc(expr_t iArg1)
{
  if (iArg1 == Zero)
    return One;

  return AddUnaryOp(UnaryOpcode::erfc, iArg1);
}

expr_t
DataTree::AddMax(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::max, iArg2);
}

expr_t
DataTree::AddMin(expr_t iArg1, expr_t iArg2)
{
  return AddBinaryOp(iArg1, BinaryOpcode::min, iArg2);
}

expr_t
DataTree::AddNormcdf(expr_t iArg1, expr_t iArg2, expr_t iArg3)
{
  return AddTrinaryOp(iArg1, TrinaryOpcode::normcdf, iArg2, iArg3);
}

expr_t
DataTree::AddNormpdf(expr_t iArg1, expr_t iArg2, expr_t iArg3)
{
  return AddTrinaryOp(iArg1, TrinaryOpcode::normpdf, iArg2, iArg3);
}

expr_t
DataTree::AddSteadyState(expr_t iArg1)
{
  return AddUnaryOp(UnaryOpcode::steadyState, iArg1);
}

expr_t
DataTree::AddSteadyStateParamDeriv(expr_t iArg1, int param_symb_id)
{
  return AddUnaryOp(UnaryOpcode::steadyStateParamDeriv, iArg1, 0, param_symb_id);
}

expr_t
DataTree::AddSteadyStateParam2ndDeriv(expr_t iArg1, int param1_symb_id, int param2_symb_id)
{
  return AddUnaryOp(UnaryOpcode::steadyStateParam2ndDeriv, iArg1, 0, param1_symb_id,
                    param2_symb_id);
}

expr_t
DataTree::AddExpectation(int iArg1, expr_t iArg2)
{
  return AddUnaryOp(UnaryOpcode::expectation, iArg2, iArg1);
}

expr_t
DataTree::AddVarExpectation(const string& model_name)
{
  if (auto it = var_expectation_node_map.find(model_name); it != var_expectation_node_map.end())
    return it->second;

  auto sp = make_unique<VarExpectationNode>(*this, node_list.size(), model_name);
  auto p = sp.get();
  node_list.push_back(move(sp));
  var_expectation_node_map.emplace(model_name, p);
  return p;
}

expr_t
DataTree::AddPacExpectation(const string& model_name)
{
  if (auto it = pac_expectation_node_map.find(model_name); it != pac_expectation_node_map.end())
    return it->second;

  auto sp = make_unique<PacExpectationNode>(*this, node_list.size(), model_name);
  auto p = sp.get();
  node_list.push_back(move(sp));
  pac_expectation_node_map.emplace(model_name, p);
  return p;
}

expr_t
DataTree::AddPacTargetNonstationary(const string& model_name)
{
  if (auto it = pac_target_nonstationary_node_map.find(model_name);
      it != pac_target_nonstationary_node_map.end())
    return it->second;

  auto sp = make_unique<PacTargetNonstationaryNode>(*this, node_list.size(), model_name);
  auto p = sp.get();
  node_list.push_back(move(sp));
  pac_target_nonstationary_node_map.emplace(model_name, p);
  return p;
}

BinaryOpNode*
DataTree::AddEqual(expr_t iArg1, expr_t iArg2)
{
  /* We know that we can safely cast to BinaryOpNode because
     BinaryOpCode::equal can never be reduced to a constant. */
  return dynamic_cast<BinaryOpNode*>(AddBinaryOp(iArg1, BinaryOpcode::equal, iArg2));
}

void
DataTree::AddLocalVariable(int symb_id, expr_t value) noexcept(false)
{
  assert(symbol_table.getType(symb_id) == SymbolType::modelLocalVariable);

  // Throw an exception if symbol already declared
  if (local_variables_table.contains(symb_id))
    throw LocalVariableException {symbol_table.getName(symb_id)};

  local_variables_table.emplace(symb_id, value);
  local_variables_vector.push_back(symb_id);
}

expr_t
DataTree::AddExternalFunction(int symb_id, vector<expr_t> arguments)
{
  assert(symbol_table.getType(symb_id) == SymbolType::externalFunction);

  if (auto it = external_function_node_map.find({arguments, symb_id});
      it != external_function_node_map.end())
    return it->second;

  auto sp = make_unique<ExternalFunctionNode>(*this, node_list.size(), symb_id, arguments);
  auto p = sp.get();
  node_list.push_back(move(sp));
  external_function_node_map.try_emplace({move(arguments), symb_id}, p);
  return p;
}

expr_t
DataTree::AddFirstDerivExternalFunction(int top_level_symb_id, const vector<expr_t>& arguments,
                                        int input_index)
{
  assert(symbol_table.getType(top_level_symb_id) == SymbolType::externalFunction);

  if (auto it
      = first_deriv_external_function_node_map.find({arguments, input_index, top_level_symb_id});
      it != first_deriv_external_function_node_map.end())
    return it->second;

  auto sp = make_unique<FirstDerivExternalFunctionNode>(*this, node_list.size(), top_level_symb_id,
                                                        arguments, input_index);
  auto p = sp.get();
  node_list.push_back(move(sp));
  first_deriv_external_function_node_map.try_emplace({arguments, input_index, top_level_symb_id},
                                                     p);
  return p;
}

expr_t
DataTree::AddSecondDerivExternalFunction(int top_level_symb_id, const vector<expr_t>& arguments,
                                         int input_index1, int input_index2)
{
  assert(symbol_table.getType(top_level_symb_id) == SymbolType::externalFunction);

  if (auto it = second_deriv_external_function_node_map.find(
          {arguments, input_index1, input_index2, top_level_symb_id});
      it != second_deriv_external_function_node_map.end())
    return it->second;

  auto sp = make_unique<SecondDerivExternalFunctionNode>(*this, node_list.size(), top_level_symb_id,
                                                         arguments, input_index1, input_index2);
  auto p = sp.get();
  node_list.push_back(move(sp));
  second_deriv_external_function_node_map.try_emplace(
      {arguments, input_index1, input_index2, top_level_symb_id}, p);
  return p;
}

expr_t
DataTree::AddSum(expr_t arg)
{
  return AddUnaryOp(UnaryOpcode::sum, arg);
}

bool
DataTree::isSymbolUsed(int symb_id) const
{
  if (ranges::any_of(views::keys(variable_node_map),
                     [=](const auto& symb_lag) { return symb_lag.first == symb_id; }))
    return true;

  if (local_variables_table.contains(symb_id))
    return true;

  return false;
}

int
DataTree::getDerivID([[maybe_unused]] int symb_id, [[maybe_unused]] int lag) const noexcept(false)
{
  throw UnknownDerivIDException();
}

SymbolType
DataTree::getTypeByDerivID([[maybe_unused]] int deriv_id) const noexcept(false)
{
  throw UnknownDerivIDException();
}

int
DataTree::getLagByDerivID([[maybe_unused]] int deriv_id) const noexcept(false)
{
  throw UnknownDerivIDException();
}

int
DataTree::getSymbIDByDerivID([[maybe_unused]] int deriv_id) const noexcept(false)
{
  throw UnknownDerivIDException();
}

int
DataTree::getTypeSpecificIDByDerivID([[maybe_unused]] int deriv_id) const
{
  throw UnknownDerivIDException();
}

void
DataTree::addAllParamDerivId([[maybe_unused]] set<int>& deriv_id_set)
{
}

bool
DataTree::isUnaryOpUsed(UnaryOpcode opcode) const
{
  return ranges::any_of(views::keys(unary_op_node_map),
                        [=](const auto& key) { return get<1>(key) == opcode; });
}

bool
DataTree::isUnaryOpUsedOnType(SymbolType type, UnaryOpcode opcode) const
{
  set<int> var;
  for (const auto& [key, value] : unary_op_node_map)
    if (get<1>(key) == opcode)
      {
        value->collectVariables(type, var);
        if (!var.empty())
          return true;
      }
  return false;
}

bool
DataTree::isBinaryOpUsed(BinaryOpcode opcode) const
{
  return ranges::any_of(views::keys(binary_op_node_map),
                        [=](const auto& key) { return get<2>(key) == opcode; });
}

bool
DataTree::isBinaryOpUsedOnType(SymbolType type, BinaryOpcode opcode) const
{
  set<int> var;
  for (const auto& [key, value] : binary_op_node_map)
    if (get<2>(key) == opcode)
      {
        value->collectVariables(type, var);
        if (!var.empty())
          return true;
      }
  return false;
}

int
DataTree::minLagForSymbol(int symb_id) const
{
  int r = 0;
  for (const auto& [symb_lag, expr] : variable_node_map)
    if (symb_lag.first == symb_id)
      r = min(r, symb_lag.second);
  return r;
}

void
DataTree::writeCHelpersDefinition(ostream& output) const
{
  if (isBinaryOpUsed(BinaryOpcode::powerDeriv))
    output << "// The k-th derivative of x^p" << '\n'
           << "inline double" << '\n'
           << "getPowerDeriv(double x, double p, int k)" << '\n'
           << "{" << '\n'
           << "  if (fabs(x) < " << power_deriv_near_zero
           << " && p >= 0 && k > p && fabs(p-nearbyint(p)) < " << power_deriv_near_zero << ')'
           << '\n'
           << "    return 0.0;" << '\n'
           << "  else" << '\n'
           << "    {" << '\n'
           << "      double dxp = pow(x, p-k);" << '\n'
           << "      for (int i = 0; i<k; i++)" << '\n'
           << "        dxp *= p--;" << '\n'
           << "      return dxp;" << '\n'
           << "    }" << '\n'
           << "}" << '\n';

  if (isUnaryOpUsed(UnaryOpcode::sign))
    output << "inline double" << '\n'
           << "sign(double x)" << '\n'
           << "{" << '\n'
           << "  return (x > 0) ? 1 : ((x < 0) ? -1 : 0);" << '\n'
           << "}" << '\n';
}

void
DataTree::writeCHelpersDeclaration(ostream& output) const
{
  if (isBinaryOpUsed(BinaryOpcode::powerDeriv))
    output << "extern inline double getPowerDeriv(double x, double p, int k);" << '\n';
  if (isUnaryOpUsed(UnaryOpcode::sign))
    output << "extern inline double sign(double x);" << '\n';
}

string
DataTree::writePythonHelpersImports(bool use_jax, bool use_numba)
{
  string imports = "import numpy as np\n"
                   "from numpy import log, exp, sqrt\n"
                   "from scipy.stats import norm\n"
                   "from typing import Tuple, Optional\n";

  if (use_jax)
    imports += "import jax\n"
               "import jax.numpy as jnp\n"
               "from jax import jit\n";

  if (use_numba)
    imports += "from numba import jit, prange\n";

  return imports;
}

string
DataTree::writePythonHelpersDecorator(bool use_jax, bool use_numba)
{
  if (use_jax)
    return "@jit\n";
  else if (use_numba)
    return "@jit(nopython=True, cache=True)\n";
  else
    return "";
}

string
DataTree::writePythonHelpersDefinition()
{
  return "# Helper functions for derivatives\n"
         "def getPowerDeriv(x, p, k):\n"
         "    \"\"\"Compute the k-th derivative of x^p\"\"\"\n"
         "    if k == 1:\n"
         "        if abs(p) < 1e-12:\n"
         "            return 0.0\n"
         "        else:\n"
         "            return p * x**(p - 1)\n"
         "    elif k == 2:\n"
         "        if abs(p) < 1e-12:\n"
         "            return 0.0\n"
         "        else:\n"
         "            return p * (p - 1) * x**(p - 2)\n"
         "    else:\n"
         "        # General case for k-th derivative\n"
         "        result = 1.0\n"
         "        for i in range(k):\n"
         "            result *= (p - i)\n"
         "        return result * x**(p - k)\n\n";
}

string
DataTree::writePythonHelpersImportFromInit()
{
  return "from . import getPowerDeriv\n";
}
