/*
   Copyright 2015 Ruben Moreno Montoliu <ruben3d at gmail dot com>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

//
// SphereAreaLight.cpp - Created on 2009.01.15
//

#include <cmath>
#include <string>
#include <boost/lexical_cast.hpp>
#include "SphereAreaLight.h"
#include "Object.h"
#include "Polar3.h"
#include "Random.h"
#include "Ray.h"
#include "Sphere.h"
#include "Material.h"

unsigned int SphereAreaLight::GlobalCounter = 0;

SphereAreaLight::SphereAreaLight(const Transform& tr, const double intensity, const Color& color,
								const double radius, const bool show, const Attenuation attenuation)
	: Light(tr, intensity, color), LightHasPosition(tr(Point(0,0,0))),
		LightHasAttenuation(attenuation), LightHasGeometry(show), m_radius(radius),
		m_minSamples(16), m_maxSamples(32)
{
	Light::addAttribute(LightHasPosition::ATTRIBUTE);
	Light::addAttribute(LightHasAttenuation::ATTRIBUTE);
	Light::addAttribute(LightHasGeometry::ATTRIBUTE);

	m_id = ++GlobalCounter;
}

bool SphereAreaLight::computeSample(const Point& P, const Vector& N,
							const std::list<Object*>& objects,
							unsigned int& haltonSeq,
							Color& I, Vector& L) const
{
	// Light surface normal at sample point
	Vector LN;

	// Generates a random sample. We don't want a sample facing away
	do
	{
		haltonSeq++;
		double rotPhi = acos(Random::Instance().haltonSeq(haltonSeq,2)*2-1);
		double rotTheta = Random::Instance().haltonSeq(haltonSeq,3)*PI*2;
		Polar3 p(1,rotPhi,rotTheta);
		LN=p;
	}
	while (N.dot(LN) <= 0);

	// Point at the light surface
	Point LP = /*getPosition() +*/ Point(LN * m_radius);

	// Ray from point of intersection to light surface
	L = LP - P;
	const double lightRayLength = L.length();
	L = L.normalize();
	Ray lightRay = T(Ray(P, L));	// <- Local to World
	lightRay.tmax = lightRayLength;

	// Test if the light ray is obstructed
	bool obstructed = false;
	std::list<Object*>::const_iterator obj;
	for (obj = objects.begin(); obj != objects.end(); obj++)
	{
		if (!(*obj)->castShadow()) continue;

		double t;
		const Ray lightRayT = (*obj)->T.applyInv(lightRay);	// World to object
		bool intersects = (*obj)->getGeometry()->intersects(lightRayT,t);
		if (intersects)
		{
			obstructed = true;
			break;
		}
	}

	I = obstructed ? Color(0,0,0) : getColor()*getIntensity()*getAttFactor(lightRayLength)*L.dot(LN);
	return obstructed;
}

void SphereAreaLight::computeIntensity(const Vector& N, const Point& P,
							const std::list<Object*>& objects,
							std::list<Color>& Is, std::list<Vector>& Ls) const
{
	bool penumbra = false;
	unsigned int haltonSeq = Random::Instance().generate(1024);

	Vector tN = T.applyInv(N);	// World to local
	Point tP = T.applyInv(P);
	Color I;
	Vector L;
	bool shadow = computeSample(tP,tN,objects,haltonSeq,I,L);
	Is.push_back(I);
	Ls.push_back(T(L));	// Local to world

	for (int i=1; i<m_minSamples; i++)
	{
		Color I;
		Vector L;
		bool blocked = computeSample(tP,tN,objects,haltonSeq,I,L);
		Is.push_back(I);
		Ls.push_back(T(L));	// Local to world

		if (blocked != shadow) penumbra = true;
	}

	if (penumbra)
	{
		for (int n=0; n<m_maxSamples; n++)
		{
			Color I;
			Vector L;
			computeSample(tP,tN,objects,haltonSeq,I,L);
			Is.push_back(I);
			Ls.push_back(T(L));	// Local to world
		}
	}
}

Geometry* SphereAreaLight::getGeometry() const
{
	return new Sphere(m_radius);
}

Material* SphereAreaLight::getMaterial() const
{
	Material *m = new Material(std::string("_sal_")+boost::lexical_cast<std::string>(m_id), getColor());
	m->setEmission(getIntensity());
	m->setSpecularIntensity(0.0);
	m->setSpecularGlossiness(0.0);
	return m;
}


#include <starpu.h>
#include <starpu_mpi.h>
#include <vector>
#include <memory>
#include <random>
#include <iostream>
#ifdef USE_MKL
#include <mkl_cblas.h>
#include <mkl_lapacke.h>
#else
#include <cblas.h>
#include <lapacke.h>
#endif
#include <Eigen/Core>
#include <Eigen/Cholesky>
#include <mpi.h>

using namespace std;
using namespace Eigen;

void potrf(void *buffers[], void *cl_arg) {
    double *A = (double *)STARPU_MATRIX_GET_PTR(buffers[0]);
    int m = STARPU_MATRIX_GET_NX(buffers[0]);
    int n = STARPU_MATRIX_GET_NY(buffers[0]);
    assert(m == n);
    LAPACKE_dpotrf(LAPACK_COL_MAJOR, 'L', m, A, m);
}
struct starpu_codelet potrf_cl = {
    .where = STARPU_CPU,
    .cpu_funcs = { potrf, NULL },
    .nbuffers = 1,
    .modes = { STARPU_RW }
};

void trsm(void *buffers[], void *cl_arg) {
    double *A = (double *)STARPU_MATRIX_GET_PTR(buffers[0]);
    double *B = (double *)STARPU_MATRIX_GET_PTR(buffers[1]);
    int m = STARPU_MATRIX_GET_NX(buffers[1]);
    int n = STARPU_MATRIX_GET_NY(buffers[1]);
    assert(STARPU_MATRIX_GET_NX(buffers[0]) == n);
    assert(STARPU_MATRIX_GET_NY(buffers[0]) == n);
    cblas_dtrsm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit, m, n, 1.0, A, n, B, m);
}
struct starpu_codelet trsm_cl = {
    .where = STARPU_CPU,
    .cpu_funcs = { trsm, NULL },
    .nbuffers = 2,
    .modes = { STARPU_R, STARPU_RW }
};

void syrk(void *buffers[], void *cl_arg) {
    double *A = (double *)STARPU_MATRIX_GET_PTR(buffers[0]);
    double *C = (double *)STARPU_MATRIX_GET_PTR(buffers[1]);
    int m = STARPU_MATRIX_GET_NX(buffers[0]);
    int k = STARPU_MATRIX_GET_NY(buffers[0]);
    assert(STARPU_MATRIX_GET_NX(buffers[1]) == m);
    assert(STARPU_MATRIX_GET_NY(buffers[1]) == m);
    cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, m, k, -1.0, A, m, 1.0, C, m);
}
struct starpu_codelet syrk_cl = {
    .where = STARPU_CPU,
    .cpu_funcs = { syrk, NULL },
    .nbuffers = 2,
    .modes = { STARPU_R, STARPU_RW }
};

void gemm(void *buffers[], void *cl_arg) {
    double *A = (double *)STARPU_MATRIX_GET_PTR(buffers[0]);
    double *B = (double *)STARPU_MATRIX_GET_PTR(buffers[1]);
    double *C = (double *)STARPU_MATRIX_GET_PTR(buffers[2]);
    int m = STARPU_MATRIX_GET_NX(buffers[2]);
    int n = STARPU_MATRIX_GET_NY(buffers[2]);
    int k = STARPU_MATRIX_GET_NY(buffers[0]);
    assert(STARPU_MATRIX_GET_NX(buffers[0]) == m);
    assert(STARPU_MATRIX_GET_NX(buffers[1]) == n);
    assert(STARPU_MATRIX_GET_NY(buffers[1]) == k);
    cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans, m, n, k, -1.0, A, m, B, n, 1.0, C, m);
  }
struct starpu_codelet gemm_cl = {
    .where = STARPU_CPU,
    .cpu_funcs = { gemm, NULL },
    .nbuffers = 3,
    .modes = { STARPU_R, STARPU_R, STARPU_RW }
};

void cholesky(const int block_size, const int num_blocks, const int rank, const int size, const int test, const int nrow, const int ncol, const bool prune, const int upper_block_size) {
    auto val = [&](int i, int j) { return  1.0/(double)((i-j)*(i-j)+1.0); };
    vector<MatrixXd*> blocks(num_blocks*num_blocks);
    vector<starpu_data_handle_t> dataA(num_blocks*num_blocks);
    auto block_2_rank = [&](int i, int j){return (i % nrow) * ncol + j % ncol;};
    const int ncores = starpu_worker_get_count_by_type(STARPU_CPU_WORKER);
    const int matrix_size = block_size * num_blocks;

    // Warmup MKL
    {
        Eigen::MatrixXd A = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd B = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd C = Eigen::MatrixXd::Identity(256,256);
        for(int i = 0; i < 10; i++) {
            cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, 256, 256, 256, 1.0, A.data(), 256, B.data(), 256, 1.0, C.data(), 256);
        }
    }

    std::mt19937 gen(2020);
    assert(upper_block_size <= 2*block_size);
    const int lower_block_size = 2*block_size - upper_block_size;
    std::uniform_int_distribution<> distrib(lower_block_size,upper_block_size); // average is block_size
    if(rank == 0) printf("lower_block_size %d, upper_block_size %d\n", upper_block_size, lower_block_size);
    std::vector<int> block_sizes(num_blocks, block_size);
    {
        int n = 0;
        for(int i = 0; i < num_blocks-1; i++) {
            int bs = std::min(matrix_size - n, distrib(gen));
            n += bs;
            block_sizes[i] = bs;
        }
        assert(matrix_size - n >= 0);
        block_sizes[num_blocks-1] = matrix_size - n;
    }
    int total = std::accumulate(block_sizes.begin(), block_sizes.end(), 0);
    assert(total == matrix_size);
    std::vector<int> block_displ(num_blocks+1, 0);
    for(int i = 1; i < num_blocks+1; i++) {
        block_displ[i] = block_displ[i-1] + block_sizes[i-1];
    }
    assert(block_displ[num_blocks] == matrix_size);
    if(rank == 0) {
        printf("block sizes: ");
        for(int i = 0; i < num_blocks; i++) {
            assert(block_sizes[i] >= 0);
            printf("%d ", block_sizes[i]);
        };
        printf("\n");
        printf("block displ: ");
        for(int i = 0; i < num_blocks+1; i++) {
            assert(block_displ[i] >= 0);
            printf("%d ", block_displ[i]);
        };
        printf("\n");
    }

    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<=ii; jj++) {
            int mpi_rank = block_2_rank(ii,jj);
            blocks[ii+jj*num_blocks] = new MatrixXd();
            if (mpi_rank == rank) {
                auto val_block = [&](int i, int j) { return val(block_displ[ii]+i,block_displ[jj]+j); };
                *blocks[ii+jj*num_blocks] = MatrixXd::NullaryExpr(block_sizes[ii], block_sizes[jj], val_block);
            }
        }
    }

    auto is_block_non_empty = [&](int i, int j){ return block_sizes[i] > 0 && block_sizes[j] > 0; };

    size_t num_pruned = 0;
    starpu_mpi_barrier(MPI_COMM_WORLD);
    double start = starpu_timing_now();

    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<=ii; jj++) {
            int mpi_rank = block_2_rank(ii,jj);
            if(is_block_non_empty(ii,jj)){
                if (mpi_rank == rank) {
                    starpu_matrix_data_register(&dataA[ii+jj*num_blocks], STARPU_MAIN_RAM, (uintptr_t)blocks[ii+jj*num_blocks]->data(), block_sizes[ii], block_sizes[ii], block_sizes[jj], sizeof(double));
                } else {
                    starpu_matrix_data_register(&dataA[ii+jj*num_blocks], -1, (uintptr_t)NULL, block_sizes[ii], block_sizes[ii], block_sizes[jj], sizeof(double));
                }
                if (dataA[ii+jj*num_blocks]) {
                    starpu_mpi_data_register(dataA[ii+jj*num_blocks], ii+jj*num_blocks, mpi_rank);
                }
            }
        }
    }

    for (int kk = 0; kk < num_blocks; ++kk) {
        // POTRF
        if( (!prune) || block_2_rank(kk,kk) == rank) {
            if(is_block_non_empty(kk,kk)){
                starpu_mpi_task_insert(MPI_COMM_WORLD,&potrf_cl,
                    STARPU_RW, dataA[kk+kk*num_blocks],
                0);
            }
        } else {
            num_pruned++;
        }

        for (int ii = kk+1; ii < num_blocks; ++ii) {
            // TRSM
            if( (!prune) || block_2_rank(kk,kk) == rank || block_2_rank(ii,kk) == rank) {
                if(is_block_non_empty(ii,kk)){
                    starpu_mpi_task_insert(MPI_COMM_WORLD,&trsm_cl,
                        STARPU_R,  dataA[kk+kk*num_blocks],
                        STARPU_RW, dataA[ii+kk*num_blocks],
                    0);
                }
            } else {
                num_pruned++;
            }
            if(block_sizes[kk] > 0) {
                starpu_mpi_cache_flush(MPI_COMM_WORLD, dataA[kk+kk*num_blocks]);
            }

            // SYRK
            if( (!prune) || block_2_rank(ii,kk) == rank || block_2_rank(ii,ii) == rank) {
                if(is_block_non_empty(ii,kk)){
                    starpu_mpi_task_insert(MPI_COMM_WORLD,&syrk_cl,
                        STARPU_R,  dataA[ii+kk*num_blocks],
                        STARPU_RW, dataA[ii+ii*num_blocks],
                    0);
                }
            } else {
                num_pruned++;
            }

            for (int jj = kk+1; jj < ii; ++jj) {
                // GEMM
                if( (!prune) || block_2_rank(ii,kk) == rank || block_2_rank(jj,kk) == rank || block_2_rank(ii,jj) == rank) {
                    if(is_block_non_empty(ii,kk) && is_block_non_empty(jj,kk)){
                        starpu_mpi_task_insert(MPI_COMM_WORLD,&gemm_cl,
                            STARPU_R,  dataA[ii+kk*num_blocks],
                            STARPU_R,  dataA[jj+kk*num_blocks],
                            STARPU_RW, dataA[ii+jj*num_blocks],
                        0);
                    }
                } else {
                    num_pruned++;
                }
            }
            if(block_sizes[kk] > 0 && block_sizes[ii] > 0) {
                starpu_mpi_cache_flush(MPI_COMM_WORLD, dataA[ii+kk*num_blocks]);
            }
        }
    }

    double end_insertion = starpu_timing_now();
    starpu_task_wait_for_all();
    starpu_mpi_barrier(MPI_COMM_WORLD);
    double end = starpu_timing_now();

    // Makes grep/import to excel easier ; just do
    // cat output | grep -P '\[0\]\>\>\>\>'
    // to extract rank 0 info
    printf(">>>>test rank nranks ncores matrix_size block_size num_blocks total_time insertion_time prune num_pruned upper_block_size\n");
    printf("[%d]>>>>chol_starpu %d %d %d %d %d %d %e %e %d %zd %d\n",rank,rank,size,ncores,matrix_size,block_size,num_blocks,(end-start)/1e6,(end_insertion-start)/1e6,prune,num_pruned,upper_block_size);

    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<=ii; jj++) {
            if(block_sizes[ii] > 0 && block_sizes[jj] > 0) {
                starpu_data_unregister(dataA[ii+jj*num_blocks]);
            }
        }
    }

    if (test) {
        printf("Testing...\n");
        for (int ii=0; ii<num_blocks; ii++) {
            for (int jj=0; jj<=ii; jj++) {
                int mpi_rank = block_2_rank(ii,jj);
                if (rank == 0 && rank != mpi_rank) {
                    blocks[ii+jj*num_blocks] = new MatrixXd(block_sizes[ii], block_sizes[jj]);
                    MPI_Recv(blocks[ii+jj*num_blocks]->data(), block_sizes[ii]*block_sizes[jj], MPI_DOUBLE, mpi_rank, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                } else if (rank == mpi_rank && rank != 0) {
                    MPI_Send(blocks[ii+jj*num_blocks]->data(), block_sizes[ii]*block_sizes[jj], MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
                }
            }
        }
        if (rank==0) {
            MatrixXd B = MatrixXd::NullaryExpr(matrix_size, matrix_size, val);
            MatrixXd L = MatrixXd::Zero(matrix_size, matrix_size);
            for (int ii=0; ii<num_blocks; ii++) {
                for (int jj=0; jj<num_blocks; jj++) {
                    if (jj<=ii)  {
                        L.block(block_displ[ii],block_displ[jj],block_sizes[ii],block_sizes[jj])=*blocks[ii+jj*num_blocks];
                    }
                }
            }
            auto L1=L.triangularView<Lower>();
            VectorXd x = VectorXd::Random(matrix_size);
            VectorXd b = B*x;
            VectorXd bref = b;
            L1.solveInPlace(b);
            L1.transpose().solveInPlace(b);
            double error = (b - x).norm() / x.norm();
            printf("\nError solve: %e\n\n", error);
            assert(error < 1e-6);
        }
    }
}

int main(int argc, char **argv)
{
    int req = MPI_THREAD_SERIALIZED;
    int prov = -1;
    MPI_Init_thread(NULL, NULL, req, &prov);
    starpu_mpi_init_conf(&argc, &argv, 0, MPI_COMM_WORLD, NULL);
    int rank, size;
    starpu_mpi_comm_rank(MPI_COMM_WORLD, &rank);
    starpu_mpi_comm_size(MPI_COMM_WORLD, &size);
    if (rank==0) {
        printf("Running on %d CPU cores per rank,", starpu_worker_get_count_by_type(STARPU_CPU_WORKER));
        printf("and %d ranks in total\n", size);
    }
    int block_size=10;
    int num_blocks=1;
    int test=0;
    int nrow=1;
    int ncol=1;
    bool prune = true;
    int upper_block_size = block_size;

    if (argc >= 2)
    {
        block_size = atoi(argv[1]);
    }
    if (argc >= 3)
    {
        num_blocks = atoi(argv[2]);
    }

    if (argc >= 4) {
        test = atoi(argv[3]);
    }

    if (argc >= 6) {
        nrow = atoi(argv[4]);
        ncol = atoi(argv[5]);
    }

    if (argc >= 7) {
        prune = atoi(argv[6]);
    }

    if (argc >= 8) {
        upper_block_size = atoi(argv[7]);
    } else {
        upper_block_size = block_size;
    }

    assert(nrow * ncol == size);
    printf("Usage: ./cholesky_mpi block_size num_blocks test nrow ncol prune upper_block_size\n");
    printf("block_size,%d\n",block_size);
    printf("num_blocks,%d\n",num_blocks);
    printf("test,%d\n",test);
    printf("nprocs_row,%d\n",nrow);
    printf("nprocs_col,%d\n",ncol);
    printf("prune,%d\n",prune);
    printf("upper_block_size,%d\n",upper_block_size);
    cholesky(block_size, num_blocks, rank, size, test, nrow, ncol, prune, upper_block_size);
    starpu_mpi_shutdown();
    MPI_Finalize();
    return 0;
}


// An interval object.
#pragma once

#include <string>

#include <boost/numeric/interval.hpp>

#define USE_FILIB_INTERVALS
#ifdef USE_FILIB_INTERVALS
#include <interval/filib_rounding.hpp>
#endif

#include <utils/eigen_ext.hpp>

namespace ipc::rigid {

namespace interval_options {
    typedef boost::numeric::interval_lib::checking_catch_nan<double>
        CheckingPolicy;
} // namespace interval_options

#ifdef USE_FILIB_INTERVALS

// Use filib rounding arithmetic
typedef boost::numeric::interval<
    double,
    boost::numeric::interval_lib::policies<
        boost::numeric::interval_lib::save_state<FILibRounding>,
        interval_options::CheckingPolicy>>
    Interval;

#elif defined(__APPLE__)

// clang-format off
#warning "Rounding modes seem to be broken with trigonometric functions on macOS, unable to compute exact interval arithmetic!"
// clang-format on
typedef boost::numeric::interval<
    double,
    boost::numeric::interval_lib::policies<
        boost::numeric::interval_lib::save_state<
            boost::numeric::interval_lib::rounded_transc_exact<double>>,
        interval_options::CheckingPolicy>>
    Interval;

#else

// Use proper rounding arithmetic
typedef boost::numeric::interval<
    double,
    boost::numeric::interval_lib::policies<
        boost::numeric::interval_lib::save_state<
            boost::numeric::interval_lib::rounded_transc_std<double>>,
        interval_options::CheckingPolicy>>
    Interval;

#endif // USE_FILIB_INTERVALS

template <typename Derived>
inline Eigen::VectorXd width(const Eigen::MatrixBase<Derived>& x)
{
    Eigen::VectorXd w(x.size());
    for (int i = 0; i < x.size(); i++) {
        w(i) = width(x(i));
    }
    return w;
}

template <typename Derived>
inline double diagonal_width(const Eigen::MatrixBase<Derived>& x)
{
    Eigen::VectorXd widths = width(x);
    double w = 0;
    for (int i = 0; i < widths.size(); i++) {
        w += widths(i) * widths(i);
    }
    return sqrt(w);
}

template <typename Derived>
inline bool zero_in(const Eigen::MatrixBase<Derived>& x)
{
    // Check if the origin is in the n-dimensional interval
    for (int i = 0; i < x.size(); i++) {
        if (!boost::numeric::zero_in(x(i))) {
            return false;
        }
    }
    return true;
}

typedef Vector2<Interval> Vector2I;
typedef Vector3<Interval> Vector3I;
typedef VectorX<Interval> VectorXI;
typedef VectorMax3<Interval> VectorMax3I;
typedef Matrix3<Interval> Matrix2I;
typedef Matrix3<Interval> Matrix3I;
typedef MatrixMax3<Interval> MatrixMax3I;
typedef MatrixX<Interval> MatrixXI;

/// @brief Format a string for an Interval
std::string fmt_interval(const Interval& i, const int precision = 16);
/// @brief Format an eigen VectorX<Interval>
std::string fmt_eigen_intervals(const VectorXI& x, const int precision = 16);
} // namespace ipc::rigid

namespace Eigen {

template <typename BinOp>
struct ScalarBinaryOpTraits<ipc::rigid::Interval, double, BinOp> {
    typedef ipc::rigid::Interval ReturnType;
};

template <typename BinOp>
struct ScalarBinaryOpTraits<double, ipc::rigid::Interval, BinOp> {
    typedef ipc::rigid::Interval ReturnType;
};

#if EIGEN_MAJOR_VERSION >= 3
namespace internal {
    template <typename X, typename S, typename P>
    struct is_convertible<X, boost::numeric::interval<S, P>> {
        enum { value = is_convertible<X, S>::value };
    };

    template <typename S, typename P1, typename P2>
    struct is_convertible<
        boost::numeric::interval<S, P1>,
        boost::numeric::interval<S, P2>> {
        enum { value = true };
    };
} // namespace internal
#endif
} // namespace Eigen


#ifndef RECURSION_SEARCH_HPP
#define RECURSION_SEARCH_HPP
/**
 * @file recursion_search.hpp
 *
 * @brief search parameters so that the random number generator's state
 * transition function has an irreducible characteristic polynomial.
 *
 * @author Mutsuo Saito (Hiroshima University)
 * @author Makoto Matsumoto (The University of Tokyo)
 *
 * Copyright (C) 2011 Mutsuo Saito, Makoto Matsumoto,
 * Hiroshima University and The University of Tokyo.
 * All rights reserved.
 *
 * The 3-clause BSD License is applied to this software, see
 * LICENSE.txt
 */
#include <NTL/GF2X.h>
#include <NTL/GF2XFactoring.h>

namespace tinymt {
    /**
     * @class Search
     * search parameters so that the generator's state transition function
     * has an irreducible characteristic polynomial.
     * 1) call start() function.
     * 2) if start() returns true, then call get_random(), get_minpoly(),
     * or get_count().
     *
     * @tparam T generators class
     * @tparam SG sequential generator
     */
    template<class T, class SG> class Search {
    public:
        void get_minpoly(NTL::GF2X& minpoly, T& rand) {
            using namespace std;
            using namespace NTL;

            vec_GF2 vec;
            int mexp = rand.get_mexp();
            vec.SetLength(mexp * 2);
            for (int i = 0; i < mexp * 2; i++) {
                vec[i] = rand.generate() & 1;
            }
            MinPolySeq(minpoly, vec, mexp);
        }

        /**
         * generate random parameters and check if the generator's
         * state transition function has an irreducible characteristic
         * polynomial. If found in \b try_count times, return true, else
         * return false.
         *
         * @param try_count
         */
        bool start(int try_count) {
            long mexp = rand.get_mexp();
            long degree;
            for (int i = 0; i < try_count; i++) {
                rand.setup_param(sg->next());
                rand.seeding(1);
                get_minpoly(minpoly, rand);
                count++;
                degree = deg(minpoly);
                if (degree != mexp) {
                    continue;
                }
                if (IterIrredTest(minpoly)) {
                    return true;
                }
            }
            return false;
        }
        /**
         * call this function after \b start() has returned true.
         * @return random number generator class with parameters.
         */
        const T& get_random() const {
            return rand;
        }

        /**
         * call this function after \b start() has returned true.
         * In this program, if minimal polynomial is irreducible,
         * then the polynomial is characteristic polynomial of
         * generator's state transition function.
         *
         * @return minimal polynomial of generated sequence.
         */
        const NTL::GF2X& get_minpoly() const {
            return minpoly;
        }

        /**
         * @return tried count after this class has created.
         */
        int get_count() const {
            return count;
        }

        /**
         * @param rand_ random number generator whose parameters are
         * searched.
         * @param seq_generator a sequential number generator which
         * gives sequential number for searching parameters.
         *
         */
        Search(const T& rand_, SG& seq_generator) :
            rand(rand_)
            {
                sg = &seq_generator;
                count = 0;
            }
    private:
        T rand;
        NTL::GF2X minpoly;
        SG *sg;
        int count;
    };
}
#endif


#include <iostream>
#include <fstream>
#include <strstream>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <ceres/ceres.h>
#include <igl/point_mesh_squared_distance.h>
#include <igl/copyleft/marching_cubes.h>

typedef double FT;
typedef Eigen::Matrix<FT, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
typedef Eigen::Matrix<FT, Eigen::Dynamic, 1> VectorX;
typedef Eigen::Matrix<FT, 3, 1> Vector3;

class UniformGrid
{
public:
	UniformGrid()
	: N(0)
	{}
	UniformGrid(int _N) {
		N = _N;
		distances.resize(N);
		for (auto& d : distances) {
			d.resize(N);
			for (auto& v : d)
				v.resize(N, 1e30);
		}
	}
	template <class T>
	T distance(const T* const p) const {
		int px = *(double*)&p[0] * N;
		int py = *(double*)&p[1] * N;
		int pz = *(double*)&p[2] * N;
		if (px < 0 || py < 0 || pz < 0 || px >= N - 1 || py >= N - 1 || pz >= N - 1) {
			T l = (T)0;
			if (px < 0)
				l = l + -p[0] * (T)N;
			else if (px >= N)
				l = l + (p[0] * (T)N - (T)(N - 1 - 1e-3));

			if (py < 0)
				l = l + -p[1] * (T)N;
			else if (py >= N)
				l = l + (p[1] * (T)N - (T)(N - 1 - 1e-3));

			if (pz < 0)
				l = l + -p[2] * (T)N;
			else if (pz >= N)
				l = l + (p[2] * (T)N - (T)(N - 1 - 1e-3));

			return l;
		}
		T wx = p[0] * (T)N - (T)px;
		T wy = p[1] * (T)N - (T)py;
		T wz = p[2] * (T)N - (T)pz;
		T w0 = ((T)1 - wx) * ((T)1 - wy) * ((T)1 - wz) * distances[pz    ][py    ][px    ];
		T w1 = wx 		   * ((T)1 - wy) * ((T)1 - wz) * distances[pz    ][py    ][px + 1];
		T w2 = ((T)1 - wx) * wy 		 * ((T)1 - wz) * distances[pz    ][py + 1][px    ];
		T w3 = wx 		   * wy 		 * ((T)1 - wz) * distances[pz    ][py + 1][px + 1];
		T w4 = ((T)1 - wx) * ((T)1 - wy) * wz 		   * distances[pz + 1][py    ][px    ];
		T w5 = wx 		   * ((T)1 - wy) * wz 		   * distances[pz + 1][py    ][px + 1];
		T w6 = ((T)1 - wx) * wy 		 * wz		   * distances[pz + 1][py + 1][px    ];
		T w7 = wx 		   * wy 		 * wz 		   * distances[pz + 1][py + 1][px + 1];
		T res = w0 + w1 + w2 + w3 + w4 + w5 + w6 + w7;
		T thres = (T)0.02;

		//commented out for deform_tune
		if (res > thres)
			res = thres;
		//

		return res;
	}
	int N;
	std::vector<std::vector<std::vector<double> > > distances;
};

struct LengthError {
  LengthError(const Eigen::Vector3d& v_, double lambda_)
  : v(v_), lambda(lambda_) {}

  template <typename T>
  bool operator()(const T* const p1,
                  const T* const p2,
                  T* residuals) const {
  	T px = p1[0] - p2[0];
  	T py = p1[1] - p2[1];
  	T pz = p1[2] - p2[2];
  	residuals[0] = px - v[0];
  	residuals[1] = py - v[1];
  	residuals[2] = pz - v[2];
    return true;
  }

   // Factory to hide the construction of the CostFunction object from
   // the client code.
   static ceres::CostFunction* Create(const Eigen::Vector3d& v, const double lambda_) {
     return (new ceres::AutoDiffCostFunction<LengthError, 3, 3, 3>(
                 new LengthError(v, lambda_)));
   }
   double lambda;
   Eigen::Vector3d v;
};

struct DistanceError {
  DistanceError(UniformGrid* grid_)
  : grid(grid_) {}

  template <typename T>
  bool operator()(const T* const p1,
                  T* residuals) const {
  	residuals[0] = grid->distance(p1);
  	residuals[1] = (T)0;
  	residuals[2] = (T)0;
    return true;
  }

   // Factory to hide the construction of the CostFunction object from
   // the client code.
   static ceres::CostFunction* Create(UniformGrid* grid) {
     return (new ceres::AutoDiffCostFunction<DistanceError, 3, 3>(
                 new DistanceError(grid)));
   }
   UniformGrid* grid;
};

class Mesh
{
public:
	Mesh() : scale(1.0) {}
	std::vector<Eigen::Vector3d> V;
	std::vector<Eigen::Vector3i> F;
	void ReadOBJ(const char* filename) {
		std::ifstream is(filename);
		char buffer[256];
		while (is.getline(buffer, 256)) {
			std::strstream str;
			str << buffer;
			str >> buffer;
			if (strcmp(buffer, "v") == 0) {
				double x, y, z;
				str >> x >> y >> z;
				V.push_back(Eigen::Vector3d(x, y, z));
			}
			else if (strcmp(buffer, "f") == 0) {
				Eigen::Vector3i f;
				for (int j = 0; j < 3; ++j) {
					str >> buffer;
					int id = 0;
					int p = 0;
					while (buffer[p] != '/') {
						id = id * 10 + (buffer[p] - '0');
						p += 1;
					}
					f[j] = id - 1;
				}
				F.push_back(f);
			}
		}
	}

	void ReadOBJ_Manifold(const char* filename) {
		std::ifstream is(filename);
		char buffer[256];
		while (is.getline(buffer, 256)) {
			std::strstream str;
			str << buffer;
			str >> buffer;
			if (strcmp(buffer, "v") == 0) {
				double x, y, z;
				str >> x >> y >> z;
				V.push_back(Eigen::Vector3d(x, y, z));
			}
			else if (strcmp(buffer, "f") == 0) {
				Eigen::Vector3i f;
				int idx_x, idx_y, idx_z;
				str >> idx_x >> idx_y >> idx_z;
				// std::cout<<idx_x <<' '<< idx_y <<' '<< idx_z <<std::endl;
				f = Eigen::Vector3i(idx_x-1, idx_y-1, idx_z-1);
				F.push_back(f);
			}
		}
	}

	void WriteOBJ(const char* filename) {
		std::ofstream os(filename);
		for (int i = 0; i < V.size(); ++i) {
			os << "v " << V[i][0] << " " << V[i][1] << " " << V[i][2] << "\n";
		}
		for (int i = 0; i < F.size(); ++i) {
			os << "f " << F[i][0] + 1 << " " << F[i][1] + 1 << " " << F[i][2] + 1 << "\n";
		}
		os.close();
	}
	double scale;
	Eigen::Vector3d pos;
	void Normalize() {
		double min_p[3], max_p[3];
		for (int j = 0; j < 3; ++j) {
			min_p[j] = 1e30;
			max_p[j] = -1e30;
			for (int i = 0; i < V.size(); ++i) {
				if (V[i][j] < min_p[j])
					min_p[j] = V[i][j];
				if (V[i][j] > max_p[j])
					max_p[j] = V[i][j];
			}
		}
		scale = std::max(max_p[0] - min_p[0], std::max(max_p[1] - min_p[1], max_p[2] - min_p[2])) * 1.1;
		for (int j = 0; j < 3; ++j)
			pos[j] = min_p[j] - 0.05 * scale;
		for (auto& v : V) {
			v = (v - pos) / scale;
		}
		for (int j = 0; j < 3; ++j) {
			min_p[j] = 1e30;
			max_p[j] = -1e30;
			for (int i = 0; i < V.size(); ++i) {
				if (V[i][j] < min_p[j])
					min_p[j] = V[i][j];
				if (V[i][j] > max_p[j])
					max_p[j] = V[i][j];
			}
		}
	}
	void ApplyTransform(Mesh& m) {
		pos = m.pos;
		scale = m.scale;
		for (auto& v : V) {
			v = (v - pos) / scale;
		}
	}
	void ConstructDistanceField(UniformGrid& grid) {
		Eigen::MatrixXd P(grid.N * grid.N * grid.N, 3);
		int offset = 0;
		for (int i = 0; i < grid.N; ++i) {
			for (int j = 0; j < grid.N; ++j) {
				for (int k = 0; k < grid.N; ++k) {
					P.row(offset) = Eigen::Vector3d(double(k) / grid.N, double(j) / grid.N, double(i) / grid.N);
					offset += 1;
				}
			}
		}

		Eigen::MatrixXd V2(V.size(), 3);
		for (int i = 0; i < V.size(); ++i)
			V2.row(i) = V[i];

		Eigen::MatrixXi F2(F.size(), 3);
		for (int i = 0; i < F.size(); ++i)
			F2.row(i) = F[i];

		Eigen::MatrixXd N(F.size(), 3);
		for (int i = 0; i < F.size(); ++i) {
			Eigen::Vector3d x = V[F[i][1]] - V[F[i][0]];
			Eigen::Vector3d y = V[F[i][2]] - V[F[i][0]];
			N.row(i) = x.cross(y).normalized();
		}

		Eigen::VectorXd sqrD;
		Eigen::VectorXi I;
		Eigen::MatrixXd C;
		igl::point_mesh_squared_distance(P,V2,F2,sqrD,I,C);

		offset = 0;

		for (int i = 0; i < grid.N; ++i) {
			for (int j = 0; j < grid.N; ++j) {
				for (int k = 0; k < grid.N; ++k) {
					Eigen::Vector3d n = N.row(I[offset]);
					Eigen::Vector3d off = P.row(offset);
					off -= V[F[I[offset]][0]];
					double d = n.dot(off);
					d = 1;
					if (d > 0)
						grid.distances[i][j][k] = sqrt(sqrD[offset]);
					else
						grid.distances[i][j][k] = sqrt(-sqrD[offset]);
					offset += 1;
				}
			}
		}

	}

	void FromDistanceField(UniformGrid& grid) {
		Eigen::VectorXd S(grid.N * grid.N * grid.N);
		Eigen::MatrixXd GV(grid.N * grid.N * grid.N, 3);
		int offset = 0;
		for (int i = 0; i < grid.N; ++i) {
			for (int j = 0; j < grid.N; ++j) {
				for (int k = 0; k < grid.N; ++k) {
					S[offset] = grid.distances[i][j][k];
					GV.row(offset) = Eigen::Vector3d(k, j, i);
					offset += 1;
				}
			}
		}
		Eigen::MatrixXd SV;
		Eigen::MatrixXi SF;
		igl::copyleft::marching_cubes(S,GV,grid.N,grid.N,grid.N,SV,SF);
		V.resize(SV.rows());
		F.resize(SF.rows());
		for (int i = 0; i < SV.rows(); ++i)
			V[i] = SV.row(i) / (double)grid.N;
		for (int i = 0; i < SF.rows(); ++i)
			F[i] = SF.row(i);
	}

	void Deform(UniformGrid& grid) {
		double lambda = 1e-3;
		ceres::Problem problem;

		//Move vertices
		std::vector<ceres::ResidualBlockId> v_block_ids;
		for (int i = 0; i < V.size(); ++i) {
			ceres::CostFunction* cost_function = DistanceError::Create(&grid);
			ceres::ResidualBlockId block_id = problem.AddResidualBlock(cost_function, 0, V[i].data());
			v_block_ids.push_back(block_id);
		}

		//Enforce rigidity
		std::vector<ceres::ResidualBlockId> edge_block_ids;
		for (int i = 0; i < F.size(); ++i) {
			for (int j = 0; j < 3; ++j) {
				Eigen::Vector3d v = (V[F[i][j]] - V[F[i][(j + 1) % 3]]);
				ceres::CostFunction* cost_function = LengthError::Create(v, lambda);
				ceres::ResidualBlockId block_id = problem.AddResidualBlock(cost_function, 0, V[F[i][j]].data(), V[F[i][(j + 1) % 3]].data());
				edge_block_ids.push_back(block_id);
			}
		}

		ceres::Solver::Options options;
		options.max_num_iterations = 100;
		options.linear_solver_type = ceres::SPARSE_SCHUR;
		options.minimizer_progress_to_stdout = true;
		ceres::Solver::Summary summary;
		ceres::Solve(options, &problem, &summary);
		std::cout << summary.FullReport() << "\n";

		//V error
		ceres::Problem::EvaluateOptions v_options;
		v_options.residual_blocks = v_block_ids;
		double v_cost;
		problem.Evaluate(v_options, &v_cost, NULL, NULL, NULL);
		std::cout<<"Vertices cost: "<<v_cost<<std::endl;

		//E error
		ceres::Problem::EvaluateOptions edge_options;
		edge_options.residual_blocks = edge_block_ids;
		double edge_cost;
		problem.Evaluate(edge_options, &edge_cost, NULL, NULL, NULL);
		std::cout<<"Rigidity cost: "<<edge_cost<<std::endl;

		double final_cost = v_cost + edge_cost;
		std::cout<<"Final cost: "<<final_cost<<std::endl;
	}

	double Get_Final_Cost(Mesh& ref){
		//normalize by number of vertices and faces
		int ref_num_v = ref.V.size();
		int source_num_f = F.size();

		MatrixX SV(V.size(), 3), RV(ref.V.size(), 3);
		Eigen::MatrixXi SF(F.size(), 3);
		for (int i = 0; i < V.size(); ++i)
			SV.row(i) = V[i];
		for (int i = 0; i < ref.V.size(); ++i)
			RV.row(i) = ref.V[i];
		for (int i = 0; i < F.size(); ++i)
			SF.row(i) = F[i];

		VectorX sqrD;
		Eigen::VectorXi I;
		MatrixX C;
		igl::point_mesh_squared_distance(RV, SV, SF,sqrD,I,C);
		FT coverage_cost = sqrD.sum() * 0.5;
		std::cout<<"Coverage cost: "<<coverage_cost << std::endl;

		FT rigidity_cost = 0.0;
		std::cout<<"Rigidity cost: "<<rigidity_cost<<std::endl;

		FT final_cost = coverage_cost/ref_num_v + rigidity_cost/source_num_f;
		std::cout<<"Final cost: "<<final_cost<<std::endl;

		return final_cost;
	}

};

//For deformation
int main(int argc, char** argv) {
	if (argc < 4) {
		printf("./deform source.obj reference.obj output.obj textfile_name\n");
		return 0;
	}
	//Deform source to fit the reference

	Mesh src, ref;
	src.ReadOBJ_Manifold(argv[1]);
	ref.ReadOBJ_Manifold(argv[2]);

	//Get number of vertices and faces
	std::cout<<"Source:\t\t"<<"Num vertices: "<<src.V.size()<<"\tNum faces: "<<src.F.size()<<std::endl;
	std::cout<<"Reference:\t"<<"Num vertices: "<<ref.V.size()<<"\tNum faces: "<<ref.F.size()<<std::endl<<std::endl;

	UniformGrid grid(100);
	ref.Normalize();
	src.ApplyTransform(ref);

	ref.ConstructDistanceField(grid);

	src.Deform(grid);

	double cost;
	cost = src.Get_Final_Cost(ref);

	double threshold = 1e-3;

	//for deform_tune
	// double threshold = 1e-3;
	// double threshold = 5e-4; #does not filter valid deformations

	if (cost > threshold){
		std::cout<<"INVALID"<<std::endl;
		// return 0;
	}

	else {
		std::cout<<"Deformed"<<std::endl;
	}

	std::string text_file_name = argv[4];
	std::ofstream outfile;
	outfile.open(text_file_name, std::fstream::in | std::fstream::out | std::fstream::app);
	if (!outfile){
		outfile.open(text_file_name,  std::fstream::in | std::fstream::out | std::fstream::trunc);
	}
	//output_objectcost
	outfile<<argv[3]<<'\t'<<cost<<std::endl;
	outfile.close();

	std::ifstream is(argv[1]);
	std::ofstream os(argv[3]);
	char buffer[1024];
	int offset = 0;
	while (is.getline(buffer, 1024)) {
		if (buffer[0] == 'v' && buffer[1] == ' ') {
			auto v = src.V[offset++] * src.scale + src.pos;
			os << "v " << v[0] << " " << v[1] << " " << v[2] << "\n";
		} else {
			os << buffer << "\n";
		}
	}

	is.close();
	os.close();
	return 0;
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_LOG_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_LOG_HPP_INCLUDED

#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/config.hpp>

#include <boost/simd/function/any.hpp>
#include <boost/simd/function/bitwise_and.hpp>
#include <boost/simd/function/bitwise_cast.hpp>
#include <boost/simd/function/dec.hpp>
#include <boost/simd/function/fma.hpp>
#include <boost/simd/function/frexp.hpp>
#include <boost/simd/function/genmask.hpp>
#include <boost/simd/function/horn.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/if_minus.hpp>
#include <boost/simd/function/if_nan_else.hpp>
#include <boost/simd/function/is_ltz.hpp>
#include <boost/simd/function/is_nez.hpp>
#include <boost/simd/function/musl.hpp>
#include <boost/simd/function/plain.hpp>
#include <boost/simd/function/sqr.hpp>
#include <boost/simd/function/tofloat.hpp>

#include <boost/simd/constant/half.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/minf.hpp>
#include <boost/simd/constant/smallestposval.hpp>
#include <boost/simd/constant/sqrt_2o_2.hpp>
#include <boost/simd/constant/two.hpp>

#include <boost/simd/detail/constant/log_2hi.hpp>
#include <boost/simd/detail/constant/log_2lo.hpp>

#include <boost/simd/detail/dispatch/meta/as_integer.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;

  BOOST_DISPATCH_OVERLOAD_IF ( log_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::single_<A0>, X>
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 const & a0) const BOOST_NOEXCEPT
    {
      return musl_(log)(a0);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( log_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::double_<A0>, X>
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 const & a0) const BOOST_NOEXCEPT
    {
      return musl_(log)(a0);
    }
  };
  BOOST_DISPATCH_OVERLOAD_IF ( log_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::musl_tag
                          , bs::pack_< bd::single_<A0>, X>
                          )
  {
    BOOST_FORCEINLINE A0 operator() (const musl_tag &, const A0& a0) const BOOST_NOEXCEPT
    {
      using uiA0 = bd::as_integer_t<A0, unsigned>;
      using iA0 = bd::as_integer_t<A0,   signed>;
      A0 x =  a0;
      iA0 k(0);
      auto isnez = is_nez(a0);
#ifndef BOOST_SIMD_NO_DENORMALS
      auto test = is_less(a0, Smallestposval<A0>())&&isnez;
      if (any(test))
      {
        k = if_minus(test, k, iA0(23));
        x = if_else(test, x*A0(0x1p23f), x);
      }
#endif
      uiA0 ix = bitwise_cast<uiA0>(x);
      /* reduce x into [sqrt(2)/2, sqrt(2)] */
      ix += 0x3f800000 - 0x3f3504f3;
    //ix +=
      k += bitwise_cast<iA0>(ix>>23) - 0x7f;
      ix = (ix&0x007fffff) + 0x3f3504f3;
      x =  bitwise_cast<A0>(ix);
      A0 f = dec(x);
      A0 s = f/(Two<A0>() + f);
      A0 z = sqr(s);
      A0 w = sqr(z);
      A0 t1= w*horn<A0, 0x3eccce13, 0x3e789e26>(w);
      A0 t2= z*horn<A0, 0x3f2aaaaa, 0x3e91e9ee>(w);
      A0 R = t2 + t1;

      A0 hfsq = Half<A0>()*sqr(f);
      A0 dk = tofloat(k);
      A0 r = fma(dk, Log_2hi<A0>(), ((fma(s, (hfsq+R), dk*Log_2lo<A0>()) - hfsq) + f));
#ifndef BOOST_SIMD_NO_INFINITIES
      A0 zz = if_else(isnez, if_else(a0 == Inf<A0>(), Inf<A0>(), r), Minf<A0>());
#else
      A0 zz = if_else(isnez, r, Minf<A0>());
#endif
      return if_nan_else(is_ltz(a0), zz);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( log_
                             , (typename A0, typename X)
                             , (detail::is_native<X>)
                             , bd::cpu_
                             , bs::musl_tag
                             , bs::pack_< bd::double_<A0>, X>
                             )
  {
    BOOST_FORCEINLINE A0 operator() (const musl_tag &, const A0& a0) const BOOST_NOEXCEPT
    {
      using uiA0 = bd::as_integer_t<A0, unsigned>;
      using iA0 = bd::as_integer_t<A0,   signed>;
      A0 x = a0;
      uiA0 hx = bitwise_cast<uiA0>(x) >> 32;
      iA0 k(0);
      auto isnez = is_nez(a0);

#ifndef BOOST_SIMD_NO_DENORMALS
      auto test = is_less(a0, Smallestposval<A0>())&&isnez;
      if (any(test))
      {
        k = if_minus(test, k, iA0(54));
        x = if_else(test, x*A0(0x1p54), x);
      }
#endif
      /* reduce x into [sqrt(2)/2, sqrt(2)] */
      hx += 0x3ff00000 - 0x3fe6a09e;
      k += bitwise_cast<iA0>(hx>>20) - 0x3ff;
      A0 dk = tofloat(k);
      hx = (hx&0x000fffff) + 0x3fe6a09e;
      x = bitwise_cast<A0>(hx<<32 | (bitwise_and(0xffffffffull, bitwise_cast<uiA0>(x))));

      A0 f = dec(x);
      A0 hfsq = Half<A0>()*sqr(f);
      A0 s = f/(Two<A0>() + f);
      A0 z = sqr(s);
      A0 w = sqr(z);
      A0 t1= w*horn<A0,
                    0x3fd999999997fa04ll,
                    0x3fcc71c51d8e78afll,
                    0x3fc39a09d078c69fll
                    > (w);
      A0 t2= z*horn<A0,
                    0x3fe5555555555593ll,
                    0x3fd2492494229359ll,
                    0x3fc7466496cb03dell,
                    0x3fc2f112df3e5244ll
                    > (w);
      A0 R = t2+t1;
      A0 r = fma(dk, Log_2hi<A0>(), ((fma(s, (hfsq+R), dk*Log_2lo<A0>()) - hfsq) + f));
#ifndef BOOST_SIMD_NO_INFINITIES
      A0 zz = if_else(isnez, if_else(a0 == Inf<A0>(), Inf<A0>(), r), Minf<A0>());
#else
      A0 zz = if_else(isnez, r, Minf<A0>());
#endif
      return if_nan_else(is_ltz(a0), zz);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( log_
                          , (typename A0, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::plain_tag
                          , bs::pack_< bd::single_<A0>, X>
                          )
  {
    BOOST_FORCEINLINE A0 operator() (const plain_tag &, const A0& a0) const BOOST_NOEXCEPT
    {
      using iA0 = bd::as_integer_t<A0,   signed>;
      A0 x =  a0;
      iA0 k(0);
      auto isnez = is_nez(a0);
#ifndef BOOST_SIMD_NO_DENORMALS
      auto test = is_less(a0, Smallestposval<A0>())&&isnez;
      if (any(test))
      {
        k = if_minus(test, k, iA0(23));
        x = if_else(test, x*A0(0x1p23f), x);
      }
#endif
      iA0 kk;
      std::tie(x, kk) = fast_(frexp)(x);
      A0  x_lt_sqrthf = genmask(Sqrt_2o_2<A0>() > x);
      k += kk+bitwise_cast<iA0>(x_lt_sqrthf);
      A0 f = dec(x+bitwise_and(x, x_lt_sqrthf));
      A0 dk = tofloat(k);
      A0 s = f/(Two<A0>() + f);
      A0 z = sqr(s);
      A0 w = sqr(z);
      A0 t1= w*horn<A0, 0x3eccce13, 0x3e789e26>(w);
      A0 t2= z*horn<A0, 0x3f2aaaaa, 0x3e91e9ee>(w);
      A0 R = t2 + t1;

      A0 hfsq = Half<A0>()*sqr(f);
      A0 r = fma(dk, Log_2hi<A0>(), ((fma(s, (hfsq+R), dk*Log_2lo<A0>()) - hfsq) + f));
#ifndef BOOST_SIMD_NO_INFINITIES
      A0 zz = if_else(isnez, if_else(a0 == Inf<A0>(), Inf<A0>(), r), Minf<A0>());
#else
      A0 zz = if_else(isnez, r, Minf<A0>());
#endif
      return if_nan_else(is_ltz(a0), zz);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( log_
                             , (typename A0, typename X)
                             , (detail::is_native<X>)
                             , bd::cpu_
                             , bs::plain_tag
                             , bs::pack_< bd::double_<A0>, X>
                             )
  {
    BOOST_FORCEINLINE A0 operator() (const plain_tag &, const A0& a0) const BOOST_NOEXCEPT
    {
//      using uiA0 = bd::as_integer_t<A0, unsigned>;
      using iA0 = bd::as_integer_t<A0,   signed>;
      A0 x =  a0;
      iA0 k(0);
      auto isnez = is_nez(a0);
#ifndef BOOST_SIMD_NO_DENORMALS
      auto test = is_less(a0, Smallestposval<A0>())&&isnez;
      if (any(test))
      {
        k = if_minus(test, k, iA0(23));
        x = if_else(test, x*A0(0x1p23f), x);
      }
#endif
      iA0 kk;
      std::tie(x, kk) = fast_(frexp)(x);
      A0  x_lt_sqrthf = genmask(Sqrt_2o_2<A0>() >  x);
      k += kk+bitwise_cast<iA0>(x_lt_sqrthf);
      A0 f = dec(x+bitwise_and(x, x_lt_sqrthf));
      A0 dk = tofloat(k);
      // compute approximation
      A0 s = f/(Two<A0>()+f);
      A0 z = sqr(s);
      A0 w = sqr(z);
      A0 t1= w*horn<A0,
        0x3fd999999997fa04ll,
        0x3fcc71c51d8e78afll,
        0x3fc39a09d078c69fll
        > (w);
      A0 t2= z*horn<A0,
        0x3fe5555555555593ll,
        0x3fd2492494229359ll,
        0x3fc7466496cb03dell,
        0x3fc2f112df3e5244ll
        > (w);
      A0 R = t2+t1;
      A0 hfsq = Half<A0>()* sqr(f);
      A0 r = fma(dk, Log_2hi<A0>(), ((fma(s, (hfsq+R), dk*Log_2lo<A0>()) - hfsq) + f));
#ifndef BOOST_SIMD_NO_INFINITIES
      A0 zz = if_else(isnez, if_else(a0 == Inf<A0>(), Inf<A0>(), r), Minf<A0>());
#else
      A0 zz = if_else(isnez, r, Minf<A0>());
#endif
      return if_nan_else(is_ltz(a0), zz);
    }
  };

} } }

#endif
  /*
   *   1. Argument Reduction: find k and f such that
   *                      x = 2^k * (1+f),
   *         where  sqrt(2)/2 < 1+f < sqrt(2) .
   *
   *   2. Approximation of log(1+f).
   *      Let s = f/(2+f) ; based on log(1+f) = log(1+s) - log(1-s)
   *               = 2s + 2/3 s**3 + 2/5 s**5 + .....,
   *               = 2s + s*R
   *      We use a special Remez algorithm on [0,0.1716] to generate
   *      a polynomial of degree 14 to approximate R The maximum error
   *      of this polynomial approximation is bounded by 2**-58.45. In
   *      other words,
   *                      2      4      6      8      10      12      14
   *          R(z) ~ Lg1*s +Lg2*s +Lg3*s +Lg4*s +Lg5*s  +Lg6*s  +Lg7*s
   *      (the values of Lg1 to Lg7 are listed in the program)
   *      and
   *          |      2          14          |     -58.45
   *          | Lg1*s +...+Lg7*s    -  R(z) | <= 2
   *          |                             |
   *      Note that 2s = f - s*f = f - hfsq + s*hfsq, where hfsq = f*f/2.
   *      In order to guarantee error in log below 1ulp, we compute log
   *      by
   *              log(1+f) = f - s*(f - R)        (if f is not too large)
   *              log(1+f) = f - (hfsq - s*(hfsq+R)).     (better accuracy)
   *
   *      3. Finally,  log(x) = k*ln2 + log(1+f).
   *                          = k*ln2_hi+(f-(hfsq-(s*(hfsq+R)+k*ln2_lo)))
   *         Here ln2 is split into two floating point number:
   *                      ln2_hi + ln2_lo,
   *         where n*ln2_hi is always exact for |n| < 2000.
   */


#include <Rcpp.h>
#include <boost/multi_array.hpp>

using namespace Rcpp;

// [[Rcpp::plugins("cpp11")]]
// [[Rcpp::depends(BH)]]

//'Estep in girt.
//'
//'@param x DataFrame.
//'@param a0 discrinimation parameter vector.
//'@param b0 difficulty parameter vector.
//'@param Xq node of theta dist.
//'@param AX weight of theta dist.
//'@param Yr node of phi dist.
//'@param BY weighit of phi dist.
//'@param D factor constant.
//'@param group a vector.
//'@param ind a design matrix for group.
//'@param resp a design matrix for person.
//'@param MLL a vector
//'@export
// [[Rcpp::export]]

List Estep_girt_mg(DataFrame x,
                NumericVector a0,
                NumericVector b0,
                NumericVector Xq,
                NumericMatrix AX, // multigroup
                NumericVector Yr,
                NumericMatrix BY, // multigroup
                double D,
                IntegerVector group,
                IntegerMatrix ind, // design matrix
                IntegerMatrix resp, // design matrix
                NumericVector MLL
){

  const int nj = x.length();// item n
  const int nn = x.nrows(); // subject n
  const int nq = Xq.length(); // node of theta
  const int nr = Yr.length(); // node of phi
  const int ng = max(group); // group n

  boost::multi_array <double, 4> knqr (boost::extents[ng][nn][nq][nr]);
  boost::multi_array <double, 4> hqr (boost::extents[ng][nn][nq][nr]);
  boost::multi_array <double, 4> Njqr (boost::extents[ng][nj][nq][nr]);
  boost::multi_array <double, 4> rjqr (boost::extents[ng][nj][nq][nr]);
  //Dimension d(nn,nq,nr);
  //NumericVector knqr(d);

  // もとのデータフレームから項目反応パタンだけを抜き出し，行列として保存
  NumericMatrix xall(nn, nj);
  for (int j=0; j<nj; j++){
    NumericVector kk = x[j];
    for(int i=0; i<nn; i++){
      double k = kk[i];
      xall(i,j) = k;
    }
  }

  int u;
  double phi, theta, a, b, tt, t, A, e, p;
  for(int g=0; g<ng; g++){
    for(int r=0; r<nr; r++){
      phi = Yr[r];
      //Rprintf("%d % / 100%\r", 100/nr*(r+1));
      for(int q=0; q<nq; q++){
        theta = Xq[q];
        for(int i=0; i<nn; i++){
          if(group[i] != g+1) continue; // 集団に属さない受験者の部分はスキップ
          t = 1.0;
          for(int j=0; j<nj; j++){
            a = a0[j];
            b = b0[j];
            A = sqrt(1+phi*phi*a*a);
            e = exp(-D*a/A*(theta-b));
            p = 1/(1+e);
            u = xall(i,j);
            if(u == 1){ // correct
              tt = p;
            } else if (u == 0) { // incorrect
              tt = 1.0-p;
            } else { // NA
              tt = 1;
            }
            t = t*tt;
          }
          knqr[g][i][q][r] = t*AX[q]*BY[r];
        }
      }
    }
  }

  // 各受検者のtheta, phiごとに事後分布の重みを計算する。
  // 規格化
  double f = 0; // 周辺対数尤度代入用
  for(int g=0; g<ng; g++){
    for(int i=0; i<nn; i++){
      if(group[i] != g+1) continue; // 集団に属さない受験者の部分はスキップ
      double Z = 0;
      for(int q=0; q<nq; q++){ // 総和を1にするための分母の計算 // sum
        for(int r=0; r<nr; r++){
          Z += knqr[g][i][q][r];
        }
      }
      f += log(Z);
      for(int q=0; q<nq; q++){ // 総和を1にするための分母の計算 // sum
        for(int r=0; r<nr; r++){
          double l = knqr[g][i][q][r];
          hqr[g][i][q][r] =  l / Z;
        }
      }
    }
  }
  MLL.push_back(f);
  if(traits::is_nan<REALSXP>(f)){
    // 対数尤度の計算に失敗したら，計算を中止する。
    stop("Can't calculate marginal log likelihood.");
  }

  for(int g=0; g<ng; g++){
    for(int j=0; j<nj; j++){ // 各分点の期待度数
      if(ind(g,j) == 0) continue;
      for(int q=0; q<nq; q++){
        for(int r=0; r<nr; r++){
          double k = 0;
          for(int i=0; i<nn; i++){ // 欠測値がある場合，項目ごとに受検者数が異なる。
            if(resp(i,j)==0) continue;
            k += hqr[g][i][q][r];
          }
          Njqr[g][j][q][r]= k;
        }
      }
    }
  }

  for(int g=0; g<ng; g++){
    for(int j=0; j<nj; j++){ // 各分点の正答受検者の期待度数
      if(ind(g,j) == 0) continue;
      for(int q=0; q<nq; q++){
        for(int r=0; r<nr; r++){
          double h = 0;
          for(int i=0; i<nn; i++){ // sum
            //double d = resp(i,j);
            if(resp(i,j) == 0) continue;
            h += xall(i,j)*hqr[g][i][q][r];
          }
          rjqr[g][j][q][r] = h;
        }
      }
    }
  }

  return List::create(_["knqr"]=knqr, _["hqr"]=hqr, _["Njqr"]=Njqr, _["rjqr"]=rjqr, _["MLL"]=MLL);

}


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

#include <boost/program_options.hpp>
#include <cstddef>
#include <string>

#include "DataStructures/ComplexModalVector.hpp"
#include "DataStructures/DataBox/DataBox.hpp"
#include "DataStructures/DataVector.hpp"
#include "DataStructures/SpinWeighted.hpp"
#include "DataStructures/Variables.hpp"
#include "DataStructures/VariablesTag.hpp"
#include "Evolution/Systems/Cce/BoundaryData.hpp"
#include "Evolution/Systems/Cce/ReducedWorldtubeModeRecorder.hpp"
#include "Evolution/Systems/Cce/SpecBoundaryData.hpp"
#include "Evolution/Systems/Cce/Tags.hpp"
#include "Evolution/Systems/Cce/WorldtubeBufferUpdater.hpp"
#include "NumericalAlgorithms/Spectral/SwshCoefficients.hpp"
#include "NumericalAlgorithms/Spectral/SwshCollocation.hpp"
#include "Parallel/Printf.hpp"
#include "Utilities/Gsl.hpp"
#include "Utilities/TMPL.hpp"

// Charm looks for this function but since we build without a main function or
// main module we just have it be empty
extern "C" void CkRegisterMainModule(void) {}

// from a time-varies-fastest set of buffers provided by
// `MetricWorldtubeH5BufferUpdater` extract the set of coefficients for a
// particular time given by `buffer_time_offset` into the `time_span` size of
// buffer.
void slice_buffers_to_libsharp_modes(
    const gsl::not_null<Variables<Cce::cce_metric_input_tags>*>
        coefficients_set,
    const Variables<Cce::cce_metric_input_tags>& coefficients_buffers,
    const size_t time_span, const size_t buffer_time_offset, const size_t l_max,
    const size_t computation_l_max) {
  SpinWeighted<ComplexModalVector, 0> spin_weighted_buffer;

  for (const auto& libsharp_mode :
       Spectral::Swsh::cached_coefficients_metadata(computation_l_max)) {
    for (size_t i = 0; i < 3; ++i) {
      for (size_t j = i; j < 3; ++j) {
        tmpl::for_each<
            tmpl::list<Cce::Tags::detail::SpatialMetric,
                       Cce::Tags::detail::Dr<Cce::Tags::detail::SpatialMetric>,
                       Tags::dt<Cce::Tags::detail::SpatialMetric>>>(
            [&i, &j, &libsharp_mode, &spin_weighted_buffer,
             &coefficients_buffers, &coefficients_set, &l_max,
             &computation_l_max, &time_span, &buffer_time_offset](auto tag_v) {
              using tag = typename decltype(tag_v)::type;
              spin_weighted_buffer.set_data_ref(
                  get<tag>(*coefficients_set).get(i, j).data(),
                  Spectral::Swsh::size_of_libsharp_coefficient_vector(
                      computation_l_max));
              if (libsharp_mode.l > l_max) {
                Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                    libsharp_mode, make_not_null(&spin_weighted_buffer), 0, 0.0,
                    0.0);

              } else {
                Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                    libsharp_mode, make_not_null(&spin_weighted_buffer), 0,
                    get<tag>(coefficients_buffers)
                        .get(i, j)[time_span *
                                       Spectral::Swsh::goldberg_mode_index(
                                           l_max, libsharp_mode.l,
                                           static_cast<int>(libsharp_mode.m)) +
                                   buffer_time_offset],
                    get<tag>(coefficients_buffers)
                        .get(i, j)[time_span *
                                       Spectral::Swsh::goldberg_mode_index(
                                           l_max, libsharp_mode.l,
                                           -static_cast<int>(libsharp_mode.m)) +
                                   buffer_time_offset]);
              }
            });
      }
      tmpl::for_each<tmpl::list<Cce::Tags::detail::Shift,
                                Cce::Tags::detail::Dr<Cce::Tags::detail::Shift>,
                                Tags::dt<Cce::Tags::detail::Shift>>>(
          [&i, &libsharp_mode, &spin_weighted_buffer, &coefficients_buffers,
           &coefficients_set, &l_max, &computation_l_max, &time_span,
           &buffer_time_offset](auto tag_v) {
            using tag = typename decltype(tag_v)::type;
            spin_weighted_buffer.set_data_ref(
                get<tag>(*coefficients_set).get(i).data(),
                Spectral::Swsh::size_of_libsharp_coefficient_vector(
                    computation_l_max));

            if (libsharp_mode.l > l_max) {
              Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                  libsharp_mode, make_not_null(&spin_weighted_buffer), 0, 0.0,
                  0.0);

            } else {
              Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                  libsharp_mode, make_not_null(&spin_weighted_buffer), 0,
                  get<tag>(coefficients_buffers)
                      .get(i)[time_span *
                                  Spectral::Swsh::goldberg_mode_index(
                                      l_max, libsharp_mode.l,
                                      static_cast<int>(libsharp_mode.m)) +
                              buffer_time_offset],
                  get<tag>(coefficients_buffers)
                      .get(i)[time_span *
                                  Spectral::Swsh::goldberg_mode_index(
                                      l_max, libsharp_mode.l,
                                      -static_cast<int>(libsharp_mode.m)) +
                              buffer_time_offset]);
            }
          });
    }
    tmpl::for_each<tmpl::list<Cce::Tags::detail::Lapse,
                              Cce::Tags::detail::Dr<Cce::Tags::detail::Lapse>,
                              Tags::dt<Cce::Tags::detail::Lapse>>>(
        [&libsharp_mode, &spin_weighted_buffer, &coefficients_buffers,
         &coefficients_set, &l_max, &computation_l_max, &time_span,
         &buffer_time_offset](auto tag_v) {
          using tag = typename decltype(tag_v)::type;
          spin_weighted_buffer.set_data_ref(
              get(get<tag>(*coefficients_set)).data(),
              Spectral::Swsh::size_of_libsharp_coefficient_vector(
                  computation_l_max));

          if (libsharp_mode.l > l_max) {
            Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                libsharp_mode, make_not_null(&spin_weighted_buffer), 0, 0.0,
                0.0);

          } else {
            Spectral::Swsh::goldberg_modes_to_libsharp_modes_single_pair(
                libsharp_mode, make_not_null(&spin_weighted_buffer), 0,
                get(get<tag>(coefficients_buffers))
                    [time_span * Spectral::Swsh::goldberg_mode_index(
                                     l_max, libsharp_mode.l,
                                     static_cast<int>(libsharp_mode.m)) +
                     buffer_time_offset],
                get(get<tag>(coefficients_buffers))
                    [time_span * Spectral::Swsh::goldberg_mode_index(
                                     l_max, libsharp_mode.l,
                                     -static_cast<int>(libsharp_mode.m)) +
                     buffer_time_offset]);
          }
        });
  }
}

// read in the data from a (previously standard) SpEC worldtube file
// `input_file`, perform the boundary computation, and dump the (considerably
// smaller) dataset associated with the spin-weighted scalars to `output_file`.
void perform_cce_worldtube_reduction(
    const std::string& input_file, const std::string& output_file,
    const size_t buffer_depth, const size_t l_max_factor,
    const bool fix_spec_normalization = false) {
  Cce::MetricWorldtubeH5BufferUpdater buffer_updater{input_file};
  const size_t l_max = buffer_updater.get_l_max();
  // Perform the boundary computation to scalars at twice the input l_max to be
  // absolutely certain that there are no problems associated with aliasing.
  const size_t computation_l_max = l_max_factor * l_max;

  // we're not interpolating, this is just a reasonable number of rows to ingest
  // at a time.
  const size_t size_of_buffer = square(l_max + 1) * (buffer_depth);
  const DataVector& time_buffer = buffer_updater.get_time_buffer();

  Variables<Cce::cce_metric_input_tags> coefficients_buffers{size_of_buffer};
  Variables<Cce::cce_metric_input_tags> coefficients_set{
      Spectral::Swsh::size_of_libsharp_coefficient_vector(computation_l_max)};

  Variables<Cce::Tags::characteristic_worldtube_boundary_tags<
      Cce::Tags::BoundaryValue>>
      boundary_data_variables{
          Spectral::Swsh::number_of_swsh_collocation_points(computation_l_max)};

  using reduced_boundary_tags =
      tmpl::list<Cce::Tags::BoundaryValue<Cce::Tags::BondiBeta>,
                 Cce::Tags::BoundaryValue<Cce::Tags::BondiU>,
                 Cce::Tags::BoundaryValue<Cce::Tags::BondiQ>,
                 Cce::Tags::BoundaryValue<Cce::Tags::BondiW>,
                 Cce::Tags::BoundaryValue<Cce::Tags::BondiJ>,
                 Cce::Tags::BoundaryValue<Cce::Tags::Dr<Cce::Tags::BondiJ>>,
                 Cce::Tags::BoundaryValue<Cce::Tags::Du<Cce::Tags::BondiJ>>,
                 Cce::Tags::BoundaryValue<Cce::Tags::BondiR>,
                 Cce::Tags::BoundaryValue<Cce::Tags::Du<Cce::Tags::BondiR>>>;

  size_t time_span_start = 0;
  size_t time_span_end = 0;
  Cce::ReducedWorldtubeModeRecorder recorder{output_file};

  ComplexModalVector output_goldberg_mode_buffer{square(computation_l_max + 1)};
  ComplexModalVector output_libsharp_mode_buffer{
      Spectral::Swsh::size_of_libsharp_coefficient_vector(computation_l_max)};

  for (size_t i = 0; i < time_buffer.size(); ++i) {
    const double time = time_buffer[i];
    Parallel::printf("reducing data at time : %f / %f \r", time,
                     time_buffer[time_buffer.size() - 1]);
    buffer_updater.update_buffers_for_time(
        make_not_null(&coefficients_buffers), make_not_null(&time_span_start),
        make_not_null(&time_span_end), time, l_max, 0, buffer_depth);

    slice_buffers_to_libsharp_modes(
        make_not_null(&coefficients_set), coefficients_buffers,
        time_span_end - time_span_start, i - time_span_start, l_max,
        computation_l_max);

    if (not buffer_updater.has_version_history() and fix_spec_normalization) {
      Cce::create_bondi_boundary_data_from_unnormalized_spec_modes(
          make_not_null(&boundary_data_variables),
          get<Cce::Tags::detail::SpatialMetric>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::SpatialMetric>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::SpatialMetric>>(
              coefficients_set),
          get<Cce::Tags::detail::Shift>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::Shift>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::Shift>>(
              coefficients_set),
          get<Cce::Tags::detail::Lapse>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::Lapse>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::Lapse>>(
              coefficients_set),
          buffer_updater.get_extraction_radius(), computation_l_max);
    } else {
      Cce::create_bondi_boundary_data(
          make_not_null(&boundary_data_variables),
          get<Cce::Tags::detail::SpatialMetric>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::SpatialMetric>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::SpatialMetric>>(
              coefficients_set),
          get<Cce::Tags::detail::Shift>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::Shift>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::Shift>>(
              coefficients_set),
          get<Cce::Tags::detail::Lapse>(coefficients_set),
          get<Tags::dt<Cce::Tags::detail::Lapse>>(coefficients_set),
          get<Cce::Tags::detail::Dr<Cce::Tags::detail::Lapse>>(
              coefficients_set),
          buffer_updater.get_extraction_radius(), computation_l_max);
    }
    // loop over the tags that we want to dump.
    tmpl::for_each<reduced_boundary_tags>(
        [&recorder, &boundary_data_variables, &output_goldberg_mode_buffer,
         &output_libsharp_mode_buffer, &l_max, &computation_l_max,
         &time](auto tag_v) {
          using tag = typename decltype(tag_v)::type;
          SpinWeighted<ComplexModalVector, tag::type::type::spin>
              spin_weighted_libsharp_view;
          spin_weighted_libsharp_view.set_data_ref(
              output_libsharp_mode_buffer.data(),
              output_libsharp_mode_buffer.size());
          Spectral::Swsh::swsh_transform(
              computation_l_max, 1, make_not_null(&spin_weighted_libsharp_view),
              get(get<tag>(boundary_data_variables)));
          SpinWeighted<ComplexModalVector, tag::type::type::spin>
              spin_weighted_goldberg_view;
          spin_weighted_goldberg_view.set_data_ref(
              output_goldberg_mode_buffer.data(),
              output_goldberg_mode_buffer.size());
          Spectral::Swsh::libsharp_to_goldberg_modes(
              make_not_null(&spin_weighted_goldberg_view),
              spin_weighted_libsharp_view, computation_l_max);

          // The goldberg format type is in strictly increasing l modes, so to
          // reduce to a smaller l_max, we can just take the first (l_max + 1)^2
          // values.
          ComplexModalVector reduced_goldberg_view{
              output_goldberg_mode_buffer.data(), square(l_max + 1)};
          recorder.append_worldtube_mode_data(
              "/" + Cce::dataset_label_for_tag<tag>(), time,
              reduced_goldberg_view, l_max, tag::type::type::spin == 0);
        });
  }
  Parallel::printf("\n");
}

/*
 * This executable is used for converting the unnecessarily large SpEC worldtube
 * data format into a far smaller representation (roughly a factor of 4) just
 * storing the worldtube scalars that are required as input for CCE.
 */
int main(int argc, char** argv) {
  boost::program_options::positional_options_description pos_desc;
  pos_desc.add("old_spec_cce_file", 1).add("output_file", 1);

  boost::program_options::options_description desc("Options");
  desc.add_options()("help,h,", "show this help message")(
      "input_file", boost::program_options::value<std::string>()->required(),
      "name of old CCE data file")(
      "output_file", boost::program_options::value<std::string>()->required(),
      "output filename")(
      "fix_spec_normalization",
      "Apply corrections associated with documented SpEC "
      "worldtube file errors")(
      "buffer_depth",
      boost::program_options::value<size_t>()->default_value(2000),
      "number of time steps to load during each call to the file-accessing "
      "routines. Higher values mean fewer, larger loads from file into RAM.")(
      "lmax_factor", boost::program_options::value<size_t>()->default_value(2),
      "the boundary computations will be performed at a resolution that is "
      "lmax_factor times the input file lmax to avoid aliasing");

  boost::program_options::variables_map vars;

  boost::program_options::store(
      boost::program_options::command_line_parser(argc, argv)
          .positional(pos_desc)
          .options(desc)
          .run(),
      vars);

  if (vars.count("help") != 0u or vars.count("input_file") == 0u or
      vars.count("output_file") == 0u) {
    Parallel::printf("%s\n", desc);
    return 0;
  }

  perform_cce_worldtube_reduction(vars["input_file"].as<std::string>(),
                                  vars["output_file"].as<std::string>(),
                                  vars["buffer_depth"].as<size_t>(),
                                  vars["lmax_factor"].as<size_t>(),
                                  vars.count("fix_spec_normalization") != 0u);
}


#include "fbstab/components/mpc_data.h"

#include <Eigen/Dense>
#include <cmath>
#include <iostream>
#include <stdexcept>

#include "tools/matrix_sequence.h"

namespace fbstab {

using MatrixXd = Eigen::MatrixXd;
using VectorXd = Eigen::VectorXd;
using Map = Eigen::Map<Eigen::MatrixXd>;
using ConstMap = Eigen::Map<const Eigen::MatrixXd>;

void MpcData::gemvH(const Eigen::VectorXd& x, double a, double b,
                    Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::gemvH: y input is null.");
  }
  if (x.size() != nz_ || y->size() != nz_) {
    throw std::runtime_error("Size mismatch in MpcData::gemvH.");
  }
  if (b == 0.0) {
    y->fill(0.0);
  } else if (b != 1.0) {
    (*y) *= b;
  }

  // Create reshaped views of input and output vectors.
  Map w(y->data(), nx_ + nu_,
        N_ + 1);  // w = reshape(y, [nx + nu, N + 1]);
  ConstMap v(x.data(), nx_ + nu_,
             N_ + 1);  // v = reshape(x, [nx + nu, N + 1]);
  for (int i = 0; i < N_ + 1; i++) {
    const auto& Q = Q_(i);
    const auto& S = S_(i);
    const auto& R = R_(i);

    // These variables alias w.
    auto yx = w.block(0, i, nx_, 1);
    auto yu = w.block(nx_, i, nu_, 1);

    // These variables alias v.
    const auto vx = v.block(0, i, nx_, 1);
    const auto vu = v.block(nx_, i, nu_, 1);

    // [yx] += a * [Q(i) S(i)'] [vx]
    // [yu]        [S(i) R(i) ] [vu]
    // Using lazyProduct is inefficient so should be avoided when possible.
    if (a == 1.0) {
      yx.noalias() += Q * vx + S.transpose() * vu;
      yu.noalias() += S * vx + R * vu;
    } else if (a == -1.0) {
      yx.noalias() -= Q * vx + S.transpose() * vu;
      yu.noalias() -= S * vx + R * vu;
    } else {
      yx += a * Q.lazyProduct(vx);
      yx += a * S.transpose().lazyProduct(vu);
      yu += a * S.lazyProduct(vx);
      yu += a * R.lazyProduct(vu);
    }
  }
}

void MpcData::gemvA(const Eigen::VectorXd& x, double a, double b,
                    Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::gemvA: y input is null.");
  }
  if (x.size() != nz_ || y->size() != nv_) {
    throw std::runtime_error("Size mismatch in MpcData::gemvA.");
  }
  if (b == 0.0) {
    y->fill(0.0);
  } else if (b != 1.0) {
    (*y) *= b;
  }
  // Create reshaped views of input and output vectors.
  ConstMap z(x.data(), nx_ + nu_, N_ + 1);
  Map w(y->data(), nc_, N_ + 1);

  for (int i = 0; i < N_ + 1; i++) {
    const auto& E = E_(i);
    const auto& L = L_(i);

    // This aliases w.
    auto yi = w.col(i);

    // These alias z.
    const auto xi = z.block(0, i, nx_, 1);
    const auto ui = z.block(nx_, i, nu_, 1);

    // yi += a*(E*vx + L*vu)
    if (a == 1.0) {
      yi.noalias() += E * xi + L * ui;
    } else if (a == -1.0) {
      yi.noalias() -= E * xi + L * ui;
    } else {
      yi += a * E.lazyProduct(xi);
      yi += a * L.lazyProduct(ui);
    }
  }
}

void MpcData::gemvG(const Eigen::VectorXd& x, double a, double b,
                    Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::gemvG: y input is null.");
  }
  if (x.size() != nz_ || y->size() != nl_) {
    throw std::runtime_error("Size mismatch in MpcData::gemvG.");
  }
  if (b == 0.0) {
    y->fill(0.0);
  } else if (b != 1.0) {
    (*y) *= b;
  }
  // Create reshaped views of input and output vectors.
  ConstMap z(x.data(), nx_ + nu_, N_ + 1);
  Map w(y->data(), nx_, N_ + 1);

  w.col(0).noalias() += -a * z.block(0, 0, nx_, 1);

  for (int i = 1; i < N_ + 1; i++) {
    const auto& A = A_(i - 1);
    const auto& B = B_(i - 1);

    // Alias for the output at stage i.
    auto yi = w.col(i);
    // Aliases for the state and control at stage i - 1.
    const auto xm1 = z.block(0, i - 1, nx_, 1);
    const auto um1 = z.block(nx_, i - 1, nu_, 1);
    // Alias for the state at stage i.
    const auto xi = z.block(0, i, nx_, 1);

    // y(i) += a*(A(i-1)*x(i-1) + B(i-1)u(i-1) - x(i))
    if (a == 1.0) {
      yi.noalias() += A * xm1 + B * um1;
      yi.noalias() -= xi;
    } else if (a == -1.0) {
      yi.noalias() -= A * xm1 + B * um1;
      yi.noalias() += xi;
    } else {
      yi += a * A.lazyProduct(xm1);
      yi += a * B.lazyProduct(um1);
      yi -= a * xi;
    }
  }
}

void MpcData::gemvGT(const Eigen::VectorXd& x, double a, double b,
                     Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::gemvGT: y input is null.");
  }
  if (x.size() != nl_ || y->size() != nz_) {
    throw std::runtime_error("Size mismatch in MpcData::gemvGT.");
  }
  if (b == 0.0) {
    y->fill(0.0);
  } else if (b != 1.0) {
    (*y) *= b;
  }

  // Create reshaped views of input and output vectors.
  ConstMap v(x.data(), nx_, N_ + 1);
  Map w(y->data(), nx_ + nu_, N_ + 1);

  for (int i = 0; i < N_; i++) {
    const auto& A = A_(i);
    const auto& B = B_(i);

    // Aliases for the dual variables at stage i and i+1;
    const auto vi = v.col(i);
    const auto vp1 = v.col(i + 1);

    // Aliases for the state and control at stage i.
    auto xi = w.block(0, i, nx_, 1);
    auto ui = w.block(nx_, i, nu_, 1);

    // x(i) += a*(-v(i) + A(i)' * v(i+1))
    // u(i) += a*B(i)' * v(i+1)
    xi.noalias() += -a * vi;
    if (a == 1.0) {
      xi.noalias() += A.transpose() * vp1;
      ui.noalias() += B.transpose() * vp1;
    } else if (a == -1.0) {
      xi.noalias() -= A.transpose() * vp1;
      ui.noalias() -= B.transpose() * vp1;
    } else {
      xi.noalias() += a * A.transpose().lazyProduct(vp1);
    }
  }
  // The i = N step of the recursion.
  w.block(0, N_, nx_, 1).noalias() += -a * v.col(N_);
}

void MpcData::gemvAT(const Eigen::VectorXd& x, double a, double b,
                     Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::gemvAT: y input is null.");
  }
  if (x.size() != nv_ || y->size() != nz_) {
    throw std::runtime_error("Size mismatch in MpcData::gemvAT.");
  }
  if (b == 0.0) {
    y->fill(0.0);
  } else if (b != 1.0) {
    (*y) *= b;
  }
  // Create reshaped views of input and output vectors.
  ConstMap v(x.data(), nc_, N_ + 1);
  Map w(y->data(), nx_ + nu_, N_ + 1);

  for (int i = 0; i < N_ + 1; i++) {
    const auto& E = E_(i);
    const auto& L = L_(i);

    auto xi = w.block(0, i, nx_, 1);
    auto ui = w.block(nx_, i, nu_, 1);

    const auto vi = v.col(i);
    // x(i) += a*E(i)' * v(i)
    // u(i) += a*L(i)' * v(i)
    if (a == 1.0) {
      xi.noalias() += E.transpose() * vi;
      ui.noalias() += L.transpose() * vi;
    } else if (a == -1.0) {
      xi.noalias() -= E.transpose() * vi;
      ui.noalias() -= L.transpose() * vi;
    } else {
      ui.noalias() += a * L.transpose().lazyProduct(vi);
      xi.noalias() += a * E.transpose().lazyProduct(vi);
    }
  }
}

void MpcData::axpyf(double a, Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::axpyf: y input is null.");
  }
  if (y->size() != nz_) {
    throw std::runtime_error("Size mismatch in MpcData::axpyf.");
  }

  // Create reshaped view of the input vector.
  Map w(y->data(), nx_ + nu_, N_ + 1);

  for (int i = 0; i < N_ + 1; i++) {
    auto xi = w.block(0, i, nx_, 1);
    auto ui = w.block(nx_, i, nu_, 1);

    xi.noalias() += a * q_(i);
    ui.noalias() += a * r_(i);
  }
}

void MpcData::axpyh(double a, Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::axpyh: y input is null.");
  }
  if (y->size() != nl_) {
    throw std::runtime_error("Size mismatch in MpcData::axpyh.");
  }
  // Create reshaped view of the input vector.
  Map w(y->data(), nx_, N_ + 1);
  w.col(0) += -a * x0_;

  for (int i = 1; i < N_ + 1; i++) {
    w.col(i) += -a * c_(i - 1);
  }
}

void MpcData::axpyb(double a, Eigen::VectorXd* y) const {
  if (y == nullptr) {
    throw std::runtime_error("In MpcData::axpyb: y input is null.");
  }
  if (y->size() != nv_) {
    throw std::runtime_error("Size mismatch in MpcData::axpyb.");
  }
  // Create reshaped view of the input vector.
  Map w(y->data(), nc_, N_ + 1);

  for (int i = 0; i < N_ + 1; i++) {
    w.col(i).noalias() += -a * d_(i);
  }
}

void MpcData::ValidateInputs() const {
  bool OK = true;
  const int N = Q_.length();
  if (N <= 0) {
    throw std::runtime_error("Horizon length must be at least 1.");
  }

  OK = OK && N == R_.length();
  OK = OK && N == S_.length();
  OK = OK && N == q_.length();
  OK = OK && N == r_.length();
  OK = OK && (N - 1) == A_.length();
  OK = OK && (N - 1) == B_.length();
  OK = OK && (N - 1) == c_.length();
  OK = OK && N == E_.length();
  OK = OK && N == L_.length();
  OK = OK && N == d_.length();
  if (!OK) {
    throw std::runtime_error(
        "Sequence length mismatch in input data to MpcData.");
  }

  const int nx = Q_.rows();
  if (x0_.size() != nx) {
    throw std::runtime_error("Size mismatch in x0 input to MpcData.");
  }
  if (Q_.cols() != nx) {
    throw std::runtime_error("Size mismatch in Q input to MpcData.");
  }
  if (S_.cols() != nx) {
    throw std::runtime_error("Size mismatch in S input to MpcData.");
  }
  if (q_.rows() != nx) {
    throw std::runtime_error("Size mismatch in q input to MpcData.");
  }
  if (E_.cols() != nx) {
    throw std::runtime_error("Size mismatch in E input to MpcData.");
  }
  if (A_.rows() != nx || A_.cols() != nx) {
    throw std::runtime_error("Size mismatch in A input to MpcData.");
  }
  if (B_.rows() != nx) {
    throw std::runtime_error("Size mismatch in B input to MpcData.");
  }
  if (c_.rows() != nx) {
    throw std::runtime_error("Size mismatch in c input to MpcData.");
  }

  const int nu = R_.rows();
  if (R_.cols() != nu) {
    throw std::runtime_error("Size mismatch in R input to MpcData.");
  }
  if (S_.rows() != nu) {
    throw std::runtime_error("Size mismatch in S input to MpcData.");
  }
  if (r_.rows() != nu) {
    throw std::runtime_error("Size mismatch in r input to MpcData.");
  }
  if (L_.cols() != nu) {
    throw std::runtime_error("Size mismatch in L input to MpcData.");
  }
  if (B_.cols() != nu) {
    throw std::runtime_error("Size mismatch in B input to MpcData.");
  }

  const int nc = E_.rows();
  if (L_.rows() != nc) {
    throw std::runtime_error("Size mismatch in L input to MpcData.");
  }
  if (d_.rows() != nc) {
    throw std::runtime_error("Size mismatch in d input to MpcData.");
  }
}

}  // namespace fbstab