#include <stdio.h>
#include <stdint.h>
#include <inttypes.h>
#include <ctype.h>
#include <unistd.h>

#include <opencv2/opencv.hpp>

#include <AprilTags/apriltag.h>
#include <AprilTags/common/image_u8.h>
#include <AprilTags/tag36h11.h>
#include <AprilTags/common/zarray.h>
#include <AprilTags/common/getopt.h>

#include <drc_utils/LcmWrapper.hpp>
#include <drc_utils/BotWrapper.hpp>
#include <lcm/lcm-cpp.hpp>

#include <lcmtypes/bot_core.hpp>
#include <lcmtypes/bot_core/images_t.hpp>
#include <lcmtypes/bot_core/image_t.hpp>

#include <bot_core/camtrans.h>
#include <bot_param/param_util.h>
#include <bot_frames_cpp/bot_frames_cpp.hpp>

#include <vector>
#include <iostream>

#include <Eigen/Dense>

struct TagMatch {
    int id;
    cv::Point2d p0, p1, p2, p3;
    Eigen::Matrix3d H;
};

Eigen::Isometry3d getRelativeTransform(TagMatch const& match, Eigen::Matrix3d const & K, double tag_size)
{
  std::vector<cv::Point3f> objPts;
  std::vector<cv::Point2f> imgPts;
  double s = tag_size/2.;
  objPts.push_back(cv::Point3f(-s,-s, 0));
  objPts.push_back(cv::Point3f( s,-s, 0));
  objPts.push_back(cv::Point3f( s, s, 0));
  objPts.push_back(cv::Point3f(-s, s, 0));

  imgPts.push_back(match.p0);
  imgPts.push_back(match.p1);
  imgPts.push_back(match.p2);
  imgPts.push_back(match.p3);

  cv::Mat rvec, tvec;
  cv::Matx33f cameraMatrix(
                           K(0,0), 0, K(0,2),
                           0, K(1,1), K(1,2),
                           0,  0,  1);

  cv::Vec4f distParam(0,0,0,0);
  cv::solvePnP(objPts, imgPts, cameraMatrix, distParam, rvec, tvec);
  cv::Matx33d r;
  cv::Rodrigues(rvec, r);
  Eigen::Matrix3d wRo;
  wRo << r(0,0), r(0,1), r(0,2), r(1,0), r(1,1), r(1,2), r(2,0), r(2,1), r(2,2);

  Eigen::Isometry3d T;
  T.linear() = wRo;
  T.translation() << tvec.at<double>(0), tvec.at<double>(1), tvec.at<double>(2);
  return T;
}

bot_core::rigid_transform_t encodeLCMFrame(Eigen::Isometry3d const & frame)
{
    Eigen::Vector3d t(frame.translation());
    Eigen::Quaterniond r(frame.rotation());

    bot_core::rigid_transform_t msg;
    msg.quat[0] = r.w();
    msg.quat[1] = r.x();
    msg.quat[2] = r.y();
    msg.quat[3] = r.z();

    msg.trans[0] = t[0];
    msg.trans[1] = t[1];
    msg.trans[2] = t[2];

    return msg;
}

class AprilTagDetector {
    public:
    AprilTagDetector(getopt_t *options) : getopt(options) {
        tf = tag36h11_create();
        tf->black_border = getopt_get_int(options, "border");
        td = apriltag_detector_create();
        apriltag_detector_add_family(td, tf);

        td->quad_decimate = getopt_get_double(getopt, "decimate");
        td->quad_sigma = getopt_get_double(getopt, "blur");
        td->nthreads = getopt_get_int(getopt, "threads");
        td->debug = getopt_get_bool(getopt, "debug");
        td->refine_edges = getopt_get_bool(getopt, "refine-edges");
        td->refine_decode = getopt_get_bool(getopt, "refine-decode");
        td->refine_pose = getopt_get_bool(getopt, "refine-pose");

        quiet = getopt_get_bool(getopt, "quiet");
        tag_size = getopt_get_double(getopt, "size");
    }

    ~AprilTagDetector() {

        apriltag_detector_destroy(td);
        tag36h11_destroy(tf);
    }

    std::vector<TagMatch> detectTags(image_u8_t *im) {

        const int hamm_hist_max = 10;

        int hamm_hist[hamm_hist_max];
        memset(hamm_hist, 0, sizeof(hamm_hist));
        zarray_t *detections = apriltag_detector_detect(td, im);

        std::vector<TagMatch> tag_matches;
        for (int i = 0; i < zarray_size(detections); i++) {
            apriltag_detection_t *det;
            zarray_get(detections, i, &det);

            if (!quiet)
                printf("detection %3d: id (%2dx%2d)-%-4d, hamming %d, goodness %8.3f, margin %8.3f\n",
                       i, det->family->d*det->family->d, det->family->h, det->id, det->hamming, det->goodness, det->decision_margin);

            for (int x = 0; x < 3 ; x++) {
                image_u8_draw_line(im, det->p[x][0], det->p[x][1], det->p[x+1][0], det->p[x+1][1], 255, 10);
            }
            TagMatch tag_match;
            tag_match.id = det->family->d*det->family->d;
            tag_match.p0 = cv::Point2d(det->p[0][0], det->p[0][1]);
            tag_match.p1 = cv::Point2d(det->p[1][0], det->p[1][1]);
            tag_match.p2 = cv::Point2d(det->p[2][0], det->p[2][1]);
            tag_match.p3 = cv::Point2d(det->p[3][0], det->p[3][1]);

            Eigen::Map<Eigen::Matrix3d> H_map(det->H->data);
            tag_match.H = H_map.transpose();
            tag_matches.push_back(tag_match);
            hamm_hist[det->hamming]++;
        }

        apriltag_detections_destroy(detections);

        if (!quiet) {
            timeprofile_display(td->tp);
            printf("nedges: %d, nsegments: %d, nquads: %d\n", td->nedges, td->nsegments, td->nquads);
            printf("Hamming histogram: ");
            for (int i = 0; i < hamm_hist_max; i++)
                printf("%5d", hamm_hist[i]);
            printf("%12.3f", timeprofile_total_utime(td->tp) / 1.0E3);
            printf("\n");
        }

        return tag_matches;
    }

    double getTagSize() const {
        return tag_size;
    }

    private:
    int quiet;
    double tag_size;
    apriltag_family_t *tf;
    apriltag_detector_t *td;
    getopt_t *getopt;
};

class CameraListener {
    public:

    void setDetector(AprilTagDetector* detector) {
        mDetector = detector;
    }

    bool setup(bool show_window) {
        mBotWrapper.reset(new drc::BotWrapper());

        while (!mBotWrapper->getBotParam()) {
            std::cout << "Re-trying ... " << std::endl;
            mBotWrapper->setDefaults();
        }

        mLcmWrapper.reset(new drc::LcmWrapper(mBotWrapper->getLcm()));
        mLcmWrapper->get()->subscribe("CAMERA", &CameraListener::onCamera, this);

        mCamTransLeft = bot_param_get_new_camtrans(mBotWrapper->getBotParam(),"CAMERA_LEFT");

        K = Eigen::Matrix3d::Identity();

        K(0,0) = bot_camtrans_get_focal_length_x(mCamTransLeft);
        K(1,1) = bot_camtrans_get_focal_length_y(mCamTransLeft);
        K(0,2) = bot_camtrans_get_principal_x(mCamTransLeft);
        K(1,2) = bot_camtrans_get_principal_y(mCamTransLeft);

        mShowWindow = show_window;
        return true;
    }

    void start() {
        mLcmWrapper->startHandleThread(true);
    }

    void onCamera(const lcm::ReceiveBuffer* buffer, const std::string& channel,
                const bot_core::images_t* msg) {
        cv::Mat image;
        decodeImage(msg, image);
        image_u8_t *image_u8 = fromCvMat(image);

        std::vector<TagMatch> tags = mDetector->detectTags(image_u8);
        cv::cvtColor(image, image, CV_GRAY2RGB);
        for (int i = 0; i < tags.size(); i++) {

            if (mShowWindow) {
                cv::line(image, tags[i].p0, tags[i].p1, cv::Scalar(255,0,0), 2, CV_AA);
                cv::line(image, tags[i].p1, tags[i].p2, cv::Scalar(0,255,0), 2, CV_AA);
                cv::line(image, tags[i].p2, tags[i].p3, cv::Scalar(0,0,255), 2, CV_AA);
                cv::line(image, tags[i].p3, tags[i].p0, cv::Scalar(0,0,255), 2, CV_AA);

                Eigen::Vector3d x_axis(2,0,1);
                Eigen::Vector3d y_axis(0,2,1);
                Eigen::Vector3d origin(0,0,1);

                Eigen::Vector3d px = tags[i].H * x_axis;
                Eigen::Vector3d py = tags[i].H * y_axis;
                Eigen::Vector3d o  = tags[i].H * origin;

                px/= px[2];
                py/= py[2];
                o/= o[2];

                cv::line(image, cv::Point2d(o[0], o[1]), cv::Point2d(px[0], px[1]), cv::Scalar(255,0,255), 1, CV_AA);
                cv::line(image, cv::Point2d(o[0], o[1]), cv::Point2d(py[0], py[1]), cv::Scalar(255,255,0), 1, CV_AA);
            }

            Eigen::Isometry3d tag_to_camera = getRelativeTransform(tags[i], K, mDetector->getTagSize());
            bot_core::rigid_transform_t tag_to_camera_msg = encodeLCMFrame(tag_to_camera);
            tag_to_camera_msg.utime = msg->utime;
            mLcmWrapper->get()->publish("APRIL_TAG_TO_CAMERA_LEFT", &tag_to_camera_msg);
            break;

        }
        if (mShowWindow) {
            cv::imshow("detections", image);
            cv::waitKey(1);
        }

        image_u8_destroy(image_u8);
    }

    void decodeImage(const bot_core::images_t* msg, cv::Mat & decoded_image) {
        bot_core::image_t* leftImage = NULL;
        for (int i = 0; i < msg->n_images; ++i) {
          if (msg->image_types[i] == bot_core::images_t::LEFT) {
            leftImage = (bot_core::image_t*)(&msg->images[i]);
            decoded_image = leftImage->pixelformat == leftImage->PIXEL_FORMAT_MJPEG ?
                cv::imdecode(cv::Mat(leftImage->data), -1) :
                cv::Mat(leftImage->height, leftImage->width, CV_8UC1, leftImage->data.data());
                if (decoded_image.channels() > 1) {
                    cv::cvtColor(decoded_image, decoded_image, CV_RGB2GRAY);
                }
                break;
          }
        }
    }

    image_u8_t *fromCvMat(const cv::Mat & img) {
        image_u8_t *image_u8 = image_u8_create_alignment(img.cols, img.rows, img.step);
        int size = img.total() * img.elemSize();
        memcpy(image_u8->buf, img.data, size * sizeof(uint8_t));
        return image_u8;
    }

    private:
    bool mShowWindow;
    AprilTagDetector *mDetector;
    drc::LcmWrapper::Ptr mLcmWrapper;
    drc::BotWrapper::Ptr mBotWrapper;
    BotCamTrans* mCamTransLeft;
    Eigen::Matrix3d K;
};

int main(int argc, char *argv[])
{
    getopt_t *getopt = getopt_create();

    getopt_add_bool(getopt, 'h', "help", 0, "Show this help");
    getopt_add_bool(getopt, 'd', "debug", 0, "Enable debugging output (slow)");
    getopt_add_bool(getopt, 'w', "window", 1, "Show the detected tags in a window");
    getopt_add_bool(getopt, 'q', "quiet", 0, "Reduce output");
    getopt_add_int(getopt, '\0', "border", "1", "Set tag family border size");
    getopt_add_int(getopt, 't', "threads", "4", "Use this many CPU threads");
    getopt_add_double(getopt, 'x', "decimate", "1.0", "Decimate input image by this factor");
    getopt_add_double(getopt, 'b', "blur", "0.0", "Apply low-pass blur to input");
    getopt_add_bool(getopt, '0', "refine-edges", 1, "Spend more time trying to align edges of tags");
    getopt_add_bool(getopt, '1', "refine-decode", 0, "Spend more time trying to decode tags");
    getopt_add_bool(getopt, '2', "refine-pose", 0, "Spend more time trying to precisely localize tags");
    getopt_add_double(getopt, 's', "size", "0.1735", "Physical side-length of the tag (meters)");

    if (!getopt_parse(getopt, argc, argv, 1) || getopt_get_bool(getopt, "help")) {
        printf("Usage: %s [options]\n", argv[0]);
        getopt_do_usage(getopt);
        exit(0);
    }

    AprilTagDetector tag_detector(getopt);
    CameraListener camera_listener;

    if (camera_listener.setup(getopt_get_bool(getopt, "window"))) {
        camera_listener.setDetector(&tag_detector);
        camera_listener.start();
    }

    return 0;
}


#define BOOST_TEST_MODULE "test_harmonic_potential"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <test/util/check_potential.hpp>
#include <mjolnir/forcefield/local/HarmonicPotential.hpp>

BOOST_AUTO_TEST_CASE(Harmonic_double)
{
    using real_type = double;
    constexpr std::size_t N = 1000;
    constexpr real_type   h = 1e-6;
    constexpr real_type tol = 1e-6;

    const real_type k  = 1.0;
    const real_type r0 = 5.0;

    mjolnir::HarmonicPotential<real_type> harmonic(k, r0);

    const real_type x_min = 0.5 * r0;
    const real_type x_max = 1.5 * r0;

    mjolnir::test::check_potential(harmonic, x_min, x_max, tol, h, N);
}

BOOST_AUTO_TEST_CASE(Harmonic_float)
{
    namespace test = mjolnir::test;
    using real_type = float;
    constexpr std::size_t N = 100;
    constexpr real_type   h = 1e-3f;
    constexpr real_type tol = 1e-3f;

    const real_type k  = 1.0f;
    const real_type r0 = 5.0f;

    mjolnir::HarmonicPotential<real_type> harmonic(k, r0);

    const real_type x_min = 0.5 * r0;
    const real_type x_max = 1.5 * r0;

    mjolnir::test::check_potential(harmonic, x_min, x_max, tol, h, N);
}


#include "pressio_apps.hpp"
#include <random>
#include <Eigen/SVD>

using gen_t	   = std::mt19937;
using rand_distr_t = std::uniform_real_distribution<double>;

constexpr double eps = 1e-7;
std::string checkStr {"PASSED"};

// range of Prandtl number
constexpr std::array<double,2> Pr_range{{1.0, 5.0}};

// range of Reynolds number
constexpr std::array<double,2> Re_range{{10., 100.0}};

void readMatrixFromFile(std::string filename,
			std::vector<std::vector<double>> & A0,
			int ncols){
  assert( A0.empty() );
  std::ifstream source;
  source.open( filename, std::ios_base::in);
  std::string line, colv;
  std::vector<double> tmpv(ncols);
  while (std::getline(source, line) ){
    std::istringstream in(line);
    for (int i=0; i<ncols; i++){
      in >> colv;
      tmpv[i] = atof(colv.c_str());
    }
    A0.emplace_back(tmpv);
  }
  source.close();
}

int main(int argc, char *argv[]){
  using fom_t	 = ::pressio::apps::SteadyLinAdvDiff2dEpetra;

  int rank;
  MPI_Init(&argc,&argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  Epetra_MpiComm Comm(MPI_COMM_WORLD);
  assert(Comm.NumProc() == 1);

  //---------------------------------------
  // generate random samples of Parameters

  // random number generator (seeded)
  unsigned int seed = 1343234343;
  std::mt19937 engine(seed);
  rand_distr_t distr(0., 1.0);
  auto genPr = [&distr, &engine](){
		 auto c1 = Pr_range[1]-Pr_range[0];
		 auto c2 = Pr_range[0];
		 return c1 * distr(engine) + c2;
	     };

  auto genRe = [&distr, &engine](){
		 auto c1 = Re_range[1]-Re_range[0];
		 auto c2 = Re_range[0];
		 return c1 * distr(engine) + c2;
	     };

  // number of sample to take
  constexpr int nSamples = 5;

  // fill for Prandtl
  std::vector<double> PrS(nSamples);
  std::generate(PrS.begin(), PrS.end(), genPr);

  // fill for Reynolds
  std::vector<double> ReS(nSamples);
  std::generate(ReS.begin(), ReS.end(), genRe);

  if(rank==0){
    auto it1 = PrS.begin();
    auto it2 = ReS.begin();
    for( ;it2<ReS.end(); it1++, it2++)
      std::cout << std::setprecision(15)
		<< *it1 << " " << *it2
		<< "\n";
  }

  //---------------------------------------
  // fix discretization for all samples
  const int Nx = 11, Ny = Nx*2-1;

  /* # of dofs is != Nx*Ny because of how we solve pdd */
  const int numDof = (Nx-2)*Ny;

  // create as many app objects as samples
  std::vector<fom_t> vecObjs;
  for (auto i=0; i<nSamples; i++){
    vecObjs.emplace_back(Comm,Nx, Ny, PrS[i], ReS[i]);
  }

  // solve all problems
  for (auto & it : vecObjs){
    it.assembleMatrix();
    it.fillRhs();
    it.solve();
  }

  // collect all solutions into matrix
  // I can do this this easily because we know # ranks = 1
  using eig_mat = Eigen::MatrixXd;
  eig_mat A(numDof, nSamples);
  int j=0;
  for (const auto & it : vecObjs){
    auto T = it.getState();
    for (auto i=0; i<numDof; i++)
      A(i,j) = (*T)[i];
    j++;
  }

  // do SVD
  Eigen::JacobiSVD<eig_mat> svd(A, Eigen::ComputeThinU);
  auto U = svd.matrixU();
  std::cout << std::setprecision(15) << U << std::endl;

  // read gold basis from file
  std::vector<std::vector<double>> goldU;
  readMatrixFromFile("gold_basis.txt", goldU, nSamples);

  // check that computed matches gold
  assert( (size_t) goldU.size() == (size_t) U.rows() );
  for (auto i=0; i<U.rows(); i++)
    for (j=0; j<nSamples; j++)
      if ( std::abs(goldU[i][j] - U(i,j)) > eps ) checkStr = "FAILED";

  MPI_Finalize();
  std::cout << checkStr <<  std::endl;
  return 0;
}


// demo from:
//      https://stackoverflow.com/questions/47653386/how-to-use-boostgeometryrtree-with-glmvec3-as-a-custom-point-type
#include <boost/geometry.hpp>
#include <boost/geometry/index/rtree.hpp>
#include <boost/geometry/geometries/box.hpp>
#include <boost/geometry/geometries/register/point.hpp>

namespace bg = boost::geometry;
namespace bgi = boost::geometry::index;

namespace glm
{
    struct vec3 {
        double x, y, z;
        vec3(double x_, double y_, double z_) { x = x_; y = y_; z = z_; }
    };
}

BOOST_GEOMETRY_REGISTER_POINT_3D(glm::vec3, double, bg::cs::cartesian, x, y, z)

#include <iostream>
int main() {
    using IndexedPoint = std::pair<glm::vec3, uint32_t>;
    using RTree = boost::geometry::index::rtree<IndexedPoint, boost::geometry::index::rstar<8>>;

    RTree rtree;
    rtree.insert({glm::vec3(1,1,1), 1});
    rtree.insert({glm::vec3(2,2,2), 2});
    rtree.insert({glm::vec3(3,3,3), 3});
    rtree.insert({glm::vec3(4,4,4), 4});

    auto q = bgi::nearest(glm::vec3(2.9, 2.9, 2.9), 99);

    auto it = rtree.qbegin(q);
    auto p = it->first;
    std::cout << "Nearest: # " << it->second << " (" << p.x << ", " << p.y << " " << p.z << ")\n";
}


#include "neural_network.h"

#include <cuda_runtime.h>
#include <helper_cuda.h>
#include <helper_functions.h>

#include <armadillo>

#include "cublas_v2.h"
#include "gpu_func.h"
#include "iomanip"
#include "mpi.h"
#include "utils/common.h"

#define MPI_SAFE_CALL(call)                                                  \
  do {                                                                       \
    int err = call;                                                          \
    if (err != MPI_SUCCESS) {                                                \
      fprintf(stderr, "MPI error %d in file '%s' at line %i", err, __FILE__, \
              __LINE__);                                                     \
      exit(1);                                                               \
    }                                                                        \
  } while (0)

real norms(NeuralNetwork& nn) {
  real norm_sum = 0;

  for (int i = 0; i < nn.num_layers; ++i) {
    norm_sum += arma::accu(arma::square(nn.W[i]));
  }

  return norm_sum;
}

void write_cpudata_tofile(NeuralNetwork& nn, int iter) {
  std::stringstream s;
  s << "Outputs/CPUmats/SequentialW0-" << iter << ".mat";
  nn.W[0].save(s.str(), arma::raw_ascii);
  std::stringstream t;
  t << "Outputs/CPUmats/SequentialW1-" << iter << ".mat";
  nn.W[1].save(t.str(), arma::raw_ascii);
  std::stringstream u;
  u << "Outputs/CPUmats/Sequentialb0-" << iter << ".mat";
  nn.b[0].save(u.str(), arma::raw_ascii);
  std::stringstream v;
  v << "Outputs/CPUmats/Sequentialb1-" << iter << ".mat";
  nn.b[1].save(v.str(), arma::raw_ascii);
}

void write_diff_gpu_cpu(NeuralNetwork& nn, int iter,
                        std::ofstream& error_file) {
  arma::Mat<real> A, B, C, D;

  std::stringstream s;
  s << "Outputs/CPUmats/SequentialW0-" << iter << ".mat";
  A.load(s.str(), arma::raw_ascii);
  real max_errW0 = arma::norm(nn.W[0] - A, "inf") / arma::norm(A, "inf");
  real L2_errW0 = arma::norm(nn.W[0] - A, 2) / arma::norm(A, 2);

  std::stringstream t;
  t << "Outputs/CPUmats/SequentialW1-" << iter << ".mat";
  B.load(t.str(), arma::raw_ascii);
  real max_errW1 = arma::norm(nn.W[1] - B, "inf") / arma::norm(B, "inf");
  real L2_errW1 = arma::norm(nn.W[1] - B, 2) / arma::norm(B, 2);

  std::stringstream u;
  u << "Outputs/CPUmats/Sequentialb0-" << iter << ".mat";
  C.load(u.str(), arma::raw_ascii);
  real max_errb0 = arma::norm(nn.b[0] - C, "inf") / arma::norm(C, "inf");
  real L2_errb0 = arma::norm(nn.b[0] - C, 2) / arma::norm(C, 2);

  std::stringstream v;
  v << "Outputs/CPUmats/Sequentialb1-" << iter << ".mat";
  D.load(v.str(), arma::raw_ascii);
  real max_errb1 = arma::norm(nn.b[1] - D, "inf") / arma::norm(D, "inf");
  real L2_errb1 = arma::norm(nn.b[1] - D, 2) / arma::norm(D, 2);

  int ow = 15;

  if (iter == 0) {
    error_file << std::left << std::setw(ow) << "Iteration" << std::left
               << std::setw(ow) << "Max Err W0" << std::left << std::setw(ow)
               << "Max Err W1" << std::left << std::setw(ow) << "Max Err b0"
               << std::left << std::setw(ow) << "Max Err b1" << std::left
               << std::setw(ow) << "L2 Err W0" << std::left << std::setw(ow)
               << "L2 Err W1" << std::left << std::setw(ow) << "L2 Err b0"
               << std::left << std::setw(ow) << "L2 Err b1"
               << "\n";
  }

  error_file << std::left << std::setw(ow) << iter << std::left << std::setw(ow)
             << max_errW0 << std::left << std::setw(ow) << max_errW1
             << std::left << std::setw(ow) << max_errb0 << std::left
             << std::setw(ow) << max_errb1 << std::left << std::setw(ow)
             << L2_errW0 << std::left << std::setw(ow) << L2_errW1 << std::left
             << std::setw(ow) << L2_errb0 << std::left << std::setw(ow)
             << L2_errb1 << "\n";
}

/* CPU IMPLEMENTATIONS */
void feedforward(NeuralNetwork& nn, const arma::Mat<real>& X,
                 struct cache& cache) {
  cache.z.resize(2);
  cache.a.resize(2);

  // std::cout << W[0].n_rows << "\n";tw
  assert(X.n_rows == nn.W[0].n_cols);
  cache.X = X;
  int N = X.n_cols;

  arma::Mat<real> z1 = nn.W[0] * X + arma::repmat(nn.b[0], 1, N);
  cache.z[0] = z1;

  arma::Mat<real> a1;
  sigmoid(z1, a1);
  cache.a[0] = a1;

  assert(a1.n_rows == nn.W[1].n_cols);
  arma::Mat<real> z2 = nn.W[1] * a1 + arma::repmat(nn.b[1], 1, N);
  cache.z[1] = z2;

  arma::Mat<real> a2;
  softmax(z2, a2);
  cache.a[1] = cache.yc = a2;
}

/*
 * Computes the gradients of the cost w.r.t each param.
 * MUST be called after feedforward since it uses the bpcache.
 * @params y : C x N one-hot column vectors
 * @params bpcache : Output of feedforward.
 * @params bpgrads: Returns the gradients for each param
 */
void backprop(NeuralNetwork& nn, const arma::Mat<real>& y, real reg,
              const struct cache& bpcache, struct grads& bpgrads) {
  bpgrads.dW.resize(2);
  bpgrads.db.resize(2);
  int N = y.n_cols;

  // std::cout << "backprop " << bpcache.yc << "\n";
  arma::Mat<real> diff = (1.0 / N) * (bpcache.yc - y);
  bpgrads.dW[1] = diff * bpcache.a[0].t() + reg * nn.W[1];
  bpgrads.db[1] = arma::sum(diff, 1);
  arma::Mat<real> da1 = nn.W[1].t() * diff;

  arma::Mat<real> dz1 = da1 % bpcache.a[0] % (1 - bpcache.a[0]);

  bpgrads.dW[0] = dz1 * bpcache.X.t() + reg * nn.W[0];
  bpgrads.db[0] = arma::sum(dz1, 1);
}

/*
 * Computes the Cross-Entropy loss function for the neural network.
 */
real loss(NeuralNetwork& nn, const arma::Mat<real>& yc,
          const arma::Mat<real>& y, real reg) {
  int N = yc.n_cols;
  real ce_sum = -arma::accu(arma::log(yc.elem(arma::find(y == 1))));

  real data_loss = ce_sum / N;
  real reg_loss = 0.5 * reg * norms(nn);
  real loss = data_loss + reg_loss;
  // std::cout << "Loss: " << loss << "\n";
  return loss;
}

/*
 * Returns a vector of labels for each row vector in the input
 */
void predict(NeuralNetwork& nn, const arma::Mat<real>& X,
             arma::Row<real>& label) {
  struct cache fcache;
  feedforward(nn, X, fcache);
  label.set_size(X.n_cols);

  for (int i = 0; i < X.n_cols; ++i) {
    arma::uword row;
    fcache.yc.col(i).max(row);
    label(i) = row;
  }
}

/*
 * Computes the numerical gradient
 */
void numgrad(NeuralNetwork& nn, const arma::Mat<real>& X,
             const arma::Mat<real>& y, real reg, struct grads& numgrads) {
  real h = 0.00001;
  struct cache numcache;
  numgrads.dW.resize(nn.num_layers);
  numgrads.db.resize(nn.num_layers);

  for (int i = 0; i < nn.num_layers; ++i) {
    numgrads.dW[i].resize(nn.W[i].n_rows, nn.W[i].n_cols);

    for (int j = 0; j < nn.W[i].n_rows; ++j) {
      for (int k = 0; k < nn.W[i].n_cols; ++k) {
        real oldval = nn.W[i](j, k);
        nn.W[i](j, k) = oldval + h;
        feedforward(nn, X, numcache);
        real fxph = loss(nn, numcache.yc, y, reg);
        nn.W[i](j, k) = oldval - h;
        feedforward(nn, X, numcache);
        real fxnh = loss(nn, numcache.yc, y, reg);
        numgrads.dW[i](j, k) = (fxph - fxnh) / (2 * h);
        nn.W[i](j, k) = oldval;
      }
    }
  }

  for (int i = 0; i < nn.num_layers; ++i) {
    numgrads.db[i].resize(nn.b[i].n_rows, nn.b[i].n_cols);

    for (int j = 0; j < nn.b[i].size(); ++j) {
      real oldval = nn.b[i](j);
      nn.b[i](j) = oldval + h;
      feedforward(nn, X, numcache);
      real fxph = loss(nn, numcache.yc, y, reg);
      nn.b[i](j) = oldval - h;
      feedforward(nn, X, numcache);
      real fxnh = loss(nn, numcache.yc, y, reg);
      numgrads.db[i](j) = (fxph - fxnh) / (2 * h);
      nn.b[i](j) = oldval;
    }
  }
}

/*
 * Train the neural network nn
 */
void train(NeuralNetwork& nn, const arma::Mat<real>& X,
           const arma::Mat<real>& y, real learning_rate, real reg,
           const int epochs, const int batch_size, bool grad_check,
           int print_every, int debug) {
  int N = X.n_cols;
  int iter = 0;
  int print_flag = 0;

  for (int epoch = 0; epoch < epochs; ++epoch) {
    int num_batches = (N + batch_size - 1) / batch_size;

    for (int batch = 0; batch < num_batches; ++batch) {
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      arma::Mat<real> X_batch = X.cols(batch * batch_size, last_col);
      arma::Mat<real> y_batch = y.cols(batch * batch_size, last_col);

      struct cache bpcache;
      feedforward(nn, X_batch, bpcache);

      struct grads bpgrads;
      backprop(nn, y_batch, reg, bpcache, bpgrads);

      if (print_every > 0 && iter % print_every == 0) {
        if (grad_check) {
          struct grads numgrads;
          numgrad(nn, X_batch, y_batch, reg, numgrads);
          assert(gradcheck(numgrads, bpgrads));
        }

        std::cout << "Loss at iteration " << iter << " of epoch " << epoch
                  << "/" << epochs << " = "
                  << loss(nn, bpcache.yc, y_batch, reg) << "\n";
      }

      // Gradient descent step
      for (int i = 0; i < nn.W.size(); ++i) {
        nn.W[i] -= learning_rate * bpgrads.dW[i];
      }

      for (int i = 0; i < nn.b.size(); ++i) {
        nn.b[i] -= learning_rate * bpgrads.db[i];
      }

      /* Debug routine runs only when debug flag is set. If print_every is zero,
         it saves for the first batch of each epoch to avoid saving too many
         large files. Note that for the first time, you have to run debug and
         serial modes together. This will run the following function and write
         out files to CPUmats folder. In the later runs (with same parameters),
         you can use just the debug flag to
         output diff b/w CPU and GPU without running CPU version */
      if (print_every <= 0) {
        print_flag = batch == 0;
      } else {
        print_flag = iter % print_every == 0;
      }

      if (debug && print_flag) {
        write_cpudata_tofile(nn, iter);
      }

      iter++;
    }
  }
}

/*
 * TODO
 * Train the neural network nn of rank 0 in parallel. Your MPI implementation
 * should mainly be in this function.
 */
struct NNcache{
 real *W1, *W2, *b1, *b2;
 real *a1, *y_pred;
 real *dW1, *dW2, *db1, *db2;
 real *da1, *dz1, *diff;// diff derivative of cross entropy
  /*
  M:number of featrues,
  H:number of neurons in hidden layer,
  C:number of classes, 10
  */
 NNcache(int M, int H, int C, int batch_size){
   cudaMalloc((void**)&W1,sizeof(real)*H*M);
   cudaMalloc((void**)&W2,sizeof(real)*C*H);
   cudaMalloc((void**)&b1,sizeof(real)*H);
   cudaMalloc((void**)&b2,sizeof(real)*C);
   cudaMalloc((void**)&a1,sizeof(real)*H*batch_size);
   cudaMalloc((void**)&y_pred,sizeof(real)*C*batch_size);
   cudaMalloc((void**)&diff,sizeof(real)*C*batch_size);
   cudaMalloc((void**)&dW1,sizeof(real)*H*M);
   cudaMalloc((void**)&dW2,sizeof(real)*C*H);
   cudaMalloc((void**)&db1,sizeof(real)*H);
   cudaMalloc((void**)&db2,sizeof(real)*C);
   cudaMalloc((void**)&da1,sizeof(real)*H*batch_size);
   cudaMalloc((void**)&dz1,sizeof(real)*H*batch_size);
 }

 ~NNcache(){
   cudaFree(W1);
   cudaFree(W2);
   cudaFree(b1);
   cudaFree(b2);
   cudaFree(a1);
   cudaFree(y_pred);
   cudaFree(diff);
   cudaFree(dW1);
   cudaFree(dW2);
   cudaFree(db1);
   cudaFree(db2);
   cudaFree(da1);
   cudaFree(dz1);
 }
};

void parallel_feedforward(NeuralNetwork &nn, real *d_X, NNcache &nncache, int size_per_proc){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];
    real alpha = 1.0, beta = 1.0;

    /*layer 1  z1 = W1 * X + arma::repmat(b1, 1, N); a1 = sigmoid(z1)*/
    gpu_repmat(nncache.b1, nncache.a1, H, size_per_proc);
    myGEMM(nncache.W1, d_X, nncache.a1, &alpha,&beta, H, size_per_proc, M );
    gpu_sigmoid(nncache.a1,H,size_per_proc);
    /*layer 2 z2 = W2 * a1 + arma::repmat(b2, 1, N); y_pred = a2 = softmax(z2)*/
    gpu_repmat(nncache.b2, nncache.y_pred, C, size_per_proc);
    myGEMM(nncache.W2, nncache.a1, nncache.y_pred, &alpha,&beta, C, size_per_proc, H);
    gpu_softmax(nncache.y_pred, C, size_per_proc);
}

void parallel_backprop(NeuralNetwork& nn, real *d_X, real *d_Y, real reg, NNcache &nncache,int batch_size, int size_per_proc, int num_procs){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];
    real ratio = 1.0/(real)batch_size;
    reg = reg /num_procs; //change it in the parallel_train functino
    //reg = 0.0;
    /*diff = (1.0 / N) * (bpcache.yc - y)*/
    gpu_addmat(nncache.y_pred,d_Y,nncache.diff, ratio, -ratio, C, size_per_proc);

    /*bpgrads.dW[2] = diff * bpcache.a[1].t() + reg * nn.W[2];*/
    real alpha = 1.0;
    cudaMemcpy(nncache.dW2, nncache.W2, sizeof(real) * C * H, cudaMemcpyDeviceToDevice);
    myGEMMT(nncache.diff,nncache.a1,nncache.dW2,&alpha,&reg,C,H,size_per_proc,false,true);

    /*db2 = arma::sum(diff, 1)*/
    gpu_row_sum(nncache.diff, nncache.db2, C, size_per_proc);

    /* da1 = nncache.W2.t() * diff;*/
    real beta = 0.0;
    myGEMMT(nncache.W2,nncache.diff,nncache.da1,&alpha,&beta,H,size_per_proc,C,true,false);

    /* dz1 = da1 .* nncache.a1 .* (1 - nncache.a1);*/
    gpu_sigmoid_backprop(nncache.da1,nncache.a1,nncache.dz1,H,size_per_proc);

    /*dW1 = dz1 * X.t() + reg * nncache.W1;*/
    cudaMemcpy(nncache.dW1, nncache.W1, sizeof(real) * H * M, cudaMemcpyDeviceToDevice);
    myGEMMT(nncache.dz1, d_X, nncache.dW1, &alpha, &reg, H, M, size_per_proc,false,true);

    /* db1 = arma::sum(dz1, 1);*/
    gpu_row_sum(nncache.dz1, nncache.db1, H, size_per_proc);
}

void parallel_gradientdecent(NeuralNetwork& nn, NNcache &nncache, real learning_rate){
    int M = nn.H[0];
    int H = nn.H[1];
    int C = nn.H[2];

    //update nncache params can change two mat params in-place
    gpu_addmat(nncache.W1,nncache.dW1,nncache.W1,1.0,-learning_rate,H,M);
    gpu_addmat(nncache.W2,nncache.dW2,nncache.W2,1.0,-learning_rate,C,H);
    gpu_addmat(nncache.b1,nncache.db1,nncache.b1,1.0,-learning_rate,H,1);
    gpu_addmat(nncache.b2,nncache.db2,nncache.b2,1.0,-learning_rate,C,1);
}

void parallel_train(NeuralNetwork& nn, const arma::Mat<real>& X,
                    const arma::Mat<real>& y, real learning_rate, real reg,
                    const int epochs, const int batch_size, bool grad_check,
                    int print_every, int debug) {
  int rank, num_procs;
  MPI_SAFE_CALL(MPI_Comm_size(MPI_COMM_WORLD, &num_procs));
  MPI_SAFE_CALL(MPI_Comm_rank(MPI_COMM_WORLD, &rank));

  int N = (rank == 0) ? X.n_cols : 0;
  MPI_SAFE_CALL(MPI_Bcast(&N, 1, MPI_INT, 0, MPI_COMM_WORLD));

  std::ofstream error_file;
  error_file.open("Outputs/CpuGpuDiff.txt");
  int print_flag = 0;

  /* HINT: You can obtain a raw pointer to the memory used by Armadillo Matrices
     for storing elements in a column major way. Or you can allocate your own
     array memory space and store the elements in a row major way. Remember to
     update the Armadillo matrices in NeuralNetwork &nn of rank 0 before
     returning from the function. */

  // TODO
  /*
  M:number of featrues,
  H:number of neurons in hidden layer,
  C:number of classes, 10
  */
  int M = nn.H[0];
  int H = nn.H[1];
  int C = nn.H[2];

  int num_batches = (N + batch_size - 1) / batch_size;
  std::vector<real *> d_X_batches(num_batches);
  std::vector<real *> d_Y_batches(num_batches);

  //subdivide input batch of images and `MPI_scatter()' to each MPI node
  for (int batch = 0; batch < num_batches; ++batch) {
      int start_col = batch*batch_size;
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      int this_batch_size = last_col - start_col + 1;
      int nsample_per_proc = (this_batch_size + num_procs -1) / num_procs;

      int scounts_X[num_procs], scounts_Y[num_procs], displs_X[num_procs], displs_Y[num_procs];

      for(int i = 0; i < num_procs; i++){
        scounts_X[i] = M*std::min(nsample_per_proc,this_batch_size - i*nsample_per_proc);
        scounts_Y[i] = C*std::min(nsample_per_proc,this_batch_size - i*nsample_per_proc);
        displs_X[i] = i*M*nsample_per_proc;
        displs_Y[i] = i*C*nsample_per_proc;
      }

      arma::Mat<real> X_batch(M,scounts_X[rank]/M);
      MPI_SAFE_CALL(MPI_Scatterv(X.colptr(start_col),scounts_X,displs_X,MPI_FP,X_batch.memptr(),
      scounts_X[rank],MPI_FP,0,MPI_COMM_WORLD));
      arma::Mat<real> Y_batch(C ,scounts_Y[rank]/C);
      MPI_SAFE_CALL(MPI_Scatterv(y.colptr(start_col),scounts_Y,displs_Y,MPI_FP,Y_batch.memptr(),
      scounts_Y[rank],MPI_FP,0,MPI_COMM_WORLD));

      // data host to device , 3 processors
      cudaMalloc((void **)&d_X_batches[batch], scounts_X[rank] * sizeof(real));
      cudaMalloc((void **)&d_Y_batches[batch], scounts_Y[rank] * sizeof(real));
      cudaMemcpy(d_X_batches[batch], X_batch.memptr(), scounts_X[rank]* sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(d_Y_batches[batch], Y_batch.memptr(), scounts_Y[rank] * sizeof(real), cudaMemcpyHostToDevice);
  }
  /* iter is a variable used to manage debugging. It increments in the inner
     loop and therefore goes from 0 to epochs*num_batches */
  int iter = 0;
  //allocate deivice memory for nn parameters and host memory for derivatives
  NNcache nncache(M,H,C,batch_size);
  real *h_dW1 = (real *)malloc(H * M * sizeof(real));
  real *h_dW2 = (real *)malloc(C * H * sizeof(real));
  real *h_db1 = (real *)malloc(H * sizeof(real));
  real *h_db2 = (real *)malloc(C * sizeof(real));

  //copy data from host to devices
  cudaMemcpy(nncache.W1, nn.W[0].memptr(),H*M*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.W2, nn.W[1].memptr(),C*H*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.b1, nn.b[0].memptr(),H*sizeof(real),cudaMemcpyHostToDevice);
  cudaMemcpy(nncache.b2, nn.b[1].memptr(),C*sizeof(real),cudaMemcpyHostToDevice);

  for (int epoch = 0; epoch < epochs; ++epoch) {

    for (int batch = 0; batch < num_batches; ++batch) {
      /*
       * Possible implementation:
       * 1. subdivide input batch of images and `MPI_scatter()' to each MPI node
       * 2. compute each sub-batch of images' contribution to network
       * coefficient updates
       * 3. reduce the coefficient updates and broadcast to all nodes with
       * `MPI_Allreduce()'
       * 4. update local network coefficient at each node
       */

      // TODO
      int start_col = batch*batch_size;
      int last_col = std::min((batch + 1) * batch_size - 1, N - 1);
      int this_batch_size = last_col - start_col + 1;
      int nsample_per_proc = (this_batch_size + num_procs -1) / num_procs;
      //used for 3 GPU
      int nsample_this_proc = std::min(nsample_per_proc,this_batch_size-rank*nsample_per_proc);

      //training
      //forwards
      parallel_feedforward(nn,d_X_batches[batch] ,nncache,nsample_this_proc);

      //backprop
      parallel_backprop(nn,d_X_batches[batch], d_Y_batches[batch],reg, nncache, this_batch_size, nsample_this_proc, num_procs);

      // cudaMemcpy's, cudaMemcpyDeviceToHost
      cudaMemcpy(h_dW1, nncache.dW1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_dW2, nncache.dW2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_db1, nncache.db1, H * sizeof(real), cudaMemcpyDeviceToHost);
      cudaMemcpy(h_db2, nncache.db2, C * sizeof(real), cudaMemcpyDeviceToHost);

      // // MPI_Allreduce
      arma::Mat<real> dW1(size(nn.W[0]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_dW1, dW1.memptr(), H * M, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Mat<real> dW2(size(nn.W[1]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_dW2, dW2.memptr(), C * H, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Col<real> db1(size(nn.b[0]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_db1, db1.memptr(), H, MPI_FP, MPI_SUM, MPI_COMM_WORLD));
      arma::Col<real> db2(size(nn.b[1]), arma::fill::zeros);
      MPI_SAFE_CALL(MPI_Allreduce(h_db2, db2.memptr(), C, MPI_FP, MPI_SUM, MPI_COMM_WORLD));

      // // cudaMemcpy's, cudaMemcpyHostToDevice
      cudaMemcpy(nncache.dW1, dW1.memptr(), H * M * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.dW2, dW2.memptr(), C * H * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.db1, db1.memptr(), H * sizeof(real), cudaMemcpyHostToDevice);
      cudaMemcpy(nncache.db2, db2.memptr(), C * sizeof(real), cudaMemcpyHostToDevice);
      //add regularization term
    //  gpu_addmat(nncache.dW1,nncache.W1,nncache.dW1,1.0,reg,H,M);
     // gpu_addmat(nncache.dW2,nncache.W2,nncache.dW2,1.0,reg,C,H);
      // // Gradient descent step
      // nn.W[0] -= learning_rate * dW1;
      // nn.W[1] -= learning_rate * dW2;
      // nn.b[0] -= learning_rate * db1;
      // nn.b[1] -= learning_rate * db2;
      parallel_gradientdecent(nn,nncache,learning_rate);

      // +-*=+-*=+-*=+-*=+-*=+-*=+-*=+-*=+*-=+-*=+*-=+-*=+-*=+-*=+-*=+-*= //
      //                    POST-PROCESS OPTIONS                          //
      // +-*=+-*=+-*=+-*=+-*=+-*=+-*=+-*=+*-=+-*=+*-=+-*=+-*=+-*=+-*=+-*= //
      if (print_every <= 0) {
        print_flag = batch == 0;
      } else {
        print_flag = iter % print_every == 0;
      }

      if (debug && rank == 0 && print_flag) {
        // TODO
        // Copy data back to the CPU
        cudaMemcpy(nn.W[0].memptr(), nncache.W1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.W[1].memptr(), nncache.W2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.b[0].memptr(), nncache.b1, H * sizeof(real), cudaMemcpyDeviceToHost);
        cudaMemcpy(nn.b[1].memptr(), nncache.b2, C * sizeof(real), cudaMemcpyDeviceToHost);

        /* The following debug routine assumes that you have already updated the
         arma matrices in the NeuralNetwork nn.  */
        write_diff_gpu_cpu(nn, iter, error_file);
      }

      iter++;
    }
  }

  // TODO
  // Copy data back to the CPU
  cudaMemcpy(nn.W[0].memptr(), nncache.W1, H * M * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.W[1].memptr(), nncache.W2, C * H * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.b[0].memptr(), nncache.b1, H * sizeof(real), cudaMemcpyDeviceToHost);
  cudaMemcpy(nn.b[1].memptr(), nncache.b2, C * sizeof(real), cudaMemcpyDeviceToHost);
  error_file.close();

  // TODO
  // Free memory
  free(h_dW1);
  free(h_dW2);
  free(h_db1);
  free(h_db2);
  for(int batch = 0; batch < num_batches; ++batch) {
    cudaFree(d_X_batches[batch]);
    cudaFree(d_Y_batches[batch]);
  }
}

// cudaMalloc's
// MPI_Scatter. See this page for details on this function: https://www.open-mpi.org/doc/v4.1/
// cudaMemcpy's, cudaMemcpyHostToDevice
// loop over epochs and batches
// here you use your GPU kernels including myGEMM, sigmoid, softmax, etc
// cudaMemcpy's, cudaMemcpyDeviceToHost
// MPI_Allreduce
// cudaMemcpy's, cudaMemcpyHostToDevice
// Gradient descent step
// At the end, you copy back the nn coefficients from GPU to CPU and run some cudaFree's.


///////////////////////////////////////////////////////////////
//  Copyright 2011 John Maddock. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_

#include <boost/multiprecision/cpp_dec_float.hpp>
#include <boost/multiprecision/debug_adaptor.hpp>
#include <iostream>

void t1()
{
   //[debug_adaptor_eg
   //=#include <boost/multiprecision/debug_adaptor.hpp>
   //=#include <boost/multiprecision/cpp_dec_float.hpp>

   using namespace boost::multiprecision;

   typedef number<debug_adaptor<cpp_dec_float<50> > > fp_type;

   fp_type denom = 1;
   fp_type sum = 1;

   for(unsigned i = 2; i < 50; ++i)
   {
      denom *= i;
      sum += 1 / denom;
   }

   std::cout << std::setprecision(std::numeric_limits<fp_type>::digits) << sum << std::endl;
   //]
}

int main()
{
   t1();
   return 0;
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2012 Peter Caspers

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 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 license for more details.
*/

/*! \file adaptiverungekutta.hpp
    \brief Runge-Kutta ODE integration

    Runge Kutta method with adaptive stepsize as described in
    Numerical Recipes in C, Chapter 16.2
*/

#ifndef quantlib_adaptive_runge_kutta_hpp
#define quantlib_adaptive_runge_kutta_hpp

#include <ql/types.hpp>
#include <ql/errors.hpp>
#include <ql/utilities/disposable.hpp>
#include <boost/function.hpp>
#include <vector>
#include <cmath>

namespace QuantLib {

    template <class T = Real>
    class AdaptiveRungeKutta {
      public:
        typedef boost::function<
          Disposable<std::vector<T> >(const Real,
                                      const std::vector<T>&)> OdeFct;
        typedef boost::function<T(const Real, const T)> OdeFct1d;

        /*! The class is constructed with the following inputs:
            - eps       prescribed error for the solution
            - h1        start step size
            - hmin      smallest step size allowed
        */

        AdaptiveRungeKutta(const Real eps=1.0e-6,
                           const Real h1=1.0e-4,
                           const Real hmin=0.0)
        : eps_(eps), h1_(h1), hmin_(hmin),
          a2(0.2), a3(0.3), a4(0.6), a5(1.0), a6(0.875),
          b21(0.2), b31(3.0/40.0), b32(9.0/40.0), b41(0.3), b42(-0.9), b43(1.2),
          b51(-11.0/54.0), b52(2.5), b53(-70.0/27.0), b54(35.0/27.0),
          b61(1631.0/55296.0), b62(175.0/512.0), b63(575.0/13824.0),
          b64(44275.0/110592.0), b65(253.0/4096.0),
          c1(37.0/378.0), c3(250.0/621.0), c4(125.0/594.0), c6(512.0/1771.0),
          dc1(c1-2825.0/27648.0), dc3(c3-18575.0/48384.0),
          dc4(c4-13525.0/55296.0), dc5(-277.0/14336.0), dc6(c6-0.25),
          ADAPTIVERK_MAXSTP(10000), ADAPTIVERK_TINY(1.0E-30),
          ADAPTIVERK_SAFETY(0.9), ADAPTIVERK_PGROW(-0.2),
          ADAPTIVERK_PSHRINK(-0.25), ADAPTIVERK_ERRCON(1.89E-4) {}

        /*! Integrate the ode from \f$ x1 \f$ to \f$ x2 \f$ with
            initial value condition \f$ f(x1)=y1 \f$.

            The ode is given by a function \f$ F: R \times K^n
            \rightarrow K^n \f$ as \f$ f'(x) = F(x,f(x)) \f$, $K=R,
            C$ */
        Disposable<std::vector<T> > operator()(const OdeFct& ode,
                                               const std::vector<T>& y1,
                                               const Real x1,
                                               const Real x2);
        T operator()(const OdeFct1d& ode,
                     const T y1,
                     const Real x1,
                     const Real x2);

    private:
        void rkqs(std::vector<T>& y,
                  const std::vector<T>& dydx,
                  Real& x,
                  const Real htry,
                  const Real eps,
                  const std::vector<Real>& yScale,
                  Real &hdid,
                  Real &hnext,
                  const OdeFct& derivs);
        void rkck(const std::vector<T>& y,
                  const std::vector<T>& dydx,
                  Real x,
                  const Real h,
                  std::vector<T>& yout,
                  std::vector<T>& yerr,
                  const OdeFct& derivs);

        const std::vector<T> yStart_;
        const Real eps_, h1_, hmin_;
        const Real a2,a3,a4,a5,a6,
                   b21,b31,b32,b41,b42,b43,b51,b52,b53,b54,b61,b62,b63,b64,b65,
                   c1,c3,c4,c6,dc1,dc3,dc4,dc5,dc6;
        const double ADAPTIVERK_MAXSTP, ADAPTIVERK_TINY, ADAPTIVERK_SAFETY,
                   ADAPTIVERK_PGROW, ADAPTIVERK_PSHRINK, ADAPTIVERK_ERRCON;
    };

    template<class T>
    Disposable<std::vector<T> > AdaptiveRungeKutta<T>::operator()(
                                                     const OdeFct& ode,
                                                     const std::vector<T>& y1,
                                                     const Real x1,
                                                     const Real x2) {
        Size n = y1.size();
        std::vector<T> y(y1);
        std::vector<Real> yScale(n);
        Real x = x1;
        Real h = h1_* (x1<=x2 ? 1 : -1);
        Real hnext,hdid;

        for (Size nstp=1; nstp<=ADAPTIVERK_MAXSTP; nstp++) {
            std::vector<T> dydx=ode(x,y);
            for (Size i=0;i<n;i++)
                yScale[i] = std::abs(y[i])+std::abs(dydx[i]*h)+ADAPTIVERK_TINY;
            if ((x+h-x2)*(x+h-x1) > 0.0)
                h=x2-x;
            rkqs(y,dydx,x,h,eps_,yScale,hdid,hnext,ode);

            if ((x-x2)*(x2-x1) >= 0.0)
                return y;

            if (std::fabs(hnext) <= hmin_)
                QL_FAIL("Step size (" << hnext << ") too small ("
                        << hmin_ << " min) in AdaptiveRungeKutta");
            h=hnext;
        }
        QL_FAIL("Too many steps (" << ADAPTIVERK_MAXSTP
                << ") in AdaptiveRungeKutta");
    }

    namespace detail {

        template <class T>
        struct OdeFctWrapper {
            typedef typename AdaptiveRungeKutta<T>::OdeFct1d OdeFct1d;
            OdeFctWrapper(const OdeFct1d& ode1d)
            : ode1d_(ode1d) {}
            Disposable<std::vector<T> > operator()(const Real x,
                                                   const std::vector<T>& y) {
                std::vector<T> res(1,ode1d_(x,y[0]));
                return res;
            }
            const OdeFct1d& ode1d_;
        };

    }

    template<class T>
    T AdaptiveRungeKutta<T>::operator()(const OdeFct1d& ode,
                                        const T y1,
                                        const Real x1,
                                        const Real x2) {
        return operator()(detail::OdeFctWrapper<T>(ode),
                          std::vector<T>(1,y1),x1,x2)[0];
    }

    template<class T>
    void AdaptiveRungeKutta<T>::rkqs(std::vector<T>& y,
                                     const std::vector<T>& dydx,
                                     Real& x,
                                     const Real htry,
                                     const Real eps,
                                     const std::vector<Real>& yScale,
                                     Real& hdid,
                                     Real& hnext,
                                     const OdeFct& derivs) {
        Size n=y.size();
        Real errmax,htemp,xnew;
        std::vector<T> yerr(n),ytemp(n);

        Real h=htry;

        for(;;) {
            rkck(y,dydx,x,h,ytemp,yerr,derivs);
            errmax=0.0;
            for (Size i=0;i<n;i++)
                errmax=std::max(errmax,std::abs(yerr[i]/yScale[i]));
            errmax/=eps;
            if (errmax>1.0) {
                htemp=ADAPTIVERK_SAFETY*h*std::pow(errmax,ADAPTIVERK_PSHRINK);
                h = (h>=0.0 ? std::max(htemp,h/10) : std::min(htemp,h/10));
                xnew=x+h;
                if (xnew==x)
                    QL_FAIL("Stepsize (" << xnew
                            << ") underflow in AdaptiveRungeKutta::rkqs");
                continue;
            } else {
                if (errmax>ADAPTIVERK_ERRCON)
                    hnext=ADAPTIVERK_SAFETY*h*std::pow(errmax,ADAPTIVERK_PGROW);
                else
                    hnext=5.0*h;
                x+=(hdid=h);
                for (Size i=0;i<n;i++)
                    y[i]=ytemp[i];
                break;
            }
        }
    }

    template <class T>
    void AdaptiveRungeKutta<T>::rkck(const std::vector<T>& y,
                                     const std::vector<T>& dydx,
                                     Real x,
                                     const Real h,
                                     std::vector<T>& yout,
                                     std::vector<T> &yerr,
                                     const OdeFct& derivs) {

        Size n=y.size();
        std::vector<T> ak2(n),ak3(n),ak4(n),ak5(n),ak6(n),ytemp(n);

        // first step
        for (Size i=0;i<n;i++)
            ytemp[i]=y[i]+b21*h*dydx[i];

        // second step
        ak2=derivs(x+a2*h,ytemp);
        for (Size i=0;i<n;i++)
            ytemp[i]=y[i]+h*(b31*dydx[i]+b32*ak2[i]);

        // third step
        ak3=derivs(x+a3*h,ytemp);
        for (Size i=0;i<n;i++)
            ytemp[i]=y[i]+h*(b41*dydx[i]+b42*ak2[i]+b43*ak3[i]);

        // fourth step
        ak4=derivs(x+a4*h,ytemp);
        for (Size i=0;i<n;i++)
            ytemp[i]=y[i]+h*(b51*dydx[i]+b52*ak2[i]+b53*ak3[i]+b54*ak4[i]);

        // fifth step
        ak5=derivs(x+a5*h,ytemp);
        for (Size i=0;i<n;i++)
            ytemp[i]=y[i]+h*(b61*dydx[i]+b62*ak2[i]+b63*ak3[i]+b64*ak4[i]+b65*ak5[i]);

        // sixth step
        ak6=derivs(x+a6*h,ytemp);
        for (Size i=0;i<n;i++) {
            yout[i]=y[i]+h*(c1*dydx[i]+c3*ak3[i]+c4*ak4[i]+c6*ak6[i]);
            yerr[i]=h*(dc1*dydx[i]+dc3*ak3[i]+dc4*ak4[i]+dc5*ak5[i]+dc6*ak6[i]);
        }
    }

}

#endif


// Copyright (C) 2013  Davis E. King (davis@dlib.net)
// License: Boost Software License   See LICENSE.txt for the full license.

#include <dlib/python.h>
#include <dlib/statistics.h>

using namespace dlib;
namespace py = pybind11;

typedef std::vector<std::pair<unsigned long,double> > sparse_vect;

struct cca_outputs
{
    matrix<double,0,1> correlations;
    matrix<double> Ltrans;
    matrix<double> Rtrans;
};

cca_outputs _cca1 (
    const std::vector<sparse_vect>& L,
    const std::vector<sparse_vect>& R,
    unsigned long num_correlations,
    unsigned long extra_rank,
    unsigned long q,
    double regularization
)
{
    pyassert(num_correlations > 0 && L.size() > 0 && R.size() > 0 && L.size() == R.size() && regularization >= 0,
        "Invalid inputs");

    cca_outputs temp;
    temp.correlations = cca(L,R,temp.Ltrans,temp.Rtrans,num_correlations,extra_rank,q,regularization);
    return temp;
}

// ----------------------------------------------------------------------------------------

unsigned long sparse_vector_max_index_plus_one (
    const sparse_vect& v
)
{
    return max_index_plus_one(v);
}

matrix<double,0,1> apply_cca_transform (
    const matrix<double>& m,
    const sparse_vect& v
)
{
    pyassert((long)max_index_plus_one(v) <= m.nr(), "Invalid Inputs");
    return sparse_matrix_vector_multiply(trans(m), v);
}

void bind_cca(py::module& m)
{
    py::class_<cca_outputs>(m, "cca_outputs")
        .def_readwrite("correlations", &cca_outputs::correlations)
        .def_readwrite("Ltrans", &cca_outputs::Ltrans)
        .def_readwrite("Rtrans", &cca_outputs::Rtrans);

    m.def("max_index_plus_one", sparse_vector_max_index_plus_one, py::arg("v"),
"ensures    \n\
    - returns the dimensionality of the given sparse vector.  That is, returns a    \n\
      number one larger than the maximum index value in the vector.  If the vector    \n\
      is empty then returns 0.   "
    );

    m.def("apply_cca_transform", apply_cca_transform, py::arg("m"), py::arg("v"),
"requires    \n\
    - max_index_plus_one(v) <= m.nr()    \n\
ensures    \n\
    - returns trans(m)*v    \n\
      (i.e. multiply m by the vector v and return the result)   "
    );

    m.def("cca", _cca1, py::arg("L"), py::arg("R"), py::arg("num_correlations"), py::arg("extra_rank")=5, py::arg("q")=2, py::arg("regularization")=0,
"requires    \n\
    - num_correlations > 0    \n\
    - len(L) > 0     \n\
    - len(R) > 0     \n\
    - len(L) == len(R)    \n\
    - regularization >= 0    \n\
    - L and R must be properly sorted sparse vectors.  This means they must list their  \n\
      elements in ascending index order and not contain duplicate index values.  You can use \n\
      make_sparse_vector() to ensure this is true.  \n\
ensures    \n\
    - This function performs a canonical correlation analysis between the vectors    \n\
      in L and R.  That is, it finds two transformation matrices, Ltrans and    \n\
      Rtrans, such that row vectors in the transformed matrices L*Ltrans and    \n\
      R*Rtrans are as correlated as possible (note that in this notation we    \n\
      interpret L as a matrix with the input vectors in its rows).  Note also that    \n\
      this function tries to find transformations which produce num_correlations    \n\
      dimensional output vectors.    \n\
    - Note that you can easily apply the transformation to a vector using     \n\
      apply_cca_transform().  So for example, like this:     \n\
        - apply_cca_transform(Ltrans, some_sparse_vector)    \n\
    - returns a structure containing the Ltrans and Rtrans transformation matrices    \n\
      as well as the estimated correlations between elements of the transformed    \n\
      vectors.    \n\
    - This function assumes the data vectors in L and R have already been centered    \n\
      (i.e. we assume the vectors have zero means).  However, in many cases it is    \n\
      fine to use uncentered data with cca().  But if it is important for your    \n\
      problem then you should center your data before passing it to cca().   \n\
    - This function works with reduced rank approximations of the L and R matrices.    \n\
      This makes it fast when working with large matrices.  In particular, we use    \n\
      the dlib::svd_fast() routine to find reduced rank representations of the input    \n\
      matrices by calling it as follows: svd_fast(L, U,D,V, num_correlations+extra_rank, q)     \n\
      and similarly for R.  This means that you can use the extra_rank and q    \n\
      arguments to cca() to influence the accuracy of the reduced rank    \n\
      approximation.  However, the default values should work fine for most    \n\
      problems.    \n\
    - The dimensions of the output vectors produced by L*#Ltrans or R*#Rtrans are \n\
      ordered such that the dimensions with the highest correlations come first. \n\
      That is, after applying the transforms produced by cca() to a set of vectors \n\
      you will find that dimension 0 has the highest correlation, then dimension 1 \n\
      has the next highest, and so on.  This also means that the list of estimated \n\
      correlations returned from cca() will always be listed in decreasing order. \n\
    - This function performs the ridge regression version of Canonical Correlation    \n\
      Analysis when regularization is set to a value > 0.  In particular, larger    \n\
      values indicate the solution should be more heavily regularized.  This can be    \n\
      useful when the dimensionality of the data is larger than the number of    \n\
      samples.    \n\
    - A good discussion of CCA can be found in the paper \"Canonical Correlation    \n\
      Analysis\" by David Weenink.  In particular, this function is implemented    \n\
      using equations 29 and 30 from his paper.  We also use the idea of doing CCA    \n\
      on a reduced rank approximation of L and R as suggested by Paramveer S.    \n\
      Dhillon in his paper \"Two Step CCA: A new spectral method for estimating    \n\
      vector models of words\".   "

        );
}


// Distributed under the MIT License.
// See LICENSE.txt for details.

#include "NumericalAlgorithms/Spectral/SwshInterpolation.hpp"

#include <array>
#include <boost/math/special_functions/binomial.hpp>
#include <cmath>
#include <complex>
#include <cstddef>

#include "DataStructures/ComplexDataVector.hpp"
#include "DataStructures/ComplexModalVector.hpp"
#include "DataStructures/DataVector.hpp"
#include "DataStructures/SpinWeighted.hpp"
#include "NumericalAlgorithms/Spectral/SwshCoefficients.hpp"
#include "NumericalAlgorithms/Spectral/SwshCollocation.hpp"
#include "NumericalAlgorithms/Spectral/SwshTransform.hpp"
#include "Utilities/GenerateInstantiations.hpp"
#include "Utilities/Gsl.hpp"
#include "Utilities/StaticCache.hpp"
#include "Utilities/TMPL.hpp"

/// \cond

namespace Spectral {
namespace Swsh {

SpinWeightedSphericalHarmonic::SpinWeightedSphericalHarmonic(
    const int spin, const size_t l, const int m) noexcept
    : spin_{spin}, l_{l}, m_{m} {
  overall_prefactor_ = 1.0;
  const double double_l = l;
  const double double_m = m;
  const double double_spin = spin;
  if (std::abs(m) > std::abs(spin)) {
    for (size_t i = 0; i < static_cast<size_t>(std::abs(m) - std::abs(spin));
         ++i) {
      const double double_i = i;
      overall_prefactor_ *= (double_l + std::abs(double_m) - double_i) /
                            (double_l - (std::abs(double_spin) + double_i));
    }
  } else if (std::abs(spin) > std::abs(m)) {
    for (size_t i = 0; i < static_cast<size_t>(std::abs(spin) - std::abs(m));
         ++i) {
      const double double_i = i;
      overall_prefactor_ *= (double_l - (std::abs(double_m) + double_i)) /
                            (double_l + std::abs(double_spin) - double_i);
    }
  }
  // if neither is greater (they are equal), then the prefactor is 1.0
  overall_prefactor_ *= (2.0 * l + 1.0) / (4.0 * M_PI);
  overall_prefactor_ = sqrt(overall_prefactor_);
  overall_prefactor_ *= (m % 2) == 0 ? 1.0 : -1.0;

  // gcc warns about the casts in ways that are impossible to satisfy
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wsign-conversion"
  if (static_cast<int>(l) < std::abs(spin)) {
    if (spin < 0) {
      r_prefactors_ =
          std::vector<double>(l + static_cast<size_t>(std::abs(spin)) + 1, 0.0);
    }
  } else {
    // the casts in the reserve are in correct order, but clang-format
    // erroneously requests a change
    // NOLINTNEXTLINE(misc-misplaced-widening-cast)
    r_prefactors_.reserve(static_cast<size_t>(static_cast<int>(l) - spin + 1));
    for (int r = 0; r <= (static_cast<int>(l) - spin); ++r) {
      if (r + spin - m >= 0 and static_cast<int>(l) - r + m >= 0) {
        r_prefactors_.push_back(
            boost::math::binomial_coefficient<double>(
                static_cast<size_t>(static_cast<int>(l) - spin),
                static_cast<size_t>(r)) *
            boost::math::binomial_coefficient<double>(
                static_cast<size_t>(static_cast<int>(l) + spin),
                static_cast<size_t>(spin - m + r)) *
            (((static_cast<int>(l) - r - spin) % 2) == 0 ? 1.0 : -1.0));
      } else {
        r_prefactors_.push_back(0.0);
      }
    }
  }
#pragma GCC diagnostic pop
}

void SpinWeightedSphericalHarmonic::evaluate(
    const gsl::not_null<ComplexDataVector*> result, const DataVector& theta,
    const DataVector& phi, const DataVector& sin_theta_over_2,
    const DataVector& cos_theta_over_2) const noexcept {
  result->destructive_resize(theta.size());
  *result = 0.0;
  DataVector theta_factor{theta.size()};
  for (int r = 0; r <= (static_cast<int>(l_) - spin_); ++r) {
    if (2 * static_cast<int>(l_) > 2 * r + spin_ - m_) {
      theta_factor = pow(cos_theta_over_2, 2 * r + spin_ - m_) *
                     pow(sin_theta_over_2,
                         2 * static_cast<int>(l_) - (2 * r + spin_ - m_));
    } else if (2 * static_cast<int>(l_) < 2 * r + spin_ - m_) {
      theta_factor = pow(cos_theta_over_2 / sin_theta_over_2,
                         2 * r + spin_ - m_ - 2 * static_cast<int>(l_)) *
                     pow(cos_theta_over_2, 2 * l_);
    } else {
      theta_factor = pow(cos_theta_over_2, 2 * l_);
    }
    *result += gsl::at(r_prefactors_, r) * theta_factor;
  }
  // optimization note: this has not been compared with a complex `exp`
  // function, and it is not obvious which should be faster in practice.
  *result *=
      overall_prefactor_ *
      (std::complex<double>(1.0, 0.0) * cos(static_cast<double>(m_) * phi) +
       std::complex<double>(0.0, 1.0) * sin(static_cast<double>(m_) * phi));
}

ComplexDataVector SpinWeightedSphericalHarmonic::evaluate(
    const DataVector& theta, const DataVector& phi,
    const DataVector& sin_theta_over_2,
    const DataVector& cos_theta_over_2) const noexcept {
  ComplexDataVector result{theta.size(), 0.0};
  evaluate(make_not_null(&result), theta, phi, sin_theta_over_2,
           cos_theta_over_2);
  return result;
}

std::complex<double> SpinWeightedSphericalHarmonic::evaluate(
    const double theta, const double phi) const noexcept {
  std::complex<double> accumulator = 0.0;
  const double cos_theta_over_two = cos(0.5 * theta);
  const double sin_theta_over_two = sin(0.5 * theta);
  double theta_factor;
  for (int r = 0; r <= (static_cast<int>(l_) - spin_); ++r) {
    if (2 * static_cast<int>(l_) > 2 * r + spin_ - m_) {
      theta_factor = pow(cos_theta_over_two, 2 * r + spin_ - m_) *
                     pow(sin_theta_over_two,
                         2 * static_cast<int>(l_) - (2 * r + spin_ - m_));
    } else if (2 * static_cast<int>(l_) < 2 * r + spin_ - m_) {
      theta_factor = pow(cos_theta_over_two / sin_theta_over_two,
                         2 * r + spin_ - m_ - 2 * static_cast<int>(l_)) *
                     pow(cos_theta_over_two, 2 * l_);
    } else {
      theta_factor = pow(cos_theta_over_two, 2 * l_);
    }
    accumulator += gsl::at(r_prefactors_, r) *
                   std::complex<double>(cos(static_cast<double>(m_) * phi),
                                        sin(static_cast<double>(m_) * phi)) *
                   theta_factor;
  }
  accumulator *= overall_prefactor_;
  return accumulator;
}

void SpinWeightedSphericalHarmonic::pup(PUP::er& p) noexcept {
  p | spin_;
  p | l_;
  p | m_;
  p | overall_prefactor_;
  p | r_prefactors_;
}

// A function for indexing a desired element in one of the caches stored
// in `ClenshawRecurrenceConstants`.
// Useful for accessing the `beta_constant`, `alpha_constant`, or
// `alpha_prefactor` recurrence constants.
size_t clenshaw_cache_index(const size_t l_max, const int spin, const int l,
                            const int m) noexcept {
  return goldberg_mode_index(l_max - 2, static_cast<size_t>(l - 2), m) -
         static_cast<size_t>(square(spin));
}

// see the detailed doxygen for `SwshInterpolator` for full mathematical details
// of the recurrence constant computations
template <int Spin>
struct ClenshawRecurrenceConstants {
  ClenshawRecurrenceConstants() = default;

  explicit ClenshawRecurrenceConstants(size_t l_max) noexcept
      : alpha_prefactor{square(l_max + 1) -
                        square(static_cast<size_t>(std::abs(Spin)))},
        alpha_constant{square(l_max + 1) -
                       square(static_cast<size_t>(std::abs(Spin)))},
        beta_constant{square(l_max + 1) -
                      square(static_cast<size_t>(std::abs(Spin)))},
        harmonic_at_l_min_prefactors{2 * l_max + 1},
        harmonic_at_l_min_plus_one_recurrence_prefactors{2 * l_max + 1},
        harmonic_m_recurrence_prefactors{2 * l_max + 1} {
    ASSERT(static_cast<int>(l_max) > Spin,
           "l_max must be greater than the spin-weight when computing "
           "ClenshawRecurrenceConstants");
    double l_plus_k;
    double l_min_plus_k;
    double a;
    double b;
    int l_min;
    double prefactor_accumulator;
    lambda.reserve(2 * l_max + 1);
    for (int m = -static_cast<int>(l_max); m <= static_cast<int>(l_max); ++m) {
      a = static_cast<double>(std::abs(Spin + m));
      b = static_cast<double>(std::abs(Spin - m));
      l_min = std::max(std::abs(m), std::abs(Spin));

      // gcc warns about an optimization that doesn't work if we overflow. None
      // of this will overflow provided l_max is not unreasonably high (less
      // than ~10^5 will not overflow).
      for (int l = l_min + 2; l <= static_cast<int>(l_max); ++l) {
        // start caching at 2 greater than the l_min for a given m. Those are
        // the last terms needed by (descending) Clenshaw sum.
        l_plus_k = static_cast<double>(l) - 0.5 * (a + b);
        alpha_prefactor[clenshaw_cache_index(l_max, Spin, l, m)] =
            0.5 * sqrt((2.0 * l + 1.0) * (2.0 * l - 1.0) /
                       (l_plus_k * (l_plus_k + a + b) * (l_plus_k + a) *
                        (l_plus_k + b)));
        alpha_constant[clenshaw_cache_index(l_max, Spin, l, m)] =
            alpha_prefactor[clenshaw_cache_index(l_max, Spin, l, m)] *
            ((square(a) - square(b)) / (2.0 * l - 2.0));
        alpha_prefactor[clenshaw_cache_index(l_max, Spin, l, m)] *= (2.0 * l);
        beta_constant[clenshaw_cache_index(l_max, Spin, l, m)] =
            -sqrt((2.0 * l + 1.0) * (l_plus_k + a - 1.0) *
                  (l_plus_k + b - 1.0) * (l_plus_k - 1.0) *
                  (l_plus_k + a + b - 1.0) /
                  ((2.0 * l - 3.0) * l_plus_k * (l_plus_k + a + b) *
                   (l_plus_k + a) * (l_plus_k + b))) *
            (2.0 * l) / (2.0 * l - 2.0);
      }
      lambda.push_back(Spin >= -m ? 0 : Spin + m);

      // pre-compute the prefactors for the lowest order harmonics for each m
      prefactor_accumulator = 1.0;
      l_min_plus_k = -0.5 * (std::abs(Spin + m) + b) + l_min;
      for (int i = 1; i <= b; ++i) {
        if (l_min_plus_k + a + i > 0.0) {
          prefactor_accumulator *= static_cast<double>(l_min_plus_k + a + i);
        }
        if (l_min_plus_k + i > 0.0) {
          prefactor_accumulator /= static_cast<double>(l_min_plus_k + i);
        }
      }
      prefactor_accumulator =
          sqrt(prefactor_accumulator * (2.0 * l_min + 1.0) / (4.0 * M_PI));
      prefactor_accumulator *=
          ((m + gsl::at(lambda, m + static_cast<int>(l_max))) % 2) == 0 ? 1.0
                                                                        : -1.0;
      // this is the right order of the casts, other orders give the wrong
      // answer NOLINTNEXTLINE(misc-misplaced-widening-cast)
      harmonic_at_l_min_prefactors[static_cast<size_t>(
          m + static_cast<int>(l_max))] = prefactor_accumulator;

      // pre-compute the prefactors for bootstrapping the second-to-lowest order
      // harmonics for each m

      // this is the right order of the casts, other orders give the wrong
      // answer NOLINTNEXTLINE(misc-misplaced-widening-cast)
      harmonic_at_l_min_plus_one_recurrence_prefactors[static_cast<size_t>(
          m + static_cast<int>(l_max))] =
          sqrt((2.0 * (l_min) + 3.0) * (l_min_plus_k + 1.0) *
               (l_min_plus_k + a + b + 1.0) /
               ((2.0 * (l_min) + 1.0) * (l_min_plus_k + a + 1.0) *
                (l_min_plus_k + b + 1.0)));
    }
    // separate loop because we'll need the lambdas entirely populated for this
    // set of prefactors
    int lambda_difference;
    for (int m = -static_cast<int>(l_max); m <= static_cast<int>(l_max); ++m) {
      if (std::abs(m) > std::abs(Spin)) {
        l_min = std::max(std::abs(m), std::abs(Spin));
        a = std::abs(Spin + m);
        b = std::abs(Spin - m);
        l_min_plus_k = -0.5 * (std::abs(Spin + m) + b) + l_min;

        prefactor_accumulator =
            sqrt((2.0 * std::abs(m) + 1.0) * (l_min_plus_k + a + b - 1.0) *
                 (l_min_plus_k + a + b) /
                 ((2.0 * std::abs(m) - 1.0) * (l_min_plus_k + a) *
                  (l_min_plus_k + b)));
        // there is an extra `1` in these expressions to account for the -1 out
        // front of the recurrence relations.
        lambda_difference = gsl::at(lambda, m + static_cast<int>(l_max)) + 1;
        if (m > 0) {
          lambda_difference -= gsl::at(lambda, m - 1 + static_cast<int>(l_max));
        } else {
          lambda_difference -= gsl::at(lambda, m + 1 + static_cast<int>(l_max));
        }
        prefactor_accumulator *= lambda_difference % 2 == 0 ? 1.0 : -1.0;
        // this is the right order of the casts, other orders give the wrong
        // answer NOLINTNEXTLINE(misc-misplaced-widening-cast)
        harmonic_m_recurrence_prefactors[static_cast<size_t>(
            m + static_cast<int>(l_max))] = prefactor_accumulator;
      }
    }
  }

  /// Serialization for Charm++.
  void pup(PUP::er& p) noexcept {  // NOLINT
    p | alpha_prefactor;
    p | alpha_constant;
    p | beta_constant;
    p | lambda;
    p | harmonic_at_l_min_prefactors;
    p | harmonic_at_l_min_plus_one_recurrence_prefactors;
    p | harmonic_m_recurrence_prefactors;
  }

  // Tables are stored in a triangular Goldberg style
  DataVector alpha_prefactor;
  DataVector alpha_constant;
  DataVector beta_constant;
  std::vector<int> lambda;
  DataVector harmonic_at_l_min_prefactors;
  DataVector harmonic_at_l_min_plus_one_recurrence_prefactors;
  DataVector harmonic_m_recurrence_prefactors;
};

// A lazy static cache interface for retrieving `ClenshawRecurrenceConstants`.
template <int Spin>
const ClenshawRecurrenceConstants<Spin>& cached_clenshaw_factors(
    const size_t l_max) noexcept {
  const static auto lazy_clenshaw_cache =
      make_static_cache<CacheRange<0, collocation_maximum_l_max>>(
          [](const size_t local_l_max) noexcept {
            return ClenshawRecurrenceConstants<Spin>{local_l_max};
          });
  return lazy_clenshaw_cache(l_max);
}

SwshInterpolator::SwshInterpolator(const DataVector& theta,
                                   const DataVector& phi,
                                   const size_t l_max) noexcept
    : l_max_{l_max},
      raw_libsharp_coefficient_buffer_{
          size_of_libsharp_coefficient_vector(l_max)},
      raw_goldberg_coefficient_buffer_{square(l_max + 1)} {
  cos_m_phi_ = std::vector<DataVector>(l_max + 1);
  sin_m_phi_ = std::vector<DataVector>(l_max + 1);
  cos_theta_ = cos(theta);
  sin_theta_ = sin(theta);
  cos_theta_over_two_ = cos(0.5 * theta);
  sin_theta_over_two_ = sin(0.5 * theta);
  // evaluate cos(m phi) and sin(m phi) via recurrence
  cos_m_phi_[0] = DataVector{phi.size(), 1.0};
  sin_m_phi_[0] = DataVector{phi.size(), 0.0};
  const DataVector m_phi_beta = sin(phi);
  const DataVector m_phi_alpha = 2.0 * square(sin(0.5 * phi));
  for (size_t m = 1; m <= l_max; ++m) {
    cos_m_phi_[m] = cos_m_phi_[m - 1] - (m_phi_alpha * cos_m_phi_[m - 1] +
                                         m_phi_beta * sin_m_phi_[m - 1]);
    sin_m_phi_[m] = sin_m_phi_[m - 1] - (m_phi_alpha * sin_m_phi_[m - 1] -
                                         m_phi_beta * cos_m_phi_[m - 1]);
  }
}

template <int Spin>
void SwshInterpolator::interpolate(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> interpolated,
    const SpinWeighted<ComplexModalVector, Spin>& goldberg_modes) const
    noexcept {
  interpolated->destructive_resize(cos_theta_.size());
  interpolated->data() = 0.0;

  // used only if s=0;
  SpinWeighted<ComplexDataVector, Spin> cached_base_harmonic;

  // used during both recurrence legs
  SpinWeighted<ComplexDataVector, Spin> current_cached_harmonic;
  SpinWeighted<ComplexDataVector, Spin> current_cached_harmonic_l_plus_one;

  // perform the Clenshaw sums over positive m >= 0.
  for (int m = 0; m <= static_cast<int>(l_max_); ++m) {
    if (std::abs(Spin) >= std::abs(m)) {
      direct_evaluation_swsh_at_l_min(make_not_null(&current_cached_harmonic),
                                      m);
      evaluate_swsh_at_l_min_plus_one(
          make_not_null(&current_cached_harmonic_l_plus_one),
          current_cached_harmonic, m);
    } else {
      evaluate_swsh_m_recurrence_at_l_min(
          make_not_null(&current_cached_harmonic), m);
      evaluate_swsh_at_l_min_plus_one(
          make_not_null(&current_cached_harmonic_l_plus_one),
          current_cached_harmonic, m);
    }
    if (Spin == 0 and m == 0) {
      cached_base_harmonic = current_cached_harmonic;
    }
    clenshaw_sum(interpolated, current_cached_harmonic,
                 current_cached_harmonic_l_plus_one, goldberg_modes, m);
  }
  // perform the Clenshaw sums over m < 0.
  for (int m = -1; m >= -static_cast<int>(l_max_); --m) {
    // initialize the recurrence for negative m
    if (m == -1 and Spin == 0) {
      current_cached_harmonic = cached_base_harmonic;
    }
    if (std::abs(Spin) >= std::abs(m)) {
      direct_evaluation_swsh_at_l_min(make_not_null(&current_cached_harmonic),
                                      m);
      evaluate_swsh_at_l_min_plus_one(
          make_not_null(&current_cached_harmonic_l_plus_one),
          current_cached_harmonic, m);
    } else {
      evaluate_swsh_m_recurrence_at_l_min(
          make_not_null(&current_cached_harmonic), m);
      evaluate_swsh_at_l_min_plus_one(
          make_not_null(&current_cached_harmonic_l_plus_one),
          current_cached_harmonic, m);
    }
    clenshaw_sum(interpolated, current_cached_harmonic,
                 current_cached_harmonic_l_plus_one, goldberg_modes, m);
  }
}

template <int Spin>
void SwshInterpolator::interpolate(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> interpolated,
    const SpinWeighted<ComplexDataVector, Spin>& libsharp_collocation) const
    noexcept {
  SpinWeighted<ComplexModalVector, Spin> libsharp_modes;
  // this function is 'const', but modifies the internal buffer. The reason to
  // allow it to be 'const' anyways is that no interface makes any assumption
  // about the starting state of the internal buffer; it is kept exclusively to
  // save allocations.
  libsharp_modes.set_data_ref(raw_libsharp_coefficient_buffer_.data(),
                              raw_libsharp_coefficient_buffer_.size());
  swsh_transform(l_max_, 1, make_not_null(&libsharp_modes),
                 libsharp_collocation);
  SpinWeighted<ComplexModalVector, Spin> goldberg_modes;
  libsharp_to_goldberg_modes(make_not_null(&goldberg_modes), libsharp_modes,
                             l_max_);
  interpolate(interpolated, goldberg_modes);
}

template <int Spin>
void SwshInterpolator::direct_evaluation_swsh_at_l_min(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> harmonic,
    const int m) const noexcept {
  const auto& clenshaw_factors = cached_clenshaw_factors<Spin>(l_max_);
  // for this evaluation, we don't worry about recurrence because it will only
  // be called for m between -s and +s, and s should always be small. In
  // principle, it is probably true that a more complicated recurrence exists
  // for this case, but would require a bit of derivation work
  harmonic->data() =
      // clang-tidy: this is the right order of the casts, other orders give the
      // wrong answer
      // NOLINTNEXTLINE(misc-misplaced-widening-cast)
      clenshaw_factors.harmonic_at_l_min_prefactors[static_cast<size_t>(
          m + static_cast<int>(l_max_))] *
      (std::complex<double>(1.0, 0.0) * gsl::at(cos_m_phi_, std::abs(m)) +
       std::complex<double>(0.0, 1.0) * (m >= 0 ? 1.0 : -1.0) *
           gsl::at(sin_m_phi_, std::abs(m))) *
      pow(sin_theta_over_two_, static_cast<size_t>(std::abs(Spin + m))) *
      pow(cos_theta_over_two_, static_cast<size_t>(std::abs(Spin - m)));
}

template <int Spin>
void SwshInterpolator::evaluate_swsh_at_l_min_plus_one(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> harmonic,
    const SpinWeighted<ComplexDataVector, Spin>& harmonic_at_l_min,
    const int m) const noexcept {
  const auto& clenshaw_factors = cached_clenshaw_factors<Spin>(l_max_);
  const double a = std::abs(Spin + m);
  const double b = std::abs(Spin - m);
  harmonic->data() =
      clenshaw_factors
          // clang-tidy: this is the right order of the casts, other orders give
          // the wrong answer
          // NOLINTNEXTLINE(misc-misplaced-widening-cast)
          .harmonic_at_l_min_plus_one_recurrence_prefactors[static_cast<size_t>(
              m + static_cast<int>(l_max_))] *
      harmonic_at_l_min.data() *
      (a + 1.0 + 0.5 * (a + b + 2.0) * (cos_theta_ - 1.0));
}

template <int Spin>
void SwshInterpolator::evaluate_swsh_m_recurrence_at_l_min(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> harmonic,
    const int m) const noexcept {
  const auto& clenshaw_factors = cached_clenshaw_factors<Spin>(l_max_);
  harmonic->data() =
      // this is the right order of the casts, other orders give the wrong
      // answer
      // NOLINTNEXTLINE(misc-misplaced-widening-cast)
      clenshaw_factors.harmonic_m_recurrence_prefactors[static_cast<size_t>(
          m + static_cast<int>(l_max_))] *
      (sin_theta_ / 2.0) * harmonic->data();
  if (m > 0) {
    harmonic->data() *= (std::complex<double>(1.0, 0.0) * cos_m_phi_[1] +
                         std::complex<double>(0.0, 1.0) * sin_m_phi_[1]);
  } else {
    harmonic->data() *= (std::complex<double>(1.0, 0.0) * cos_m_phi_[1] -
                         std::complex<double>(0.0, 1.0) * sin_m_phi_[1]);
  }
}

template <int Spin>
void SwshInterpolator::clenshaw_sum(
    const gsl::not_null<SpinWeighted<ComplexDataVector, Spin>*> interpolation,
    const SpinWeighted<ComplexDataVector, Spin>& l_min_harmonic,
    const SpinWeighted<ComplexDataVector, Spin>& l_min_plus_one_harmonic,
    const SpinWeighted<ComplexModalVector, Spin>& goldberg_modes,
    const int m) const noexcept {
  // Since we need various combinations of the three-term recurrence constants
  // up to two orders higher, we write recurrence results to a cyclic
  // three-element cache
  std::array<ComplexDataVector, 3> recurrence_cache;
  recurrence_cache[2] = ComplexDataVector{interpolation->size(), 0.0};
  recurrence_cache[1] = ComplexDataVector{interpolation->size(), 0.0};
  recurrence_cache[0] = ComplexDataVector{interpolation->size(), 0.0};
  const auto& clenshaw_factors = cached_clenshaw_factors<Spin>(l_max_);

  for (auto l = static_cast<int>(l_max_);
       l > std::max(std::abs(Spin), std::abs(m)); l--) {
    // We want to define some cache_offset so that we can index the three
    // elements of recurrence_cache with indices cache_offset%3,
    // (cache_offset+1)%3, and (cache_offset+2)%3, and so that cache_offset
    // decreases by one on each iteration. The "obvious" way to do this is to
    // choose cache_offset = l - l_max, so that cache_offset starts at zero and
    // then decreases each iteration. However, this gives negative values of
    // cache_offset, and C++ modular arithmetic doesn't behave the way we'd want
    // for that process at negative values. But note that adding any multiple of
    // 3 to the "obvious" value of cache_offset will give identical indexing,
    // and choosing this multiple of 3 large enough (i.e. larger than l_max)
    // guarantees that the new cache_offset is positive for all l. So we choose
    // to add 3*l_max to the "obvious" value of cache_offset; in other words we
    // define cache_offset = l + 2 l_max.
    // In future, if this trick needs to be re-implemented in another use-case,
    // it should instead be factored out into a separate rotating cache utility.
    const int cache_offset = (l + 2 * static_cast<int>(l_max_));
    // gcc warns about the casts in ways that are impossible to satisfy
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wsign-conversion"
    gsl::at(recurrence_cache, (cache_offset) % 3) =
        goldberg_modes.data()[square(static_cast<size_t>(l)) +
                              static_cast<size_t>(l + m)];
    if (l < static_cast<int>(l_max_)) {
      gsl::at(recurrence_cache, (cache_offset) % 3) +=
          (clenshaw_factors
               .alpha_constant[clenshaw_cache_index(l_max_, Spin, l + 1, m)] +
           cos_theta_ * clenshaw_factors.alpha_prefactor[clenshaw_cache_index(
                            l_max_, Spin, l + 1, m)]) *
          gsl::at(recurrence_cache, (cache_offset + 1) % 3);
    }
    if (l < static_cast<int>(l_max_) - 1) {
      gsl::at(recurrence_cache, (cache_offset) % 3) +=
          clenshaw_factors
              .beta_constant[clenshaw_cache_index(l_max_, Spin, l + 2, m)] *
          gsl::at(recurrence_cache, (cache_offset + 2) % 3);
    }
  }
  const int l_min = std::max(std::abs(Spin), std::abs(m));
  const int cache_offset = (l_min + 2 * static_cast<int>(l_max_));

  if (l_max_ >=
      static_cast<size_t>(std::max(std::abs(Spin), std::abs(m))) + 2) {
    *interpolation +=
        l_min_harmonic *
            goldberg_modes.data()[square(static_cast<size_t>(l_min)) +
                                  static_cast<size_t>(l_min + m)] +
        l_min_plus_one_harmonic *
            gsl::at(recurrence_cache, (cache_offset + 1) % 3) +
        l_min_harmonic * gsl::at(recurrence_cache, (cache_offset + 2) % 3) *
            clenshaw_factors.beta_constant[clenshaw_cache_index(
                l_max_, Spin, std::max(std::abs(Spin), std::abs(m)) + 2, m)];
  } else {
    *interpolation +=
        l_min_harmonic *
            goldberg_modes.data()[square(static_cast<size_t>(l_min)) +
                                  static_cast<size_t>(l_min + m)] +
        l_min_plus_one_harmonic *
            gsl::at(recurrence_cache, (cache_offset + 1) % 3);
  }
#pragma GCC diagnostic pop
}

void SwshInterpolator::pup(PUP::er& p) noexcept {
  p | l_max_;
  p | cos_theta_;
  p | sin_theta_;
  p | cos_theta_over_two_;
  p | sin_theta_over_two_;
  p | sin_m_phi_;
  p | cos_m_phi_;
  p | raw_libsharp_coefficient_buffer_;
  p | raw_goldberg_coefficient_buffer_;
}

#define GET_SPIN(data) BOOST_PP_TUPLE_ELEM(0, data)

#define INTERPOLATION_INSTANTIATION(r, data)                                  \
  template struct ClenshawRecurrenceConstants<GET_SPIN(data)>;                \
  template const ClenshawRecurrenceConstants<GET_SPIN(data)>&                 \
  cached_clenshaw_factors<GET_SPIN(data)>(const size_t l_max) noexcept;       \
  template void SwshInterpolator::interpolate<GET_SPIN(data)>(                \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          interpolated,                                                       \
      const SpinWeighted<ComplexModalVector, GET_SPIN(data)>& goldberg_modes) \
      const noexcept;                                                         \
  template void SwshInterpolator::interpolate<GET_SPIN(data)>(                \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          interpolated,                                                       \
      const SpinWeighted<ComplexDataVector, GET_SPIN(data)>&                  \
          libsharp_collocation) const noexcept;                               \
  template void                                                               \
  SwshInterpolator::direct_evaluation_swsh_at_l_min<GET_SPIN(data)>(          \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          harmonic,                                                           \
      const int m) const noexcept;                                            \
  template void                                                               \
  SwshInterpolator::evaluate_swsh_at_l_min_plus_one<GET_SPIN(data)>(          \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          harmonic,                                                           \
      const SpinWeighted<ComplexDataVector, GET_SPIN(data)>&                  \
          harmonic_at_l_min,                                                  \
      const int m) const noexcept;                                            \
  template void                                                               \
  SwshInterpolator::evaluate_swsh_m_recurrence_at_l_min<GET_SPIN(data)>(      \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          harmonic,                                                           \
      const int m) const noexcept;                                            \
  template void SwshInterpolator::clenshaw_sum<GET_SPIN(data)>(               \
      const gsl::not_null<SpinWeighted<ComplexDataVector, GET_SPIN(data)>*>   \
          interpolation,                                                      \
      const SpinWeighted<ComplexDataVector, GET_SPIN(data)>& l_min_harmonic,  \
      const SpinWeighted<ComplexDataVector, GET_SPIN(data)>&                  \
          l_min_plus_one_harmonic,                                            \
      const SpinWeighted<ComplexModalVector, GET_SPIN(data)>& goldberg_modes, \
      const int m) const noexcept;

GENERATE_INSTANTIATIONS(INTERPOLATION_INSTANTIATION, (-2, -1, 0, 1, 2))

#undef INTERPOLATION_INSTANTIATION
#undef GET_SPIN

}  // namespace Swsh
}  // namespace Spectral
/// \endcond


#include "multi_cg/multi_cg.hpp"

#include <Eigen/Core>

#include <iostream>

using namespace Eigen;

struct BlockVector;

struct BlockVector {
    MatrixXcd vec;

    typedef std::complex<double> value_type;

    void block_axpy(std::vector<std::complex<double>> alphas, BlockVector const &X, size_t num) {
        DiagonalMatrix<std::complex<double>,Dynamic,Dynamic> D = Map<VectorXcd>(alphas.data(), num).asDiagonal();
        vec.leftCols(num) += X.vec.leftCols(num) * D;
    }

    void block_axpy_scatter(std::vector<std::complex<double>> alphas, BlockVector const &X, std::vector<size_t> ids) {
        for (size_t i = 0; i < ids.size(); ++i) {
            vec.col(ids[i]) += alphas[i] * X.vec.col(i);
        }
    }

    // rhos[i] = dot(X[i], Y[i])
    void block_dot(BlockVector const &Y, std::vector<std::complex<double>> &rhos, size_t num) {
        VectorXcd result = (vec.leftCols(num).adjoint() * Y.vec.leftCols(num)).diagonal();
        VectorXcd::Map(rhos.data(), result.size()) = result;
    }

    // X[:, i] = Z[:, i] + alpha[i] * X[:, i] for i < num_unconverged
    void block_xpby(BlockVector const &Z, std::vector<std::complex<double>> alphas, size_t num) {
        DiagonalMatrix<std::complex<double>,Dynamic,Dynamic> D = Map<VectorXcd>(alphas.data(), num).asDiagonal();
        vec.leftCols(num) = Z.vec.leftCols(num) + vec.leftCols(num) * D;
    }

    void copy(BlockVector const &X, size_t num) {
        vec.leftCols(num) = X.vec.leftCols(num);
    }

    void fill(std::complex<double> val) {
        vec.fill(val);
    }

    auto cols() {
        return vec.cols();
    }

    void repack(std::vector<size_t> const &ids) {
        for (size_t i = 0; i < ids.size(); ++i) {
            auto j = ids[i];
            if (j != i) {
                vec.col(i) = vec.col(j);
            }
        }
    }
};

// This is a linear but special operator A(X)
// producing AX + XD where D_ii = shifts[i] is a diagonal matrix.
// So column-wise it performs (A + shift[i])X[:, i]
// the multiply function basically does a gemv on every column with a different shift
// so alpha * A(X) + beta * Y.
struct PosDefMatrixShifted {
    DiagonalMatrix<double, Dynamic, Dynamic> A;
    VectorXd shifts;

    void multiply(double alpha, BlockVector const &u, double beta, BlockVector &v, size_t num) {
        v.vec.leftCols(num) = alpha * A * u.vec.leftCols(num) + alpha * u.vec.leftCols(num) * shifts.head(num).asDiagonal() + beta * v.vec.leftCols(num);
    }

    void repack(std::vector<size_t> const &ids) {
        for (size_t i = 0; i < ids.size(); ++i) {
            auto j = ids[i];
            if (j != i) {
                shifts[i] = shifts[j];
            }
        }
    }
};

struct IdentityPreconditioner {
    void apply(BlockVector &C, BlockVector const &B) {
        C = B;
    }
    void repack(std::vector<size_t> const &ids) {
        // nothing to do;
    }
};

int main() {
    size_t m = 40;
    size_t n = 10;

    auto A_shifts = VectorXd::LinSpaced(n, 1, n);
    auto A_diag = VectorXd::LinSpaced(m, 1, m);

    auto A = PosDefMatrixShifted{
        A_diag.asDiagonal(),
        A_shifts
    };

    auto P = IdentityPreconditioner{};

    auto U = BlockVector{MatrixXcd::Zero(m, n)};
    auto C = BlockVector{MatrixXcd::Zero(m, n)};
    auto X = BlockVector{MatrixXcd::Random(m, n)};
    auto B = BlockVector{MatrixXcd::Random(m, n)};
    auto R = B;

    auto tol = 1e-10;

    auto resnorms = sirius::cg::multi_cg(
        A, P,
        X, R, U, C,
        100, tol, false
    );

    // check the residual norms according to the algorithm
    for (size_t i = 0; i < resnorms.size(); ++i) {
        std::cout << "shift " << i << " needed " << resnorms[i].size() << " iterations " << std::abs(resnorms[i].back()) << "\n";

        if (std::abs(resnorms[i].back()) > tol) {
            return 1;
        }
    }

    // True residual norms might be different! because of rounding errors.
    VectorXd true_resnorms = (A_diag.asDiagonal() * X.vec + X.vec * A_shifts.asDiagonal() - B.vec).colwise().norm();

    for (Eigen::Index i = 0; i < true_resnorms.size(); ++i) {
        std::cout << "true resnorm " << i << ": " << true_resnorms[i] << '\n';
        if (true_resnorms[i] > tol * 100) {
            return 2;
        }
    }
}


///////////////////////////////////////////////////////////////
//  Copyright 2018 John Maddock. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt

//[eigen_eg
#include <iostream>
#include <boost/multiprecision/cpp_complex.hpp>
#include <boost/multiprecision/eigen.hpp>
#include <Eigen/Dense>

int main()
{
   using namespace Eigen;
   typedef boost::multiprecision::cpp_complex_quad complex_type;
   //
   // We want to solve Ax = b for x,
   // define A and b first:
   //
   Matrix<complex_type, 2, 2> A, b;
   A << complex_type(2, 3), complex_type(-1, -2), complex_type(-1, -4), complex_type(3, 6);
   b << 1, 2, 3, 1;
   std::cout << "Here is the matrix A:\n" << A << std::endl;
   std::cout << "Here is the right hand side b:\n" << b << std::endl;
   //
   // Solve for x:
   //
   Matrix<complex_type, 2, 2> x = A.fullPivHouseholderQr().solve(b);
   std::cout << "The solution is:\n" << x << std::endl;
   //
   // Compute the error in the solution by using the norms of Ax - b and b:
   //
   complex_type::value_type relative_error = (A*x - b).norm() / b.norm();
   std::cout << "The relative error is: " << relative_error << std::endl;
   return 0;
}
//]

/*
//[eigen_out
Here is the matrix A:
(2,3) (-1,-2)
(-1,-4)   (3,6)
Here is the right hand side b:
1 2
3 1
The solution is:
(0.6,-0.6)   (0.7,-0.7)
(0.64,-0.68) (0.58,-0.46)
The relative error is: 2.63132e-34
//]
*/


#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl_ros/point_cloud.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <boost/thread/thread.hpp>
#include "mongodb_store/SetParam.h"

ros::ServiceClient client;

bool is_inlier(const Eigen::Vector3f& point, const Eigen::Vector4f plane, double threshold)
{
    return fabs(point.dot(plane.segment<3>(0)) + plane(3)) < threshold;
}

void plot_best_plane(const pcl::PointCloud<pcl::PointXYZ>& points, const Eigen::Vector4f plane, double threshold)
{
    pcl::PointCloud<pcl::PointXYZ>::Ptr inlier_cloud(new pcl::PointCloud<pcl::PointXYZ>());
    pcl::PointCloud<pcl::PointXYZ>::Ptr outlier_cloud(new pcl::PointCloud<pcl::PointXYZ>());

    for (int i = 0; i < points.size(); ++i) {
        if (is_inlier(points[i].getVector3fMap(), plane, threshold)) {
            inlier_cloud->push_back(points[i]);
        }
        else {
            outlier_cloud->push_back(points[i]);
        }
    }

    pcl::visualization::PCLVisualizer viewer("3D Viewer");
    viewer.setBackgroundColor(0, 0, 0);
    viewer.addCoordinateSystem(1.0);
    viewer.initCameraParameters();

    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> inlier_color_handler(inlier_cloud, 255, 0, 0);
    viewer.addPointCloud(inlier_cloud, inlier_color_handler, "inliers");

    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> outlier_color_handler(outlier_cloud, 0, 0, 255);
    viewer.addPointCloud(outlier_cloud, outlier_color_handler, "outliers");

    while (!viewer.wasStopped())
    {
        viewer.spinOnce(100);
        boost::this_thread::sleep(boost::posix_time::microseconds(100000));
    }

}

void compute_plane(Eigen::Vector4f& plane, const pcl::PointCloud<pcl::PointXYZ>& points, int* inds)
{
    Eigen::Vector3f first = points[inds[1]].getVector3fMap() - points[inds[0]].getVector3fMap();
    Eigen::Vector3f second = points[inds[2]].getVector3fMap() - points[inds[0]].getVector3fMap();
    Eigen::Vector3f normal = first.cross(second);
    normal.normalize();
    plane.segment<3>(0) = normal;
    plane(3) = -normal.dot(points[inds[0]].getVector3fMap());
}

void extract_height_and_angle(const Eigen::Vector4f& plane)
{
    ROS_INFO("Ground plane: %f, %f, %f, %f", plane(0), plane(1), plane(2), plane(3));
    double dist = fabs(plane(3)/plane(2)); // distance along z axis
    double height = fabs(plane(3)/plane.segment<3>(0).squaredNorm()); // height
    ROS_INFO("Distance to plane along camera axis: %f", dist);
    ROS_INFO("Height above ground: %f", height);
    double angle = asin(height/dist);
    ROS_INFO("Angle radians: %f", angle);
    ROS_INFO("Angle degrees: %f", 180.0f*angle/M_PI);

    mongodb_store::SetParam srv;
    char buffer[250];

    // store height above ground in datacentre
    ros::param::set("/chest_xtion_height", height);
    sprintf(buffer, "{\"path\":\"/chest_xtion_height\",\"value\":%f}", height);
    srv.request.param = buffer;
    if (!client.call(srv)) {
        ROS_ERROR("Failed to call set height, is config manager running?");
    }

    // store angle between camera and horizontal plane
    ros::param::set("/chest_xtion_angle", angle);
    sprintf(buffer, "{\"path\":\"/chest_xtion_angle\",\"value\":%f}", angle);
    srv.request.param = buffer;
    if (!client.call(srv)) {
        ROS_ERROR("Failed to call set angle, is config manager running?");
    }
}

void callback(const sensor_msgs::PointCloud2::ConstPtr& msg)
{
    ROS_INFO("Got a pointcloud, calibrating...");
    pcl::PointCloud<pcl::PointXYZ> cloud;
    pcl::fromROSMsg(*msg, cloud);

    int nbr = cloud.size();

    int max = 1000; // ransac iterations
    double threshold = 0.02; // threshold for plane inliers

    Eigen::Vector4f best_plane; // best plane parameters found
    int best_inliers = -1; // best number of inliers

    int inds[3];

    Eigen::Vector4f plane;
    int inliers;
    for (int i = 0; i < max; ++i) {

        for (int j = 0; j < 3; ++j) {
            inds[j] = rand() % nbr; // get a random point
        }

        // check that the points aren't the same
        if (inds[0] == inds[1] || inds[0] == inds[2] || inds[1] == inds[2]) {
            continue;
        }

        compute_plane(plane, cloud, inds);
        inliers = 0;
        for (int j = 0; j < nbr; j += 30) { // count number of inliers
            if (is_inlier(cloud[j].getVector3fMap(), plane, threshold)) {
                ++inliers;
            }
        }

        if (inliers > best_inliers) {
            best_plane = plane;
            best_inliers = inliers;
        }
    }

    extract_height_and_angle(best_plane); // find parameters and feed them to datacentre
    plot_best_plane(cloud, best_plane, threshold); // visually evaluate plane fit

    exit(0);
}

int main(int argc, char** argv)
{
    ros::init(argc, argv, "calibrate_chest");
	ros::NodeHandle n;

    std::string camera_topic = "chest_xtion";
    client = n.serviceClient<mongodb_store::SetParam>("/config_manager/set_param");
	ros::Subscriber sub = n.subscribe(camera_topic + "/depth/points", 1, callback);

    ros::spin();

	return 0;
}


#include <iostream>
#include <pybind11/eigen.h>
#include <Eigen/Dense>
namespace py = pybind11;

template <typename T>
using RMatrix = Eigen::Matrix<T, -1, -1, Eigen::RowMajor>;

template <typename T>
void block_test(Eigen::Ref<RMatrix<T>> m, T a) //中身参照
{
	//プロパティ取得
	py::print(m.rows(), m.cols());
	//blockで4x4取り出して処理する
	py::print(m.block(0, 0, 4, 4));
	py::print(m.block(4, 0, 4, 4));

	Eigen::MatrixXd A = m.block(4, 0, 4, 4); // blockで取り出した行列はコピー
	py::print(A.rows(), A.cols());
	A = A * a; // ブロック元には影響しない
	py::print(A);
	m.block<4, 4>(2, 0) = Eigen::Matrix4d::Identity(); // blockに直接代入すると書き換えられる <>はサイズ固定する場合に使用
}
PYBIND11_MODULE(QT_TARGET, m)
{
	m.def("block_test", &block_test<double>, "");
}