/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2015  Virginia Tech
 *
 *  This program 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.
 *
 *  This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  justus
 *  Department of Chemistry, Virginia Tech
 *
 *  sparce.cpp
 *  Feb 5, 2015
 *
 */

#include <iostream>
#include <tiledarray.h>
#include <iomanip>

int main(int argc, char** argv) {
  int rc = 0;

  try {
    // Initialize runtime
    TiledArray::World& world = TiledArray::initialize(argc, argv);

    // Get command line arguments
    if(argc < 2) {
      std::cout << "Usage: ta_sparse matrix_size block_size [repetitions]\n";
      return 0;
    }
    const long matrix_size = atol(argv[1]);
    const long block_size = atol(argv[2]);
    if(matrix_size <= 0) {
      std::cerr << "Error: matrix size must be greater than zero.\n";
      return 1;
    }
    if(block_size <= 0) {
      std::cerr << "Error: block size must be greater than zero.\n";
      return 1;
    }
    if((matrix_size % block_size) != 0ul) {
      std::cerr << "Error: matrix size must be evenly divisible by block size.\n";
      return 1;
    }
    const long repeat = (argc >= 4 ? atol(argv[3]) : 4);
    if(repeat <= 0) {
      std::cerr << "Error: number of repetitions must be greater than zero.\n";
      return 1;
    }

    // Print information about the test
    const std::size_t num_blocks = matrix_size / block_size;
    const double app_flop = 2.0 * matrix_size * matrix_size * matrix_size;
    const float tile_norm = std::sqrt(float(block_size * block_size));
    std::vector<double> speeds, times, app_speeds, real_sparsity;

    if(world.rank() == 0)
      std::cout << "TiledArray: growing, block-sparse matrix multiply test..."
                << "\nGit HASH: " << TILEDARRAY_REVISION
                << "\nNumber of nodes     = " << world.size()
                << "\nBlock size          = " << block_size << "x" << block_size
                << "\nMemory per matrix   = " << double(matrix_size * matrix_size * sizeof(double)) / 1.0e9 << " GB"
                << "\nAverage blocks/node = " << double(num_blocks * num_blocks) / double(world.size()) << "\n";


    for(unsigned int sparsity = 100u; sparsity > 0u; sparsity -= 10u) {
      TiledArray::TSpArrayD::wait_for_lazy_cleanup(world);

      // Compute the number of blocks and matrix size for the sparse matrix
      const double sparse_fraction = double(sparsity) / 100.0;
      const long sparse_num_blocks =
          std::sqrt(double(num_blocks * num_blocks) / sparse_fraction);
      const long sparse_matrix_size = sparse_num_blocks * block_size;
      const long sparse_block_count = sparse_fraction * double(sparse_num_blocks * sparse_num_blocks);

      if(world.rank() == 0)
        std::cout << "\nSparsity = " << sparsity << "%"
            << "\nMatrix size = " << sparse_matrix_size << "x" << sparse_matrix_size << "\n";

      // Construct TiledRange
      std::vector<unsigned int> blocking;
      blocking.reserve(sparse_num_blocks + 1);
      for(long i = 0l; i <= sparse_matrix_size; i += block_size)
        blocking.push_back(i);

      const std::vector<TiledArray::TiledRange1> blocking2(2,
          TiledArray::TiledRange1(blocking.begin(), blocking.end()));

      const TiledArray::TiledRange
        trange(blocking2.begin(), blocking2.end());

      // Construct tile norm tensors
      TiledArray::Tensor<float>
          a_tile_norms(trange.tiles_range(), 0.0f),
          b_tile_norms(trange.tiles_range(), 0.0f);

      // Fill tile norm tensors
      if(world.rank() == 0) {
        if(sparsity == 100u) {
          std::fill(a_tile_norms.begin(), a_tile_norms.end(), tile_norm);
          std::fill(b_tile_norms.begin(), b_tile_norms.end(), tile_norm);
        } else {
          world.srand(time(NULL));
          for(long count = 0l; count < sparse_block_count; ++count) {
            // Find a new zero tile index.
            std::size_t index = 0ul;
            do {
              index = world.rand() % trange.tiles_range().volume();
            } while(a_tile_norms[index] > TiledArray::SparseShape<float>::threshold());

            // Set index tile of matrix matrix a.
            a_tile_norms[index] = tile_norm;

            // Find a new zero tile index.
            do {
              index = world.rand() % trange.tiles_range().volume();
            } while(b_tile_norms[index] > TiledArray::SparseShape<float>::threshold());

            b_tile_norms[index] = tile_norm;
          }
        }
      }

      // Construct the argument shapes
      TiledArray::SparseShape<float>
          a_shape(world, a_tile_norms, trange),
          b_shape(world, b_tile_norms, trange);

      // Construct and initialize arrays
      TiledArray::TSpArrayD a(world, trange, a_shape);
      TiledArray::TSpArrayD b(world, trange, b_shape);
      TiledArray::TSpArrayD c;
      a.fill_local(1.0);
      b.fill_local(1.0);

      // Start clock
      if(world.rank() == 0)
        std::cout << "Starting iterations:\n";

      double total_time = 0.0, flop = 0.0;

        // Do matrix multiplication
      for(int i = 0; i < repeat; ++i) {
        const double start = madness::wall_time();
        c("m,n") = a("m,k") * b("k,n");
        const double time = madness::wall_time() - start;
        total_time += time;
        if(flop < 1.0)
          flop = 2.0 * c("m,n").sum();
        if(world.rank() == 0)
          std::cout << "Iteration " << i + 1 << "   time=" << time << " s,   speed="
              << flop / time / 1.0e9 << " GFLOPS,   apparent speed=" << app_flop / time / 1.0e9 << " GFLOPS\n";
      }

      // Compute results
      speeds.push_back(double(repeat) * flop / total_time / 1.0e9);
      times.push_back(total_time / repeat);
      app_speeds.push_back(double(repeat) * app_flop / total_time / 1.0e9);
      real_sparsity.push_back(100.0 * double(sparse_block_count) / double(sparse_num_blocks * sparse_num_blocks));

      // Print results for this iteration
      if(world.rank() == 0) {
        std::cout << "\nSparsity               = " << real_sparsity.back() << "%\n"
                    << "Average wall time      = " << times.back() << " s\n"
                    << "Average speed          = " << speeds.back() << " GFLOPS\n"
                    << "Average apparent speed = " << app_speeds.back() << " GFLOPS\n";
      }
    }

    // Print out comma separated list of all results
    if(world.rank() == 0) {
      std::cout << "\n\nResults:\n"
          "sparsity (%), time (s), speed (GFLOPS), apparent speed (GFLOPS)\n";
      for(unsigned int i = 0; i < 10; ++i) {
        std::cout << real_sparsity[i] << ", " << times[i] << ", " << speeds[i] << ", " << app_speeds[i] << "\n";
      }
    }

    TiledArray::finalize();

  } catch(TiledArray::Exception& e) {
    std::cerr << "!! TiledArray exception: " << e.what() << "\n";
    rc = 1;
  } catch(madness::MadnessException& e) {
    std::cerr << "!! MADNESS exception: " << e.what() << "\n";
    rc = 1;
  } catch(SafeMPI::Exception& e) {
    std::cerr << "!! SafeMPI exception: " << e.what() << "\n";
    rc = 1;
  } catch(std::exception& e) {
    std::cerr << "!! std exception: " << e.what() << "\n";
    rc = 1;
  } catch(...) {
    std::cerr << "!! exception: unknown exception\n";
    rc = 1;
  }

  return rc;
}
