/*
 * Copyright © 2008-2023 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/>.
 */

// GP, based on work by O.Kamenik

#include "k_ord_dynare.hh"
#include "dynamic_abstract_class.hh"
#include "dynare_exception.hh"

#include <cassert>
#include <utility>

KordpDynare::KordpDynare(const std::vector<std::string>& endo, const std::vector<std::string>& exo,
                         int nexog, int npar, Vector& ysteady, TwoDMatrix& vcov, Vector& inParams,
                         int nstat, int npred, int nforw, int nboth, const ConstVector& nnzd,
                         int nsteps, int norder, Journal& jr,
                         std::unique_ptr<DynamicModelAC> dynamicModelFile_arg,
                         const std::vector<int>& dr_order, const ConstTwoDMatrix& llincidence) :
    nStat {nstat},
    nBoth {nboth},
    nPred {npred},
    nForw {nforw},
    nExog {nexog},
    nPar {npar},
    nYs {npred + nboth},
    nYss {nboth + nforw},
    nY {nstat + npred + nboth + nforw},
    nJcols {nExog + nY + nYs + nYss},
    NNZD {nnzd},
    nSteps {nsteps},
    nOrder {norder},
    journal {jr},
    ySteady {ysteady},
    params {inParams},
    vCov {vcov},
    md {1},
    dnl {endo},
    denl {exo},
    dsnl {*this, dnl, denl},
    ll_Incidence {llincidence},
    dynamicModelFile {std::move(dynamicModelFile_arg)}
{
  computeJacobianPermutation(dr_order);
}

void
KordpDynare::solveDeterministicSteady()
{
  JournalRecordPair pa(journal);
  pa << "Non-linear solver for deterministic steady state skipped" << endrec;
}

void
KordpDynare::evaluateSystem(Vector& out, [[maybe_unused]] const ConstVector& yy,
                            [[maybe_unused]] const Vector& xx)
{
  // This method is only called when checking the residuals at steady state (Approximation::check),
  // so return zero residuals
  out.zeros();
}

void
KordpDynare::evaluateSystem(Vector& out, [[maybe_unused]] const ConstVector& yym,
                            [[maybe_unused]] const ConstVector& yy,
                            [[maybe_unused]] const ConstVector& yyp,
                            [[maybe_unused]] const Vector& xx)
{
  // This method is only called when checking the residuals at steady state (Approximation::check),
  // so return zero residuals
  out.zeros();
}

void
KordpDynare::calcDerivativesAtSteady()
{
  assert(md.begin() == md.end());

  std::vector<TwoDMatrix> dyn_md; // Model derivatives, in Dynare form

  dyn_md.emplace_back(nY, nJcols); // Allocate Jacobian
  dyn_md.back().zeros();

  for (int i = 2; i <= nOrder; i++)
    {
      // Higher order derivatives, as sparse (3-column) matrices
      dyn_md.emplace_back(static_cast<int>(NNZD[i - 1]), 3);
      dyn_md.back().zeros();
    }

  Vector xx(nexog());
  xx.zeros();

  Vector out(nY);
  out.zeros();
  Vector llxSteady(nJcols - nExog);
  LLxSteady(ySteady, llxSteady);

  dynamicModelFile->eval(llxSteady, xx, params, ySteady, out, dyn_md);

  for (int i = 1; i <= nOrder; i++)
    populateDerivativesContainer(dyn_md, i);
}

void
KordpDynare::populateDerivativesContainer(const std::vector<TwoDMatrix>& dyn_md, int ord)
{
  const TwoDMatrix& g = dyn_md[ord - 1];

  // model derivatives FSSparseTensor instance
  auto mdTi = std::make_unique<FSSparseTensor>(ord, nJcols, nY);

  IntSequence s(ord, 0);

  if (ord == 1)
    for (int i = 0; i < g.ncols(); i++)
      {
        for (int j = 0; j < g.nrows(); j++)
          {
            double x = g.get(j, dynppToDyn[s[0]]);
            if (x != 0.0)
              mdTi->insert(s, j, x);
          }
        s[0]++;
      }
  else // ord ≥ 2
    for (int i = 0; i < g.nrows(); i++)
      {
        int j = static_cast<int>(g.get(i, 0)) - 1;
        int i1 = static_cast<int>(g.get(i, 1)) - 1;
        if (j < 0 || i1 < 0)
          continue; // Discard empty entries (see comment in DynamicModelAC::unpackSparseMatrix())

        for (int k = 0; k < ord; k++)
          {
            s[k] = dynToDynpp[i1 % nJcols];
            i1 /= nJcols;
          }

        if (ord == 2 && !s.isSorted())
          continue; // Skip symmetric elements (only needed at order 2)
        else if (ord > 2)
          s.sort(); // For higher order, canonicalize the multi-index

        double x = g.get(i, 2);
        mdTi->insert(s, j, x);
      }

  md.insert(std::move(mdTi));
}

/* Returns ySteady extended with leads and lags suitable for passing to
   <model>_dynamic */
void
KordpDynare::LLxSteady(const Vector& yS, Vector& llxSteady)
{
  if (yS.length() == nJcols - nExog)
    throw DynareException(__FILE__, __LINE__, "ySteady already of right size");

  /* Create temporary square 2D matrix size nEndo×nEndo (sparse)
     for the lag, current and lead blocks of the jacobian */
  if (llxSteady.length() != nJcols - nExog)
    throw DynareException(__FILE__, __LINE__, "llxSteady has wrong size");

  for (int ll_row = 0; ll_row < ll_Incidence.nrows(); ll_row++)
    // populate (non-sparse) vector with ysteady values
    for (int i = 0; i < nY; i++)
      if (ll_Incidence.get(ll_row, i))
        llxSteady[static_cast<int>(ll_Incidence.get(ll_row, i)) - 1] = yS[i];
}

/*
   Computes mapping between Dynare and Dynare++ orderings of the (dynamic)
   variable indices in derivatives.

   If one defines:
   – y (resp. x) as the vector of all endogenous (size nY), in DR-order (resp.
     declaration order)
   – y⁻ (resp. x⁻) as the vector of endogenous that appear at previous period (size nYs),
     in DR-order (resp. declaration order)
   – y⁺ (resp. x⁺) as the vector of endogenous that appear at future period (size nYss) in
     DR-order (resp. declaration order)
   – u as the vector of exogenous (size nExog)

   In Dynare, the ordering is (x⁻, x, x⁺, u).
   In Dynare++, the ordering is (y⁺, y, y⁻, u).

   dr_order is typically equal to oo_.dr.order_var.
*/
void
KordpDynare::computeJacobianPermutation(const std::vector<int>& dr_order)
{
  // Compute restricted inverse DR-orderings: x⁻→y⁻ and x⁺→y⁺
  std::vector<int> dr_inv_order_forw(nBoth + nForw), dr_inv_order_pred(nBoth + nPred);
  std::iota(dr_inv_order_forw.begin(), dr_inv_order_forw.end(), 0);
  std::sort(dr_inv_order_forw.begin(), dr_inv_order_forw.end(), [&](int i, int j) {
    return dr_order[nStat + nPred + i] < dr_order[nStat + nPred + j];
  });
  std::iota(dr_inv_order_pred.begin(), dr_inv_order_pred.end(), 0);
  std::sort(dr_inv_order_pred.begin(), dr_inv_order_pred.end(),
            [&](int i, int j) { return dr_order[nStat + i] < dr_order[nStat + j]; });

  // Compute restricted DR-orderings: y⁻→x⁻ and y⁺→x⁺
  std::vector<int> dr_order_forw(nBoth + nForw), dr_order_pred(nBoth + nPred);
  for (int i = 0; i < nBoth + nForw; i++)
    dr_order_forw[dr_inv_order_forw[i]] = i;
  for (int i = 0; i < nBoth + nPred; i++)
    dr_order_pred[dr_inv_order_pred[i]] = i;

  // Compute Dynare++ → Dynare ordering
  dynppToDyn.resize(nJcols);
  int j = 0;
  for (; j < nYss; j++)
    dynppToDyn[j] = dr_order_forw[j] + nYs + nY; // Forward variables
  for (; j < nYss + nY; j++)
    dynppToDyn[j] = dr_order[j - nYss] + nYs; // Variables in current period
  for (; j < nYss + nY + nYs; j++)
    dynppToDyn[j] = dr_order_pred[j - nY - nYss]; // Predetermined variables
  for (; j < nJcols; j++)
    dynppToDyn[j] = j; // Exogenous

  // Compute Dynare → Dynare++ ordering
  dynToDynpp.resize(nJcols);
  for (int i = 0; i < nJcols; i++)
    dynToDynpp[dynppToDyn[i]] = i;
}

DynareNameList::DynareNameList(std::vector<std::string> names_arg) : names(std::move(names_arg))
{
}

DynareStateNameList::DynareStateNameList(const KordpDynare& dynare, const DynareNameList& dnl,
                                         const DynareNameList& denl)
{
  for (int i = 0; i < dynare.nYs; i++)
    names.emplace_back(dnl.getName(i + dynare.nStat));
  for (int i = 0; i < dynare.nExog; i++)
    names.emplace_back(denl.getName(i));
}