/*
 * 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 <array>
#include <cassert>
#include <cstdlib>
#include <iostream>
#include <ranges>
#include <utility>

#include "Shocks.hh"
#include "Utils.hh"

static auto print_matlab_period_range = []<class T>(ostream& output, const T& arg) {
  if constexpr (is_same_v<T, pair<int, int>>)
    output << arg.first << ":" << arg.second;
  else if constexpr (is_same_v<T, pair<string, string>>)
    output << "(" << arg.first << "):(" << arg.second << ")";
  else
    static_assert(always_false_v<T>, "Non-exhaustive visitor!");
};

static auto print_json_period_range = []<class T>(ostream& output, const T& arg) {
  if constexpr (is_same_v<T, pair<int, int>>)
    output << R"("period1": )" << arg.first << ", "
           << R"("period2": )" << arg.second;
  else if constexpr (is_same_v<T, pair<string, string>>)
    output << R"("period1": ")" << arg.first << R"(", )"
           << R"("period2": ")" << arg.second << '"';
  else
    static_assert(always_false_v<T>, "Non-exhaustive visitor!");
};

static auto print_matlab_learnt_in = [](ostream& output, const auto& p) { output << p; };

static auto print_json_learnt_in = []<class T>(ostream& output, const T& p) {
  if constexpr (is_same_v<T, int>)
    output << p;
  else if constexpr (is_same_v<T, string>)
    output << '"' << p << '"';
  else
    static_assert(always_false_v<T>, "Non-exhaustive visitor!");
};

AbstractShocksStatement::AbstractShocksStatement(bool overwrite_arg, ShockType type_arg,
                                                 det_shocks_t det_shocks_arg,
                                                 const SymbolTable& symbol_table_arg) :
    overwrite {overwrite_arg},
    type {type_arg},
    det_shocks {move(det_shocks_arg)},
    symbol_table {symbol_table_arg}
{
}

void
AbstractShocksStatement::writeDetShocks(ostream& output) const
{
  if (det_shocks.empty())
    return;

  output << "M_.det_shocks = [ M_.det_shocks;" << '\n';
  for (const auto& [id, shock_vec] : det_shocks)
    for (bool exo_det = (symbol_table.getType(id) == SymbolType::exogenousDet);
         const auto& [period_range, value] : shock_vec)
      {
        output << "struct('exo_det'," << boolalpha << exo_det << ",'exo_id',"
               << symbol_table.getTypeSpecificID(id) + 1 << ",'type','" << typeToString(type) << "'"
               << ",'periods',";
        visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
        output << ",'value',";
        value->writeOutput(output);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
AbstractShocksStatement::writeJsonDetShocks(ostream& output) const
{
  output << R"("deterministic_shocks": [)";
  for (bool printed_something {false}; const auto& [id, shock_vec] : det_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : shock_vec)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << "]";
}

string
AbstractShocksStatement::typeToString(ShockType type)
{
  switch (type)
    {
    case ShockType::level:
      return "level";
    case ShockType::multiplySteadyState:
      return "multiply_steady_state";
    case ShockType::multiplyInitialSteadyState:
      return "multiply_initial_steady_state";
    }
  __builtin_unreachable(); // Silence GCC warning
}

ShocksStatement::ShocksStatement(bool overwrite_arg, det_shocks_t det_shocks_arg,
                                 var_and_std_shocks_t var_shocks_arg,
                                 var_and_std_shocks_t std_shocks_arg,
                                 covar_and_corr_shocks_t covar_shocks_arg,
                                 covar_and_corr_shocks_t corr_shocks_arg,
                                 skew_shocks_t skew_shocks_arg,
                                 const SymbolTable& symbol_table_arg) :
    AbstractShocksStatement {overwrite_arg, ShockType::level, move(det_shocks_arg),
                             symbol_table_arg},
    var_shocks {move(var_shocks_arg)},
    std_shocks {move(std_shocks_arg)},
    covar_shocks {move(covar_shocks_arg)},
    corr_shocks {move(corr_shocks_arg)},
    skew_shocks {move(skew_shocks_arg)}
{
}

void
ShocksStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                             [[maybe_unused]] bool minimal_workspace) const
{
  output << "%" << '\n' << "% SHOCKS instructions" << '\n' << "%" << '\n';

  if (overwrite)
    {
      output << "M_.det_shocks = struct([]);" << '\n';

      output << "M_.Sigma_e = zeros(" << symbol_table.exo_nbr() << ", " << symbol_table.exo_nbr()
             << ");" << '\n'
             << "M_.Correlation_matrix = eye(" << symbol_table.exo_nbr() << ", "
             << symbol_table.exo_nbr() << ");" << '\n';

      if (has_calibrated_measurement_errors())
        output << "M_.H = zeros(" << symbol_table.observedVariablesNbr() << ", "
               << symbol_table.observedVariablesNbr() << ");" << '\n'
               << "M_.Correlation_matrix_ME = eye(" << symbol_table.observedVariablesNbr() << ", "
               << symbol_table.observedVariablesNbr() << ");" << '\n';
      else
        output << "M_.H = 0;" << '\n' << "M_.Correlation_matrix_ME = 1;" << '\n';
    }

  writeDetShocks(output);
  writeVarAndStdShocks(output);
  writeCovarAndCorrShocks(output);
  writeSkewShocks(output);

  /* M_.sigma_e_is_diagonal is initialized to 1 by ModFile.cc.
     If there are no off-diagonal elements, and we are not in overwrite mode,
     then we don't reset it to 1, since there might be previous shocks blocks
     with off-diagonal elements. */
  if (covar_shocks.size() + corr_shocks.size() > 0)
    output << "M_.sigma_e_is_diagonal = 0;" << '\n';
  else if (overwrite)
    output << "M_.sigma_e_is_diagonal = 1;" << '\n';
}

void
ShocksStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "shocks")"
         << R"(, "overwrite": )" << boolalpha << overwrite;
  if (!det_shocks.empty())
    {
      output << ", ";
      writeJsonDetShocks(output);
    }
  output << R"(, "variance": [)";
  for (bool printed_something {false}; auto& [id, value] : var_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"name": ")" << symbol_table.getName(id) << R"(", )"
             << R"("variance": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "stderr": [)";
  for (bool printed_something {false}; auto& [id, value] : std_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"name": ")" << symbol_table.getName(id) << R"(", )"
             << R"("stderr": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "covariance": [)";
  for (bool printed_something {false}; auto& [ids, value] : covar_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << "{"
             << R"("name": ")" << symbol_table.getName(ids.first) << R"(", )"
             << R"("name2": ")" << symbol_table.getName(ids.second) << R"(", )"
             << R"("covariance": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "correlation": [)";
  for (bool printed_something {false}; auto& [ids, value] : corr_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << "{"
             << R"("name": ")" << symbol_table.getName(ids.first) << R"(", )"
             << R"("name2": ")" << symbol_table.getName(ids.second) << R"(", )"
             << R"("correlation": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "skewness": [)";
  for (bool printed_something {false}; auto& [ids, value] : skew_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      const auto& [id1, id2, id3] = ids;
      output << "{"
             << R"("name": ")" << symbol_table.getName(id1) << R"(", )"
             << R"("name2": ")" << symbol_table.getName(id2) << R"(", )"
             << R"("name3": ")" << symbol_table.getName(id3) << R"(", )"
             << R"("skewness": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << "}";
}

void
ShocksStatement::writeVarOrStdShock(ostream& output, const pair<int, expr_t>& it, bool stddev) const
{
  SymbolType type = symbol_table.getType(it.first);
  assert(type == SymbolType::exogenous || symbol_table.isObservedVariable(it.first));

  int id;
  if (type == SymbolType::exogenous)
    {
      output << "M_.Sigma_e(";
      id = symbol_table.getTypeSpecificID(it.first) + 1;
    }
  else
    {
      output << "M_.H(";
      id = symbol_table.getObservedVariableIndex(it.first) + 1;
    }

  output << id << ", " << id << ") = ";
  if (stddev)
    output << "(";
  it.second->writeOutput(output);
  if (stddev)
    output << ")^2";
  output << ";" << '\n';
}

void
ShocksStatement::writeVarAndStdShocks(ostream& output) const
{
  for (const auto& it : var_shocks)
    writeVarOrStdShock(output, it, false);

  for (const auto& it : std_shocks)
    writeVarOrStdShock(output, it, true);
}

void
ShocksStatement::writeCovarOrCorrShock(ostream& output, const pair<pair<int, int>, expr_t>& it,
                                       bool corr) const
{
  SymbolType type1 = symbol_table.getType(it.first.first);
  SymbolType type2 = symbol_table.getType(it.first.second);
  assert((type1 == SymbolType::exogenous && type2 == SymbolType::exogenous)
         || (symbol_table.isObservedVariable(it.first.first)
             && symbol_table.isObservedVariable(it.first.second)));
  string matrix, corr_matrix;
  int id1, id2;
  if (type1 == SymbolType::exogenous)
    {
      matrix = "M_.Sigma_e";
      corr_matrix = "M_.Correlation_matrix";
      id1 = symbol_table.getTypeSpecificID(it.first.first) + 1;
      id2 = symbol_table.getTypeSpecificID(it.first.second) + 1;
    }
  else
    {
      matrix = "M_.H";
      corr_matrix = "M_.Correlation_matrix_ME";
      id1 = symbol_table.getObservedVariableIndex(it.first.first) + 1;
      id2 = symbol_table.getObservedVariableIndex(it.first.second) + 1;
    }

  output << matrix << "(" << id1 << ", " << id2 << ") = ";
  it.second->writeOutput(output);
  if (corr)
    output << "*sqrt(" << matrix << "(" << id1 << ", " << id1 << ")*" << matrix << "(" << id2
           << ", " << id2 << "))";
  output << ";" << '\n'
         << matrix << "(" << id2 << ", " << id1 << ") = " << matrix << "(" << id1 << ", " << id2
         << ");" << '\n';

  if (corr)
    {
      output << corr_matrix << "(" << id1 << ", " << id2 << ") = ";
      it.second->writeOutput(output);
      output << ";" << '\n'
             << corr_matrix << "(" << id2 << ", " << id1 << ") = " << corr_matrix << "(" << id1
             << ", " << id2 << ");" << '\n';
    }
}

void
ShocksStatement::writeCovarAndCorrShocks(ostream& output) const
{
  for (const auto& it : covar_shocks)
    writeCovarOrCorrShock(output, it, false);

  for (const auto& it : corr_shocks)
    writeCovarOrCorrShock(output, it, true);
}

void
ShocksStatement::writeSkewShock(ostream& output, const pair<tuple<int, int, int>, expr_t>& it) const
{
  const auto& [id1, id2, id3] = it.first;

  // All three must be exogenous (already checked in checkPass)
  int tsid1 = symbol_table.getTypeSpecificID(id1) + 1;
  int tsid2 = symbol_table.getTypeSpecificID(id2) + 1;
  int tsid3 = symbol_table.getTypeSpecificID(id3) + 1;

  // Store in sorted (increasing) order, like model derivative matrices
  array idx {tsid1, tsid2, tsid3};
  ranges::sort(idx);

  // Remove existing row for this sorted triple (overwrite semantics)
  output << "M_.Skew_e(M_.Skew_e(:,1)==" << idx[0] << " & M_.Skew_e(:,2)==" << idx[1]
         << " & M_.Skew_e(:,3)==" << idx[2] << ",:) = [];" << '\n';

  // Append single row [i, j, k, value] with i <= j <= k
  output << "M_.Skew_e = [M_.Skew_e; " << idx[0] << ", " << idx[1] << ", " << idx[2] << ", ";
  it.second->writeOutput(output);
  output << "];" << '\n';
}

void
ShocksStatement::writeSkewShocks(ostream& output) const
{
  for (const auto& it : skew_shocks)
    writeSkewShock(output, it);
  // Remove any zero-valued entries from sparse Skew_e (value in column 4)
  if (!skew_shocks.empty())
    output << "M_.Skew_e(M_.Skew_e(:,4)==0,:) = [];" << '\n';
}

void
ShocksStatement::checkPass(ModFileStructure& mod_file_struct,
                           [[maybe_unused]] WarningConsolidation& warnings)
{
  mod_file_struct.shocks_present = true;

  /* Error out if variables are not of the right type. This must be done here
     and not at parsing time (see #448).
     Also Determine if there is a calibrated measurement error */
  for (const auto& [id, val] : var_shocks)
    {
      if (symbol_table.getType(id) != SymbolType::exogenous && !symbol_table.isObservedVariable(id))
        {
          cerr << "shocks: setting a variance on '" << symbol_table.getName(id)
               << "' is not allowed, because it is neither an exogenous variable nor an observed "
                  "endogenous variable"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  for (const auto& [id, val] : std_shocks)
    {
      if (symbol_table.getType(id) != SymbolType::exogenous && !symbol_table.isObservedVariable(id))
        {
          cerr << "shocks: setting a standard error on '" << symbol_table.getName(id)
               << "' is not allowed, because it is neither an exogenous variable nor an observed "
                  "endogenous variable"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  for (const auto& [ids, val] : covar_shocks)
    {
      auto& [symb_id1, symb_id2] = ids;

      if (!((symbol_table.getType(symb_id1) == SymbolType::exogenous
             && symbol_table.getType(symb_id2) == SymbolType::exogenous)
            || (symbol_table.isObservedVariable(symb_id1)
                && symbol_table.isObservedVariable(symb_id2))))
        {
          cerr << "shocks: setting a covariance between '" << symbol_table.getName(symb_id1)
               << "' and '" << symbol_table.getName(symb_id2)
               << "'is not allowed; covariances can only be specified for exogenous or observed "
                  "endogenous variables of same type"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  for (const auto& [ids, val] : corr_shocks)
    {
      auto& [symb_id1, symb_id2] = ids;

      if (!((symbol_table.getType(symb_id1) == SymbolType::exogenous
             && symbol_table.getType(symb_id2) == SymbolType::exogenous)
            || (symbol_table.isObservedVariable(symb_id1)
                && symbol_table.isObservedVariable(symb_id2))))
        {
          cerr << "shocks: setting a correlation between '" << symbol_table.getName(symb_id1)
               << "' and '" << symbol_table.getName(symb_id2)
               << "'is not allowed; correlations can only be specified for exogenous or observed "
                  "endogenous variables of same type"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  for (const auto& [ids, val] : skew_shocks)
    {
      const auto& [symb_id1, symb_id2, symb_id3] = ids;

      if (!(symbol_table.getType(symb_id1) == SymbolType::exogenous
            && symbol_table.getType(symb_id2) == SymbolType::exogenous
            && symbol_table.getType(symb_id3) == SymbolType::exogenous))
        {
          cerr << "shocks: setting skewness for '" << symbol_table.getName(symb_id1) << "', '"
               << symbol_table.getName(symb_id2) << "', '" << symbol_table.getName(symb_id3)
               << "' is not allowed; skewness can only be specified for exogenous variables"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  // Determine if there is a calibrated measurement error
  mod_file_struct.calibrated_measurement_errors |= has_calibrated_measurement_errors();

  // Fill in mod_file_struct.parameters_with_shocks_values (related to #469)
  for (const auto& [id, val] : var_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [id, val] : std_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [ids, val] : covar_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [ids, val] : corr_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [ids, val] : skew_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
}

bool
ShocksStatement::has_calibrated_measurement_errors() const
{
  for (const auto& [id, val] : var_shocks)
    if (symbol_table.isObservedVariable(id))
      return true;

  for (const auto& [id, val] : std_shocks)
    if (symbol_table.isObservedVariable(id))
      return true;

  for (const auto& [ids, val] : covar_shocks)
    if (symbol_table.isObservedVariable(ids.first) || symbol_table.isObservedVariable(ids.second))
      return true;

  for (const auto& [ids, val] : corr_shocks)
    if (symbol_table.isObservedVariable(ids.first) || symbol_table.isObservedVariable(ids.second))
      return true;

  return false;
}

MShocksStatement::MShocksStatement(bool overwrite_arg, bool relative_to_initval_arg,
                                   det_shocks_t det_shocks_arg,
                                   const SymbolTable& symbol_table_arg) :
    AbstractShocksStatement {overwrite_arg,
                             relative_to_initval_arg ? ShockType::multiplyInitialSteadyState
                                                     : ShockType::multiplySteadyState,
                             move(det_shocks_arg), symbol_table_arg},
    relative_to_initval {relative_to_initval_arg}
{
}

void
MShocksStatement::checkPass(ModFileStructure& mod_file_struct,
                            [[maybe_unused]] WarningConsolidation& warnings)
{
  mod_file_struct.mshocks_present = true;
}

void
MShocksStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                              [[maybe_unused]] bool minimal_workspace) const
{
  output << "%" << '\n' << "% MSHOCKS instructions" << '\n' << "%" << '\n';

  if (overwrite)
    output << "M_.det_shocks = struct([]);" << '\n';

  writeDetShocks(output);
}

void
MShocksStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "mshocks")"
         << R"(, "overwrite": )" << boolalpha << overwrite << R"(, "relative_to_initval": )"
         << boolalpha << relative_to_initval;
  if (!det_shocks.empty())
    {
      output << ", ";
      writeJsonDetShocks(output);
    }
  output << "}";
}

ShocksSurpriseStatement::ShocksSurpriseStatement(
    bool overwrite_arg, AbstractShocksStatement::det_shocks_t surprise_shocks_arg,
    const SymbolTable& symbol_table_arg) :
    overwrite {overwrite_arg},
    surprise_shocks {move(surprise_shocks_arg)},
    symbol_table {symbol_table_arg}
{
}

void
ShocksSurpriseStatement::checkPass(ModFileStructure& mod_file_struct,
                                   [[maybe_unused]] WarningConsolidation& warnings)
{
  mod_file_struct.shocks_surprise_present = true;
}

void
ShocksSurpriseStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                     [[maybe_unused]] bool minimal_workspace) const
{
  if (overwrite)
    output << "M_.surprise_shocks = [" << '\n';
  else
    output << "M_.surprise_shocks = [ M_.surprise_shocks;" << '\n';
  for (const auto& [id, shock_vec] : surprise_shocks)
    for (const auto& [period_range, value] : shock_vec)
      {
        auto [period1, period2] = get<pair<int, int>>(period_range);
        output << "struct('exo_id'," << symbol_table.getTypeSpecificID(id) + 1 << ",'periods',"
               << period1 << ":" << period2 << ",'value',";
        value->writeOutput(output);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
ShocksSurpriseStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "shocks")"
         << R"(, "surprise": true)"
         << R"(, "surprise_shocks": [)";
  for (bool printed_something {false}; const auto& [id, shock_vec] : surprise_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : shock_vec)
        {
          auto [period1, period2] = get<pair<int, int>>(period_range);
          if (exchange(printed_something2, true))
            output << ", ";
          output << R"({"period1": )" << period1 << ", "
                 << R"("period2": )" << period2 << ", "
                 << R"("value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << "]}";
}

ShocksLearntInStatement::ShocksLearntInStatement(variant<int, string> learnt_in_period_arg,
                                                 bool overwrite_arg,
                                                 learnt_shocks_t learnt_shocks_arg,
                                                 const SymbolTable& symbol_table_arg) :
    learnt_in_period {move(learnt_in_period_arg)},
    overwrite {overwrite_arg},
    learnt_shocks {move(learnt_shocks_arg)},
    symbol_table {symbol_table_arg}
{
}

void
ShocksLearntInStatement::checkPass(ModFileStructure& mod_file_struct,
                                   [[maybe_unused]] WarningConsolidation& warnings)
{
  mod_file_struct.shocks_learnt_in_present = true;
}

string
ShocksLearntInStatement::typeToString(LearntShockType type)
{
  switch (type)
    {
    case LearntShockType::level:
      return "level";
    case LearntShockType::add:
      return "add";
    case LearntShockType::multiply:
      return "multiply";
    case LearntShockType::multiplySteadyState:
      return "multiply_steady_state";
    case LearntShockType::multiplyInitialSteadyState:
      return "multiply_initial_steady_state";
    }
  __builtin_unreachable(); // Silence GCC warning
}

void
ShocksLearntInStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                     [[maybe_unused]] bool minimal_workspace) const
{
  if (overwrite)
    {
      output << "if ~isempty(M_.learnt_shocks)" << '\n'
             << "  M_.learnt_shocks = M_.learnt_shocks(cellfun(@(x) ~isa(x, '";
      if (holds_alternative<int>(learnt_in_period))
        output << "numeric";
      else
        output << "dates";
      output << "') || x ~= ";
      /* NB: date expression not parenthesized since it can only contain a + operator, which has
         higher precedence than ~= and || */
      visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
      output << ", {M_.learnt_shocks.learnt_in}));" << '\n' << "end" << '\n';
    }

  output << "M_.learnt_shocks = [ M_.learnt_shocks;" << '\n';
  for (const auto& [id, shock_vec] : learnt_shocks)
    for (const auto& [type, period_range, value] : shock_vec)
      {
        output << "struct('learnt_in',";
        visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
        output << ",'exo_id'," << symbol_table.getTypeSpecificID(id) + 1 << ",'periods',";
        visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
        output << ",'type','" << typeToString(type) << "'"
               << ",'value',";
        value->writeOutput(output);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
ShocksLearntInStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "shocks")"
         << R"(, "learnt_in": )";
  visit([&](const auto& p) { print_json_learnt_in(output, p); }, learnt_in_period);
  output << R"(, "overwrite": )" << boolalpha << overwrite << R"(, "learnt_shocks": [)";
  for (bool printed_something {false}; const auto& [id, shock_vec] : learnt_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [type, period_range, value] : shock_vec)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "type": ")" << typeToString(type) << R"(", )"
                 << R"("value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << "]}";
}

HeterogeneousShocksStatement::HeterogeneousShocksStatement(
    int heterogeneity_dimension_arg, bool overwrite_arg, var_and_std_shocks_t var_shocks_arg,
    var_and_std_shocks_t std_shocks_arg, covar_and_corr_shocks_t covar_shocks_arg,
    covar_and_corr_shocks_t corr_shocks_arg, const SymbolTable& symbol_table_arg,
    const HeterogeneityTable& heterogeneity_table_arg) :
    heterogeneity_dimension {heterogeneity_dimension_arg},
    overwrite {overwrite_arg},
    var_shocks {move(var_shocks_arg)},
    std_shocks {move(std_shocks_arg)},
    covar_shocks {move(covar_shocks_arg)},
    corr_shocks {move(corr_shocks_arg)},
    symbol_table {symbol_table_arg},
    heterogeneity_table {heterogeneity_table_arg}
{
}

void
HeterogeneousShocksStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                          [[maybe_unused]] bool minimal_workspace) const
{
  if (overwrite)
    output << sigmaeName() << " = zeros(" << symbol_table.het_exo_nbr(heterogeneity_dimension)
           << ", " << symbol_table.het_exo_nbr(heterogeneity_dimension) << ");" << '\n';

  writeVarAndStdShocks(output);
  writeCovarAndCorrShocks(output);
}

void
HeterogeneousShocksStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "shocks")"
         << R"(, "heterogeneity": ")" << heterogeneity_table.getName(heterogeneity_dimension)
         << R"(", "overwrite": )" << boolalpha << overwrite;
  output << R"(, "variance": [)";
  for (bool printed_something {false}; auto& [id, value] : var_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"name": ")" << symbol_table.getName(id) << R"(", )"
             << R"("variance": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "stderr": [)";
  for (bool printed_something {false}; auto& [id, value] : std_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"name": ")" << symbol_table.getName(id) << R"(", )"
             << R"("stderr": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "covariance": [)";
  for (bool printed_something {false}; auto& [ids, value] : covar_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << "{"
             << R"("name": ")" << symbol_table.getName(ids.first) << R"(", )"
             << R"("name2": ")" << symbol_table.getName(ids.second) << R"(", )"
             << R"("covariance": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << R"(, "correlation": [)";
  for (bool printed_something {false}; auto& [ids, value] : corr_shocks)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << "{"
             << R"("name": ")" << symbol_table.getName(ids.first) << R"(", )"
             << R"("name2": ")" << symbol_table.getName(ids.second) << R"(", )"
             << R"("correlation": ")";
      value->writeJsonOutput(output, {}, {});
      output << R"("})";
    }
  output << "]"
         << "}";
}

void
HeterogeneousShocksStatement::writeVarOrStdShock(ostream& output, const pair<int, expr_t>& it,
                                                 bool stddev) const
{
  SymbolType type = symbol_table.getType(it.first);
  assert(type == SymbolType::heterogeneousExogenous);

  int id {symbol_table.getTypeSpecificID(it.first) + 1};
  output << sigmaeName() << "(" << id << ", " << id << ") = ";
  if (stddev)
    output << "(";
  it.second->writeOutput(output);
  if (stddev)
    output << ")^2";
  output << ";" << '\n';
}

void
HeterogeneousShocksStatement::writeVarAndStdShocks(ostream& output) const
{
  for (const auto& it : var_shocks)
    writeVarOrStdShock(output, it, false);

  for (const auto& it : std_shocks)
    writeVarOrStdShock(output, it, true);
}

void
HeterogeneousShocksStatement::writeCovarOrCorrShock(ostream& output,
                                                    const pair<pair<int, int>, expr_t>& it,
                                                    bool corr) const
{
  assert(symbol_table.getType(it.first.first) == SymbolType::heterogeneousExogenous
         && symbol_table.getType(it.first.second) == SymbolType::heterogeneousExogenous);
  int id1 {symbol_table.getTypeSpecificID(it.first.first) + 1},
      id2 {symbol_table.getTypeSpecificID(it.first.second) + 1};

  output << sigmaeName() << "(" << id1 << ", " << id2 << ") = ";
  it.second->writeOutput(output);
  if (corr)
    output << "*sqrt(" << sigmaeName() << "(" << id1 << ", " << id1 << ")*" << sigmaeName() << "("
           << id2 << ", " << id2 << "))";
  output << ";" << '\n'
         << sigmaeName() << "(" << id2 << ", " << id1 << ") = " << sigmaeName() << "(" << id1
         << ", " << id2 << ");" << '\n';
}

void
HeterogeneousShocksStatement::writeCovarAndCorrShocks(ostream& output) const
{
  for (const auto& it : covar_shocks)
    writeCovarOrCorrShock(output, it, false);

  for (const auto& it : corr_shocks)
    writeCovarOrCorrShock(output, it, true);
}

void
HeterogeneousShocksStatement::checkPass(ModFileStructure& mod_file_struct,
                                        [[maybe_unused]] WarningConsolidation& warnings)
{
  /* Error out if variables are not of the right type. This must be done here
     and not at parsing time (see #448). */
  for (const auto& [id, val] : var_shocks)
    if (symbol_table.getType(id) != SymbolType::heterogeneousExogenous)
      {
        cerr << "shocks: setting a variance on '" << symbol_table.getName(id)
             << "' is not allowed, because it is not a heterogeneous exogenous variable" << '\n';
        exit(EXIT_FAILURE);
      }

  for (const auto& [id, val] : std_shocks)
    if (symbol_table.getType(id) != SymbolType::heterogeneousExogenous)
      {
        cerr << "shocks: setting a standard error on '" << symbol_table.getName(id)
             << "' is not allowed, because it is not a heterogeneous exogenous variable" << '\n';
        exit(EXIT_FAILURE);
      }

  for (const auto& [ids, val] : covar_shocks)
    {
      auto& [symb_id1, symb_id2] = ids;

      if (!(symbol_table.getType(symb_id1) == SymbolType::heterogeneousExogenous
            && symbol_table.getType(symb_id2) == SymbolType::heterogeneousExogenous))
        {
          cerr << "shocks: setting a covariance between '" << symbol_table.getName(symb_id1)
               << "' and '" << symbol_table.getName(symb_id2)
               << "'is not allowed; covariances can only be specified for heterogeneous exogenous "
                  "variables"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  for (const auto& [ids, val] : corr_shocks)
    {
      auto& [symb_id1, symb_id2] = ids;

      if (!(symbol_table.getType(symb_id1) == SymbolType::heterogeneousExogenous
            && symbol_table.getType(symb_id2) == SymbolType::heterogeneousExogenous))
        {
          cerr << "shocks: setting a correlation between '" << symbol_table.getName(symb_id1)
               << "' and '" << symbol_table.getName(symb_id2)
               << "'is not allowed; covariances can only be specified for heterogeneous exogenous "
                  "variables"
               << '\n';
          exit(EXIT_FAILURE);
        }
    }

  // Fill in mod_file_struct.parameters_with_shocks_values (related to #469)
  for (const auto& [id, val] : var_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [id, val] : std_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [ids, val] : covar_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
  for (const auto& [ids, val] : corr_shocks)
    val->collectVariables(SymbolType::parameter, mod_file_struct.parameters_within_shocks_values);
}

ConditionalForecastPathsStatement::ConditionalForecastPathsStatement(
    AbstractShocksStatement::det_shocks_t paths_arg, const SymbolTable& symbol_table_arg) :
    paths {move(paths_arg)}, symbol_table {symbol_table_arg}, path_length {computePathLength(paths)}
{
}

int
ConditionalForecastPathsStatement::computePathLength(
    const AbstractShocksStatement::det_shocks_t& paths)
{
  int length {0};
  for (const auto& [ignore, elems] : paths)
    for (auto& [period_range, value] : elems)
      {
        const auto& [period1, period2] = get<pair<int, int>>(period_range);
        // Period1 < Period2, as enforced in ParsingDriver::add_period()
        length = max(length, period2);
      }
  return length;
}

void
ConditionalForecastPathsStatement::writeOutput(ostream& output,
                                               [[maybe_unused]] const string& basename,
                                               [[maybe_unused]] bool minimal_workspace) const
{
  assert(path_length > 0);
  output << "constrained_vars_ = [];" << '\n'
         << "constrained_paths_ = NaN(" << paths.size() << ", " << path_length << ");" << '\n';

  for (int k {1}; const auto& [id, elems] : paths)
    {
      if (k == 1)
        output << "constrained_vars_ = " << symbol_table.getTypeSpecificID(id) + 1 << ";" << '\n';
      else
        output << "constrained_vars_ = [constrained_vars_; "
               << symbol_table.getTypeSpecificID(id) + 1 << "];" << '\n';
      for (const auto& [period_range, value] : elems)
        {
          const auto& [period1, period2] = get<pair<int, int>>(period_range);
          output << "constrained_paths_(" << k << "," << period1 << ":" << period2 << ")=";
          value->writeOutput(output);
          output << ";" << '\n';
        }
      k++;
    }
}

void
ConditionalForecastPathsStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "conditional_forecast_paths")"
         << R"(, "paths": [)";
  for (bool printed_something {false}; const auto& [id, elems] : paths)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : elems)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << "]}";
}

PerfectForesightControlledPathsStatement::PerfectForesightControlledPathsStatement(
    paths_t paths_arg, variant<int, string> learnt_in_period_arg,
    const SymbolTable& symbol_table_arg) :
    paths {move(paths_arg)},
    learnt_in_period {move(learnt_in_period_arg)},
    symbol_table {symbol_table_arg}
{
}

void
PerfectForesightControlledPathsStatement::checkPass(ModFileStructure& mod_file_struct,
                                                    [[maybe_unused]] WarningConsolidation& warnings)
{
  mod_file_struct.perfect_foresight_controlled_paths_present = true;
  if (!(holds_alternative<int>(learnt_in_period) && get<int>(learnt_in_period) == 1))
    mod_file_struct.perfect_foresight_controlled_paths_learnt_in_present = true;
}

void
PerfectForesightControlledPathsStatement::writeOutput(ostream& output,
                                                      [[maybe_unused]] const string& basename,
                                                      [[maybe_unused]] bool minimal_workspace) const
{
  output << "M_.perfect_foresight_controlled_paths = [ M_.perfect_foresight_controlled_paths;"
         << '\n';
  for (const auto& [exogenize_id, constraints, endogenize_id] : paths)
    for (const auto& [period_range, value] : constraints)
      {
        output << "struct('exogenize_id'," << symbol_table.getTypeSpecificID(exogenize_id) + 1
               << ",'periods',";
        visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
        output << ",'value',";
        value->writeOutput(output);
        output << ",'endogenize_id'," << symbol_table.getTypeSpecificID(endogenize_id) + 1
               << ",'learnt_in',";
        visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
PerfectForesightControlledPathsStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "perfect_foresight_controlled_paths", "learnt_in": )";
  visit([&](const auto& p) { print_json_learnt_in(output, p); }, learnt_in_period);
  output << R"(, "paths": [)";
  for (bool printed_something {false};
       const auto& [exogenize_id, constraints, endogenize_id] : paths)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"exogenize": ")" << symbol_table.getName(exogenize_id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : constraints)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << R"(], "endogenize": ")" << symbol_table.getName(endogenize_id) << R"("})";
    }
  output << "]}";
}

MomentCalibration::MomentCalibration(constraints_t constraints_arg,
                                     const SymbolTable& symbol_table_arg) :
    constraints {move(constraints_arg)}, symbol_table {symbol_table_arg}
{
}

void
MomentCalibration::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                               [[maybe_unused]] bool minimal_workspace) const
{
  output << "options_.endogenous_prior_restrictions.moment = {" << '\n';
  for (const auto& c : constraints)
    {
      output << "'" << symbol_table.getName(c.endo1) << "', "
             << "'" << symbol_table.getName(c.endo2) << "', " << c.lags << ", "
             << "[ ";
      c.lower_bound->writeOutput(output);
      output << ", ";
      c.upper_bound->writeOutput(output);
      output << " ];" << '\n';
    }
  output << "};" << '\n';
}

void
MomentCalibration::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "moment_calibration")"
         << R"(, "moment_calibration_criteria": [)";
  for (bool printed_something {false}; const auto& c : constraints)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"endogenous1": ")" << symbol_table.getName(c.endo1) << R"(")"
             << R"(, "endogenous2": ")" << symbol_table.getName(c.endo2) << R"(")"
             << R"(, "lags": ")" << c.lags << R"(")"
             << R"(, "lower_bound": ")";
      c.lower_bound->writeJsonOutput(output, {}, {});
      output << R"(")"
             << R"(, "upper_bound": ")";
      c.upper_bound->writeJsonOutput(output, {}, {});
      output << R"(")"
             << "}";
    }
  output << "]"
         << "}";
}

IrfCalibration::IrfCalibration(constraints_t constraints_arg, const SymbolTable& symbol_table_arg,
                               OptionsList options_list_arg) :
    constraints {move(constraints_arg)},
    symbol_table {symbol_table_arg},
    options_list {move(options_list_arg)}
{
}

void
IrfCalibration::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                            [[maybe_unused]] bool minimal_workspace) const
{
  options_list.writeOutput(output);

  output << "options_.endogenous_prior_restrictions.irf = {" << '\n';
  for (const auto& c : constraints)
    {
      output << "'" << symbol_table.getName(c.endo) << "', "
             << "'" << symbol_table.getName(c.exo) << "', " << c.periods << ", "
             << "[ ";
      c.lower_bound->writeOutput(output);
      output << ", ";
      c.upper_bound->writeOutput(output);
      output << " ];" << '\n';
    }
  output << "};" << '\n';
}

void
IrfCalibration::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "irf_calibration")";
  if (!options_list.empty())
    {
      output << ", ";
      options_list.writeJsonOutput(output);
    }

  output << R"(, "irf_restrictions": [)";
  for (bool printed_something {false}; const auto& c : constraints)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"endogenous": ")" << symbol_table.getName(c.endo) << R"(")"
             << R"(, "exogenous": ")" << symbol_table.getName(c.exo) << R"(")"
             << R"(, "periods": ")" << c.periods << R"(")"
             << R"(, "lower_bound": ")";
      c.lower_bound->writeJsonOutput(output, {}, {});
      output << R"(")";
      output << R"(, "upper_bound": ")";
      c.upper_bound->writeJsonOutput(output, {}, {});
      output << R"(")"
             << "}";
    }
  output << "]"
         << "}";
}

ShockGroupsStatement::ShockGroupsStatement(group_t shock_groups_arg, string name_arg) :
    shock_groups {move(shock_groups_arg)}, name {move(name_arg)}
{
}

void
ShockGroupsStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                  [[maybe_unused]] bool minimal_workspace) const
{
  int i = 1;
  for (auto it = shock_groups.begin(); it != shock_groups.end(); ++it)
    {
      bool unique_label {true};
      for (auto it1 = it + 1; it1 != shock_groups.end(); ++it1)
        if (it->name == it1->name)
          {
            unique_label = false;
            cerr << "Warning: shock group label '" << it->name << "' has been reused. "
                 << "Only using the last definition." << '\n';
            break;
          }

      if (unique_label)
        {
          output << "M_.shock_groups." << name << ".group" << i << ".label = '" << it->name << "';"
                 << '\n'
                 << "M_.shock_groups." << name << ".group" << i << ".shocks = {";
          for (const auto& it1 : it->list)
            output << " '" << it1 << "'";
          output << "};" << '\n';
          i++;
        }
    }
}

void
ShockGroupsStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "shock_groups", "name": ")" << name << R"(", "groups": [)";
  bool printed_something {false};
  for (auto it = shock_groups.begin(); it != shock_groups.end(); ++it)
    {
      bool unique_label {true};
      for (auto it1 = it + 1; it1 != shock_groups.end(); ++it1)
        if (it->name == it1->name)
          {
            unique_label = false;
            break;
          }

      if (unique_label)
        {
          if (exchange(printed_something, true))
            output << ", ";
          output << R"({"group_name": ")" << it->name << R"(",)"
                 << R"("shocks": [)";
          for (bool printed_something2 {false}; const auto& it1 : it->list)
            {
              if (exchange(printed_something2, true))
                output << ", ";
              output << R"(")" << it1 << R"(")";
            }
          output << "]}";
        }
    }
  output << "]}";
}

Init2shocksStatement::Init2shocksStatement(vector<pair<int, int>> init2shocks_arg, string name_arg,
                                           const SymbolTable& symbol_table_arg) :
    init2shocks {move(init2shocks_arg)}, name {move(name_arg)}, symbol_table {symbol_table_arg}
{
}

void
Init2shocksStatement::checkPass([[maybe_unused]] ModFileStructure& mod_file_struct,
                                [[maybe_unused]] WarningConsolidation& warnings)
{
  for (size_t i = 0; i < init2shocks.size(); i++)
    for (size_t j = i + 1; j < init2shocks.size(); j++)
      if (init2shocks.at(i).first == init2shocks.at(j).first)
        {
          cerr << "Init2shocks(" << name << "): enogenous variable '"
               << symbol_table.getName(init2shocks.at(i).first)
               << "' appears more than once in the init2shocks statement" << '\n';
          exit(EXIT_FAILURE);
        }
}

void
Init2shocksStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                  [[maybe_unused]] bool minimal_workspace) const
{
  output << "M_.init2shocks." << name << " = {" << '\n';
  for (const auto& [id1, id2] : init2shocks)
    output << "{'" << symbol_table.getName(id1) << "', '" << symbol_table.getName(id2) << "'};"
           << '\n';
  output << "};" << '\n';
}

void
Init2shocksStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "init2shocks", "name": ")" << name << R"(", "groups": [)";
  for (bool printed_something {false}; const auto& [id1, id2] : init2shocks)
    {
      if (exchange(printed_something, true))
        output << ",";
      output << R"({"endogenous": ")" << symbol_table.getName(id1) << R"(", )"
             << R"( "exogenous": ")" << symbol_table.getName(id2) << R"("})";
    }
  output << "]}";
}

HeteroskedasticShocksStatement::HeteroskedasticShocksStatement(
    bool overwrite_arg, heteroskedastic_shocks_t values_arg, heteroskedastic_shocks_t scales_arg,
    heteroskedastic_shocks_t endo_values_arg, heteroskedastic_shocks_t endo_scales_arg,
    const SymbolTable& symbol_table_arg) :
    overwrite {overwrite_arg},
    exo_values {move(values_arg)},
    exo_scales {move(scales_arg)},
    endo_values {move(endo_values_arg)},
    endo_scales {move(endo_scales_arg)},
    symbol_table {symbol_table_arg}
{
}

void
HeteroskedasticShocksStatement::writeOutput(ostream& output,
                                            [[maybe_unused]] const string& basename,
                                            [[maybe_unused]] bool minimal_workspace) const
{
  for (const auto& [matrix, type, data] :
       {tuple {"Q", "value", &exo_values}, tuple {"Q", "scale", &exo_scales},
        tuple {"H", "value", &endo_values}, tuple {"H", "scale", &endo_scales}})
    {
      // NB: The first initialization of the fields is done in ModFile::writeMOutput()
      if (overwrite)
        output << "M_.heteroskedastic_shocks." << matrix << type << "_orig = struct([]);" << '\n';

      if (!data->empty())
        {
          output << "M_.heteroskedastic_shocks." << matrix << type
                 << "_orig = [M_.heteroskedastic_shocks." << matrix << type << "_orig;" << '\n';
          for (const auto& [symb_id, vec] : *data)
            for (int tsid = symbol_table.getTypeSpecificID(symb_id);
                 const auto& [period_range, value] : vec)
              {
                output << "struct('" << (matrix == "Q"s ? "exo" : "endo") << "_id', " << tsid + 1
                       << ",'periods',";
                visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
                output << ",'" << type << "',";
                value->writeOutput(output);
                output << ");" << '\n';
              }
          output << "];" << '\n';
        }
    }
}

void
HeteroskedasticShocksStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "heteroskedastic_shocks")"
         << R"(, "overwrite": )" << boolalpha << overwrite;

  for (const auto& [symb_type, shock_type, data] :
       {tuple {"exo", "values", &exo_values}, tuple {"exo", "scales", &exo_scales},
        tuple {"endo", "values", &endo_values}, tuple {"endo", "scales", &endo_scales}})
    {
      output << R"(, "shocks_)" << symb_type << '_' << shock_type << R"(": [)";
      for (bool printed_something {false}; const auto& [symb_id, vec] : *data)
        {
          if (exchange(printed_something, true))
            output << ", ";
          output << R"({"var": ")" << symbol_table.getName(symb_id) << R"(", ")" << shock_type
                 << R"(": [)";
          for (bool printed_something2 {false}; const auto& [period_range, value] : vec)
            {
              if (exchange(printed_something2, true))
                output << ", ";
              output << "{";
              visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
              output << R"(, "value": ")";
              value->writeJsonOutput(output, {}, {});
              output << R"("})";
            }
          output << "]}";
        }
      output << ']';
    }
  output << "}";
}

ShockPathsStatement::ShockPathsStatement(
    variant<int, string> learnt_in_period_arg, bool overwrite_arg, exo_paths_t exo_paths_arg,
    PerfectForesightControlledPathsStatement::paths_t controlled_paths_arg,
    const SymbolTable& symbol_table_arg) :
    learnt_in_period {move(learnt_in_period_arg)},
    overwrite {overwrite_arg},
    exo_paths {move(exo_paths_arg)},
    controlled_paths {move(controlled_paths_arg)},
    symbol_table {symbol_table_arg}
{
}

void
ShockPathsStatement::checkPass(ModFileStructure& mod_file_struct,
                               [[maybe_unused]] WarningConsolidation& warnings)
{
  index = ++mod_file_struct.shock_paths_number;
  if (!(holds_alternative<int>(learnt_in_period) && get<int>(learnt_in_period) == 1))
    mod_file_struct.shock_paths_learnt_in_present = true;

  for (const auto& [symb_id, shock_vec] : exo_paths)
    for (const auto& [period_range, value] : shock_vec)
      {
        set<pair<int, int>> self_vars;
        value->collectSelfVariables(self_vars);
        for (const auto& [symb_id2, lag] : self_vars)
          {
            if (symb_id == symb_id2 && lag == 0)
              {
                cerr << "ERROR: in the definition of '" << symbol_table.getName(symb_id)
                     << "' in a 'shock_paths' block, the use of 'self."
                     << symbol_table.getName(symb_id)
                     << "' without a lag is not allowed, since it is a circular reference" << '\n';
                exit(EXIT_FAILURE);
              }
            assert(lag <= 0); // Already checked in ParsingDriver::add_self_variable()
            // NB: the restriction on lag is needed
          }
      }
}

void
ShockPathsStatement::writeOutput(ostream& output, const string& basename,
                                 [[maybe_unused]] bool minimal_workspace) const
{
  if (overwrite)
    {
      output << "if ~isempty(M_.shock_paths)" << '\n'
             << "  M_.shock_paths = M_.shock_paths(cellfun(@(x) ~isa(x, '";
      if (holds_alternative<int>(learnt_in_period))
        output << "numeric";
      else
        output << "dates";
      output << "') || x ~= ";
      /* NB: date expression not parenthesized since it can only contain a + operator, which has
         higher precedence than ~= and || */
      visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
      output << ", {M_.shock_paths.learnt_in}));" << '\n' << "end" << '\n';
    }

  /* Whether this block has a date somewhere:
     – in the “learnt_in” option
     – in the “periods” statement
     – in the “learnt_in()” namespace of a variable inside the expressions */
  bool contains_date {
      holds_alternative<string>(learnt_in_period)
      || ranges::any_of(exo_paths,
                        [](const auto& symb_vec) {
                          return ranges::any_of(symb_vec.second, [](const auto& period_value) {
                            auto& [period, value] = period_value;
                            return holds_alternative<string>(period.get_first())
                                   || holds_alternative<string>(period.get_last())
                                   || value->containsDate();
                          });
                        })
      || ranges::any_of(controlled_paths, [](const auto& it) {
           return ranges::any_of(get<1>(it), [](const auto& period_value) {
             auto& [period_range, value] = period_value;
             return holds_alternative<pair<string, string>>(period_range) || value->containsDate();
           });
         })};

  // Whether this block has an “end” period (and thus triggers a permanent shock)
  bool contains_endval {ranges::any_of(exo_paths, [](const auto& symb_vec) {
    return ranges::any_of(symb_vec.second, [](const auto& period_value) {
      auto& [period, value] = period_value;
      return holds_alternative<monostate>(period.get_last());
    });
  })};

  output << "M_.shock_paths = [M_.shock_paths; struct('learnt_in', ";
  visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
  output << ", 'evaluation_function', '" << basename << "." << evaluationFunctionName()
         << "', 'contains_date', " << boolalpha << contains_date << ", 'contains_endval', "
         << contains_endval << ")];" << '\n';

  writeEvaluationFunctionFile(basename);

  // Write the controlled paths
  /* FIXME: for the time being, we populate M_.perfect_foresight_controlled_paths as does the block
     of the same name, but in the future we will write this information into the generated
     shock_paths evaluation functions, so as to support all syntaxes (self, database, prev,
     learnt_in(LAG)…).
     In particular, this is the reason why we have some code duplication with the
     PerfectForesightControlledPathsStatement class, since it is meant to be temporary. */
  if (overwrite)
    {
      output << "if ~isempty(M_.perfect_foresight_controlled_paths)" << '\n'
             << "  M_.perfect_foresight_controlled_paths = "
                "M_.perfect_foresight_controlled_paths(cellfun(@(x) ~isa(x, '";
      if (holds_alternative<int>(learnt_in_period))
        output << "numeric";
      else
        output << "dates";
      output << "') || x ~= ";
      /* NB: date expression not parenthesized since it can only contain a + operator, which has
         higher precedence than ~= and || */
      visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
      output << ", {M_.perfect_foresight_controlled_paths.learnt_in}));" << '\n' << "end" << '\n';
    }
  output << "M_.perfect_foresight_controlled_paths = [ M_.perfect_foresight_controlled_paths;"
         << '\n';
  for (const auto& [exogenize_id, constraints, endogenize_id] : controlled_paths)
    for (const auto& [period_range, value] : constraints)
      {
        output << "struct('exogenize_id'," << symbol_table.getTypeSpecificID(exogenize_id) + 1
               << ",'periods',";
        visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
        output << ",'value',";
        value->writeOutput(output);
        output << ",'endogenize_id'," << symbol_table.getTypeSpecificID(endogenize_id) + 1
               << ",'learnt_in',";
        visit([&](const auto& p) { print_matlab_learnt_in(output, p); }, learnt_in_period);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
ShockPathsStatement::writeJsonOutput(ostream& output) const
{
  auto print_json_period_t = [&output](const auto& arg) {
    if (holds_alternative<int>(arg))
      output << get<int>(arg);
    else if (holds_alternative<string>(arg))
      output << '"' << get<string>(arg) << '"';
    else // monostate
      output << R"("end")";
  };

  output << R"({"statementName": "shock_paths")"
         << R"(, "learnt_in": )";
  visit([&](const auto& p) { print_json_learnt_in(output, p); }, learnt_in_period);
  output << R"(, "overwrite": )" << boolalpha << overwrite << R"(, "exo_paths": [)";
  for (bool printed_something {false}; const auto& [id, shock_vec] : exo_paths)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : shock_vec)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << R"({"period1" :)";
          print_json_period_t(period_range.get_first());
          output << R"(, "period2" :)";
          print_json_period_t(period_range.get_last());
          output << R"(", "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << R"(], "controlled_paths": [)";
  for (bool printed_something {false};
       const auto& [exogenize_id, constraints, endogenize_id] : controlled_paths)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"exogenize": ")" << symbol_table.getName(exogenize_id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value] : constraints)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << R"(], "endogenize": ")" << symbol_table.getName(endogenize_id) << R"("})";
    }
  output << "]}";
}

void
ShockPathsStatement::writeEvaluationFunctionFile(const string& basename) const
{
  filesystem::path filename {packageDir(basename) / (evaluationFunctionName() + ".m")};
  ofstream output {filename, ios::out | ios::binary};
  if (!output.is_open())
    {
      cerr << "ERROR: Can't open file " << filename.string() << " for writing" << '\n';
      exit(EXIT_FAILURE);
    }

  // M_ is there for parameters, oo_ for initval namespace
  output << "function exo_paths = " << evaluationFunctionName()
         << "(exo_paths, p, periods, first_simulation_period, M_, oo_)" << '\n'
         << "info_period = ";
  if (holds_alternative<int>(learnt_in_period))
    output << get<int>(learnt_in_period);
  else
    output << get<string>(learnt_in_period) << "-first_simulation_period+1";
  output << ";" << '\n';

  auto print_matlab_period_shock_paths = [&output](const auto& p) {
    if (holds_alternative<int>(p))
      output << get<int>(p);
    else if (holds_alternative<string>(p))
      output << get<string>(p) << "-first_simulation_period+1";
    else // monostate
      output << "periods+1";
  };

  for (const auto& [symb_id, shock_vec] : exo_paths)
    for (const auto& [period_range, value] : shock_vec)
      {
        auto p1 = period_range.get_first();
        auto p2 = period_range.get_last();
        if (p1 == p2)
          {
            output << "if p == ";
            print_matlab_period_shock_paths(p1);
          }
        else if (holds_alternative<monostate>(p2))
          {
            output << "if p >= ";
            print_matlab_period_shock_paths(p1);
          }
        else if (holds_alternative<int>(p1) && get<int>(p1) == 1)
          {
            output << "if p <= ";
            print_matlab_period_shock_paths(p2);
          }
        else
          {
            output << "if p >= ";
            print_matlab_period_shock_paths(p1);
            output << " && p <= ";
            print_matlab_period_shock_paths(p2);
          }
        output << '\n'
               << "exo_paths(" << symbol_table.getTypeSpecificID(symb_id) + 1
               << ",p,info_period) = ";
        value->writeOutput(output);
        output << ";" << '\n' << "end" << '\n';
      }

  output << "end" << '\n';
  output.close();
}

FilterTunesStatement::FilterTunesStatement(tunes_t tunes_arg, const SymbolTable& symbol_table_arg) :
    tunes {move(tunes_arg)}, symbol_table {symbol_table_arg}
{
}

void
FilterTunesStatement::checkPass([[maybe_unused]] ModFileStructure& mod_file_struct,
                                [[maybe_unused]] WarningConsolidation& warnings)
{
  for (int id : views::keys(tunes))
    if (symbol_table.isObservedVariable(id))
      {
        cerr << "filter_tunes: variable " << symbol_table.getName(id)
             << " is declared as an observable" << endl;
        exit(EXIT_FAILURE);
      }
}

void
FilterTunesStatement::writeOutput(ostream& output, [[maybe_unused]] const string& basename,
                                  [[maybe_unused]] bool minimal_workspace) const
{
  output << "M_.filter_tunes = [ M_.filter_tunes;" << '\n';
  for (const auto& [id, tunes_vec] : tunes)
    for (const auto& [period_range, value, std_err] : tunes_vec)
      {
        output << "struct('endo_id'," << symbol_table.getTypeSpecificID(id) + 1 << ",'periods',";
        visit([&](const auto& p) { print_matlab_period_range(output, p); }, period_range);
        output << ",'value',";
        value->writeOutput(output);
        output << ",'stderr',";
        std_err->writeOutput(output);
        output << ");" << '\n';
      }
  output << "];" << '\n';
}

void
FilterTunesStatement::writeJsonOutput(ostream& output) const
{
  output << R"({"statementName": "filter_tunes", "tunes": [)";
  for (bool printed_something {false}; const auto& [id, tunes_vec] : tunes)
    {
      if (exchange(printed_something, true))
        output << ", ";
      output << R"({"var": ")" << symbol_table.getName(id) << R"(", )"
             << R"("values": [)";
      for (bool printed_something2 {false}; const auto& [period_range, value, std_err] : tunes_vec)
        {
          if (exchange(printed_something2, true))
            output << ", ";
          output << "{";
          visit([&](const auto& p) { print_json_period_range(output, p); }, period_range);
          output << R"(, "value": ")";
          value->writeJsonOutput(output, {}, {});
          output << R"(", "stderr": ")";
          std_err->writeJsonOutput(output, {}, {});
          output << R"("})";
        }
      output << "]}";
    }
  output << "]}";
}
