/*
 * Copyright © 2004-2011 Ondra Kamenik
 * Copyright © 2019-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/>.
 */

#include "BlockDiagonal.hh"
#include "int_power.hh"

#include <iostream>
#include <utility>

BlockDiagonal::BlockDiagonal(const ConstVector& d, int d_size) :
    QuasiTriangular(d, d_size), row_len(d_size), col_len(d_size)
{
  for (int i = 0; i < d_size; i++)
    {
      row_len[i] = d_size;
      col_len[i] = 0;
    }
}

BlockDiagonal::BlockDiagonal(const QuasiTriangular& t) :
    QuasiTriangular(t), row_len(t.nrows()), col_len(t.nrows())
{
  for (int i = 0; i < t.nrows(); i++)
    {
      row_len[i] = t.nrows();
      col_len[i] = 0;
    }
}

/* Put zeroes to right upper submatrix whose first column is defined
   by ‘edge’ */
void
BlockDiagonal::setZerosToRU(diag_iter edge)
{
  int iedge = edge->getIndex();
  for (int i = 0; i < iedge; i++)
    for (int j = iedge; j < ncols(); j++)
      get(i, j) = 0.0;
}

/* Updates row_len and col_len so that there are zeroes in upper right part, i.e.
   ⎛T1  0⎞
   ⎝ 0 T2⎠. The first column of T2 is given by diagonal iterator ‘edge’.

   Note the semantics of row_len and col_len. row_len[i] is distance
   of the right-most non-zero element of i-th row from the left, and
   col_len[j] is distance of top-most non-zero element of j-th column
   to the top. (First element has distance 1).
*/
void
BlockDiagonal::setZeroBlockEdge(diag_iter edge)
{
  setZerosToRU(edge);

  int iedge = edge->getIndex();
  for (auto run = diag_begin(); run != edge; ++run)
    if (int ind {run->getIndex()}; row_len[ind] > iedge)
      {
        row_len[ind] = iedge;
        if (!run->isReal())
          row_len[ind + 1] = iedge;
      }
  for (auto run = edge; run != diag_end(); ++run)
    if (int ind {run->getIndex()}; col_len[ind] < iedge)
      {
        col_len[ind] = iedge;
        if (!run->isReal())
          col_len[ind + 1] = iedge;
      }
}

BlockDiagonal::const_col_iter
BlockDiagonal::col_begin(const DiagonalBlock& b) const
{
  int jbar = b.getIndex();
  int d_size = diagonal.getSize();
  return {&getData()[jbar * d_size + col_len[jbar]], d_size, b.isReal(), col_len[jbar]};
}

BlockDiagonal::col_iter
BlockDiagonal::col_begin(const DiagonalBlock& b)
{
  int jbar = b.getIndex();
  int d_size = diagonal.getSize();
  return {&getData()[jbar * d_size + col_len[jbar]], d_size, b.isReal(), col_len[jbar]};
}

BlockDiagonal::const_row_iter
BlockDiagonal::row_end(const DiagonalBlock& b) const
{
  int jbar = b.getIndex();
  int d_size = diagonal.getSize();
  return {&getData()[d_size * row_len[jbar] + jbar], d_size, b.isReal(), row_len[jbar]};
}

BlockDiagonal::row_iter
BlockDiagonal::row_end(const DiagonalBlock& b)
{
  int jbar = b.getIndex();
  int d_size = diagonal.getSize();
  return {&getData()[d_size * row_len[jbar] + jbar], d_size, b.isReal(), row_len[jbar]};
}

int
BlockDiagonal::getNumZeros() const
{
  int sum = 0;
  for (int i = 0; i < diagonal.getSize(); i++)
    sum += diagonal.getSize() - row_len[i];
  return sum;
}

QuasiTriangular::const_diag_iter
BlockDiagonal::findBlockStart(const_diag_iter from) const
{
  if (from != diag_end())
    {
      ++from;
      while (from != diag_end() && col_len[from->getIndex()] != from->getIndex())
        ++from;
    }
  return from;
}

int
BlockDiagonal::getLargestBlock() const
{
  int largest = 0;
  auto start = diag_begin();
  auto end = findBlockStart(start);
  while (start != diag_end())
    {
      int si = start->getIndex();
      int ei = diagonal.getSize();
      if (end != diag_end())
        ei = end->getIndex();
      largest = std::max(largest, ei - si);
      start = end;
      end = findBlockStart(start);
    }
  return largest;
}

void
BlockDiagonal::savePartOfX(int si, int ei, const KronVector& x, Vector& work)
{
  for (int i = si; i < ei; i++)
    {
      ConstKronVector xi(x, i);
      Vector target(work, (i - si) * xi.length(), xi.length());
      target = xi;
    }
}

void
BlockDiagonal::multKronBlock(const_diag_iter start, const_diag_iter end, KronVector& x,
                             Vector& work) const
{
  int si = start->getIndex();
  int ei = diagonal.getSize();
  if (end != diag_end())
    ei = end->getIndex();
  savePartOfX(si, ei, x, work);

  for (auto di = start; di != end; ++di)
    if (int jbar {di->getIndex()}; di->isReal())
      {
        KronVector xi(x, jbar);
        xi.zeros();
        Vector wi(work, (jbar - si) * xi.length(), xi.length());
        xi.add(*(di->getAlpha()), wi);
        for (const_row_iter ri = row_begin(*di); ri != row_end(*di); ++ri)
          {
            int col = ri.getCol();
            Vector wj(work, (col - si) * xi.length(), xi.length());
            xi.add(*ri, wj);
          }
      }
    else
      {
        KronVector xi(x, jbar);
        KronVector xii(x, jbar + 1);
        xi.zeros();
        xii.zeros();
        Vector wi(work, (jbar - si) * xi.length(), xi.length());
        Vector wii(work, (jbar + 1 - si) * xi.length(), xi.length());
        xi.add(*(di->getAlpha()), wi);
        xi.add(di->getBeta1(), wii);
        xii.add(di->getBeta2(), wi);
        xii.add(*(di->getAlpha()), wii);
        for (const_row_iter ri = row_begin(*di); ri != row_end(*di); ++ri)
          {
            int col = ri.getCol();
            Vector wj(work, (col - si) * xi.length(), xi.length());
            xi.add(ri.a(), wj);
            xii.add(ri.b(), wj);
          }
      }
}

void
BlockDiagonal::multKronBlockTrans(const_diag_iter start, const_diag_iter end, KronVector& x,
                                  Vector& work) const
{
  int si = start->getIndex();
  int ei = diagonal.getSize();
  if (end != diag_end())
    ei = end->getIndex();
  savePartOfX(si, ei, x, work);

  for (auto di = start; di != end; ++di)
    if (int jbar {di->getIndex()}; di->isReal())
      {
        KronVector xi(x, jbar);
        xi.zeros();
        Vector wi(work, (jbar - si) * xi.length(), xi.length());
        xi.add(*(di->getAlpha()), wi);
        for (const_col_iter ci = col_begin(*di); ci != col_end(*di); ++ci)
          {
            int row = ci.getRow();
            Vector wj(work, (row - si) * xi.length(), xi.length());
            xi.add(*ci, wj);
          }
      }
    else
      {
        KronVector xi(x, jbar);
        KronVector xii(x, jbar + 1);
        xi.zeros();
        xii.zeros();
        Vector wi(work, (jbar - si) * xi.length(), xi.length());
        Vector wii(work, (jbar + 1 - si) * xi.length(), xi.length());
        xi.add(*(di->getAlpha()), wi);
        xi.add(di->getBeta2(), wii);
        xii.add(di->getBeta1(), wi);
        xii.add(*(di->getAlpha()), wii);
        for (const_col_iter ci = col_begin(*di); ci != col_end(*di); ++ci)
          {
            int row = ci.getRow();
            Vector wj(work, (row - si) * xi.length(), xi.length());
            xi.add(ci.a(), wj);
            xii.add(ci.b(), wj);
          }
      }
}

void
BlockDiagonal::multKron(KronVector& x) const
{
  int largest = getLargestBlock();
  Vector work(largest * x.getN() * power(x.getM(), x.getDepth() - 1));
  auto start = diag_begin();
  auto end = findBlockStart(start);
  while (start != diag_end())
    {
      multKronBlock(start, end, x, work);
      start = end;
      end = findBlockStart(start);
    }
}

void
BlockDiagonal::multKronTrans(KronVector& x) const
{
  int largest = getLargestBlock();
  Vector work(largest * x.getN() * power(x.getM(), x.getDepth() - 1));
  auto start = diag_begin();
  auto end = findBlockStart(start);
  while (start != diag_end())
    {
      multKronBlockTrans(start, end, x, work);
      start = end;
      end = findBlockStart(start);
    }
}

void
BlockDiagonal::printInfo() const
{
  std::cout << "Block sizes:";
  int num_blocks = 0;
  auto start = diag_begin();
  auto end = findBlockStart(start);
  while (start != diag_end())
    {
      int si = start->getIndex();
      int ei = diagonal.getSize();
      if (end != diag_end())
        ei = end->getIndex();
      std::cout << ' ' << ei - si;
      num_blocks++;
      start = end;
      end = findBlockStart(start);
    }
  std::cout << std::endl
            << "Num blocks: " << num_blocks << std::endl
            << "There are " << getNumZeros() << " zeros out of " << getNumOffdiagonal()
            << std::endl;
}

int
BlockDiagonal::getNumBlocks() const
{
  int num_blocks = 0;
  auto start = diag_begin();
  auto end = findBlockStart(start);
  while (start != diag_end())
    {
      num_blocks++;
      start = end;
      end = findBlockStart(start);
    }
  return num_blocks;
}
