/*
 * Copyright (C) 2010 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 <http://www.gnu.org/licenses/>.
 */

#include <cassert>
#include <algorithm>

#include "SteadyStateModel.hh"

SteadyStateModel::SteadyStateModel(SymbolTable &symbol_table_arg, NumericalConstants &num_constants, ExternalFunctionsTable &external_functions_table_arg, const StaticModel &static_model_arg) :
  DataTree(symbol_table_arg, num_constants, external_functions_table), static_model(static_model_arg)
{
}

void
SteadyStateModel::addDefinition(int symb_id, expr_t expr)
{
  assert(symbol_table.getType(symb_id) == eEndogenous
         || symbol_table.getType(symb_id) == eModFileLocalVariable
         || symbol_table.getType(symb_id) == eParameter);

  // Add the variable
  recursive_order.push_back(symb_id);
  def_table[symb_id] = AddEqual(AddVariable(symb_id), expr);
}

void
SteadyStateModel::checkPass(bool ramsey_policy) const
{
  for (vector<int>::const_iterator it = recursive_order.begin();
       it != recursive_order.end(); ++it)
    {
      // Check that symbol is not already defined
      if (find(recursive_order.begin(), it, *it) != it)
        {
          cerr << "ERROR: in the 'steady_state' block, variable '" << symbol_table.getName(*it) << "' is declared twice" << endl;
          exit(EXIT_FAILURE);
        }

      // Check that expression has no undefined symbol
      if (!ramsey_policy)
        {
          set<pair<int, int> > used_symbols;
          expr_t expr = def_table.find(*it)->second;
          expr->collectVariables(eEndogenous, used_symbols);
          expr->collectVariables(eModFileLocalVariable, used_symbols);
          for(set<pair<int, int> >::const_iterator it2 = used_symbols.begin();
              it2 != used_symbols.end(); ++it2)
            if (find(recursive_order.begin(), it, it2->first) == it
                && *it != it2->first)
              {
                cerr << "ERROR: in the 'steady_state' block, variable '" << symbol_table.getName(it2->first) << "' is undefined in the declaration of variable '" << symbol_table.getName(*it) << "'" << endl;
                exit(EXIT_FAILURE);
              }
        }
    }
}

void
SteadyStateModel::writeSteadyStateFile(const string &basename, bool ramsey_policy) const
{
  if (recursive_order.size() == 0)
    return;

  string filename = basename + "_steadystate.m";

  ofstream output;
  output.open(filename.c_str(), ios::out | ios::binary);
  if (!output.is_open())
    {
      cerr << "ERROR: Can't open file " << filename << " for writing" << endl;
      exit(EXIT_FAILURE);
    }

  output << "function [ys_, check_] = " << basename << "_steadystate(";
  if (ramsey_policy)
    output << "ys_";
  else
    output << "ys_orig_";
  output << ", exo_)" << endl
         << "% Steady state generated by Dynare preprocessor" << endl;
  if (!ramsey_policy)
    output << "    ys_=zeros(" << symbol_table.orig_endo_nbr() << ",1);" << endl;
  output << "    global M_" << endl;

  for(size_t i = 0; i < recursive_order.size(); i++)
    {
      output << "    ";
      map<int, expr_t>::const_iterator it = def_table.find(recursive_order[i]);
      it->second->writeOutput(output, oSteadyStateFile);
      output << ";" << endl;
    }
  output << "    % Auxiliary equations" << endl;
  static_model.writeAuxVarInitval(output, oSteadyStateFile);
  output << "    check_=0;" << endl
         << "end" << endl;
}