/*
@copyright Louis Dionne 2015
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_HANA_EXAMPLE_CPPCON_2014_MATRIX_DET_HPP
#define BOOST_HANA_EXAMPLE_CPPCON_2014_MATRIX_DET_HPP

#include <boost/hana/constant.hpp>
#include <boost/hana/foldable.hpp>
#include <boost/hana/functional/always.hpp>
#include <boost/hana/functional/compose.hpp>
#include <boost/hana/functional/fix.hpp>
#include <boost/hana/functional/on.hpp>
#include <boost/hana/functional/partial.hpp>
#include <boost/hana/functional/placeholder.hpp>
#include <boost/hana/functor.hpp>
#include <boost/hana/integral_constant.hpp>
#include <boost/hana/iterable.hpp>
#include <boost/hana/sequence.hpp>
#include <boost/hana/logical.hpp>
#include <boost/hana/monoid.hpp>
#include <boost/hana/range.hpp>
#include <boost/hana/ring.hpp>
#include <boost/hana/tuple.hpp>

#include <utility>

#include "matrix.hpp"

namespace cppcon {
    namespace detail {
        auto remove_at = [](auto n, auto xs) {
            using namespace boost::hana;
            using L = datatype_t<decltype(xs)>;
            auto with_indices = zip(xs, to<L>(range(size_t<0>, length(xs))));
            auto removed = filter(with_indices, compose(n != _, last));
            return transform(removed, head);
        };
    }

    auto det = boost::hana::fix([](auto det, auto&& m) -> decltype(auto) {
        using namespace boost::hana;
        auto matrix_minor = [=](auto&& m, auto i, auto j) -> decltype(auto) {
            return det(unpack(
                transform(
                    detail::remove_at(i, rows(std::forward<decltype(m)>(m))),
                    partial(detail::remove_at, j)
                ),
                matrix
            ));
        };

        auto cofactor = [=](auto&& m, auto i, auto j) {
            return power(int_<-1>, plus(i, j)) *
                    matrix_minor(std::forward<decltype(m)>(m), i, j);
        };

        return eval_if(m.size() == size_t<1>,
            always(m.at(size_t<0>, size_t<0>)),
            [=](auto _) {
                auto cofactors_1st_row = unpack(_(range)(size_t<0>, m.ncolumns()),
                    on(make<Tuple>, partial(cofactor, m, size_t<0>))
                );
                return detail::tuple_scalar_product(head(rows(m)), cofactors_1st_row);
            }
        );
    });
} // end namespace cppcon

#endif // !BOOST_HANA_EXAMPLE_CPPCON_2014_MATRIX_DET_HPP


#include <gmp.h>
#include <fc/crypto/alt_bn128.hpp>

#include <libff/algebra/curves/alt_bn128/alt_bn128_g1.hpp>
#include <libff/algebra/curves/alt_bn128/alt_bn128_g2.hpp>
#include <libff/algebra/curves/alt_bn128/alt_bn128_pairing.hpp>
#include <libff/algebra/curves/alt_bn128/alt_bn128_pp.hpp>

#include <libff/common/profiling.hpp>
#include <boost/throw_exception.hpp>
#include <algorithm>

namespace fc {

    using Scalar = libff::bigint<libff::alt_bn128_q_limbs>;

    void initLibSnark() noexcept {
        static bool s_initialized = []() noexcept {
            libff::inhibit_profiling_info = true;
            libff::inhibit_profiling_counters = true;
            libff::alt_bn128_pp::init_public_params();
            return true;
        }();
        (void)s_initialized;
    }

    Scalar to_scalar(const bytes& be) noexcept {
        mpz_t m;
        mpz_init(m);
        mpz_import(m, be.size(), /*order=*/1, /*size=*/1, /*endian=*/0, /*nails=*/0, &be[0]);
        Scalar out{m};
        mpz_clear(m);
        return out;
    }

    // Notation warning: Yellow Paper's p is the same libff's q.
    // Returns x < p (YP notation).
    static bool valid_element_of_fp(const Scalar& x) noexcept {
        return mpn_cmp(x.data, libff::alt_bn128_modulus_q.data, libff::alt_bn128_q_limbs) < 0;
    }

    std::variant<alt_bn128_error, libff::alt_bn128_G1> decode_g1_element(const bytes& bytes64_be) noexcept {
        if(bytes64_be.size() != 64) {
            return alt_bn128_error::input_len_error;
        }

        bytes sub1(bytes64_be.begin(), bytes64_be.begin()+32);
        bytes sub2(bytes64_be.begin()+32, bytes64_be.begin()+64);

        Scalar x{to_scalar(sub1)};
        Scalar y{to_scalar(sub2)};

        if (!valid_element_of_fp(x) || !valid_element_of_fp(y)) {
            return alt_bn128_error::operand_component_invalid;
        }

        if (x.is_zero() && y.is_zero()) {
            return alt_bn128_error::operand_at_origin;
        }

        libff::alt_bn128_G1 point{x, y, libff::alt_bn128_Fq::one()};
        if (!point.is_well_formed()) {
            return alt_bn128_error::operand_not_in_curve;
        }
        return point;
    }

    std::variant<alt_bn128_error, libff::alt_bn128_Fq2> decode_fp2_element(const bytes& bytes64_be) noexcept {
        if(bytes64_be.size() != 64) {
            return alt_bn128_error::input_len_error;
        }

        // big-endian encoding
        bytes sub1(bytes64_be.begin()+32, bytes64_be.begin()+64);
        bytes sub2(bytes64_be.begin(), bytes64_be.begin()+32);

        Scalar c0{to_scalar(sub1)};
        Scalar c1{to_scalar(sub2)};

        if (!valid_element_of_fp(c0) || !valid_element_of_fp(c1)) {
            return alt_bn128_error::operand_component_invalid;
        }

        return libff::alt_bn128_Fq2{c0, c1};
    }

    std::variant<alt_bn128_error, libff::alt_bn128_G2> decode_g2_element(const bytes& bytes128_be) noexcept {

        bytes sub1(bytes128_be.begin(), bytes128_be.begin()+64);
        auto maybe_x = decode_fp2_element(sub1);
        if (std::holds_alternative<alt_bn128_error>(maybe_x)) {
            return std::get<alt_bn128_error>(maybe_x);
        }

        bytes sub2(bytes128_be.begin()+64, bytes128_be.begin()+128);
        auto maybe_y = decode_fp2_element(sub2);
        if (std::holds_alternative<alt_bn128_error>(maybe_y)) {
            return std::get<alt_bn128_error>(maybe_y);
        }

        const auto& x = std::get<libff::alt_bn128_Fq2>(maybe_x);
        const auto& y = std::get<libff::alt_bn128_Fq2>(maybe_y);

        if (x.is_zero() && y.is_zero()) {
            return alt_bn128_error::operand_at_origin;
        }

        libff::alt_bn128_G2 point{x, y, libff::alt_bn128_Fq2::one()};
        if (!point.is_well_formed()) {
            return alt_bn128_error::operand_not_in_curve;
        }

        if (!(libff::alt_bn128_G2::order() * point).is_zero()) {
            // wrong order, doesn't belong to the subgroup G2
            return alt_bn128_error::operand_outside_g2;
        }

        return point;
    }

    bytes encode_g1_element(libff::alt_bn128_G1 p) noexcept {
        bytes out(64, '\0');
        if (p.is_zero()) {
            return out;
        }

        p.to_affine_coordinates();

        auto x{p.X.as_bigint()};
        auto y{p.Y.as_bigint()};

        std::memcpy(&out[0], y.data, 32);
        std::memcpy(&out[32], x.data, 32);

        std::reverse(out.begin(), out.end());
        return out;
    }

    std::variant<alt_bn128_error, bytes> alt_bn128_add(const bytes& op1, const bytes& op2) {
        fc::initLibSnark();

        auto maybe_x = decode_g1_element(op1);
        if (std::holds_alternative<alt_bn128_error>(maybe_x)) {
            return std::get<alt_bn128_error>(maybe_x);
        }

        auto maybe_y = decode_g1_element(op2);
        if (std::holds_alternative<alt_bn128_error>(maybe_y)) {
            return std::get<alt_bn128_error>(maybe_y);
        }

        const auto& x = std::get<libff::alt_bn128_G1>(maybe_x);
        const auto& y = std::get<libff::alt_bn128_G1>(maybe_y);

        libff::alt_bn128_G1 g1Sum = x + y;
        return encode_g1_element(g1Sum);
    }

    std::variant<alt_bn128_error, bytes> alt_bn128_mul(const bytes& g1_point, const bytes& scalar) {
        initLibSnark();

        auto maybe_x = decode_g1_element(g1_point);
        if (std::holds_alternative<alt_bn128_error>(maybe_x)) {
            return std::get<alt_bn128_error>(maybe_x);
        }

        auto& x = std::get<libff::alt_bn128_G1>(maybe_x);

        if(scalar.size() != 32) {
            return alt_bn128_error::invalid_scalar_size;
        }

        Scalar n{to_scalar(scalar)};

        libff::alt_bn128_G1 g1Product = n * x;
        return encode_g1_element(g1Product);
    }

    static constexpr size_t kSnarkvStride{192};

    std::variant<alt_bn128_error, bool>  alt_bn128_pair(const bytes& g1_g2_pairs, const yield_function_t& yield) {
        if (g1_g2_pairs.size() % kSnarkvStride != 0) {
            return alt_bn128_error::pairing_list_size_error;
        }

        size_t k{g1_g2_pairs.size() / kSnarkvStride};

        initLibSnark();
        using namespace libff;

        static const auto one{alt_bn128_Fq12::one()};
        auto accumulator{one};

        for (size_t i{0}; i < k; ++i) {
            auto offset = i * kSnarkvStride;

            bytes sub1(g1_g2_pairs.begin()+offset, g1_g2_pairs.begin()+offset+64);
            auto maybe_a = decode_g1_element(sub1);
            if (std::holds_alternative<alt_bn128_error>(maybe_a)) {
                return std::get<alt_bn128_error>(maybe_a);
            }

            bytes sub2(g1_g2_pairs.begin()+offset+64, g1_g2_pairs.begin()+offset+64+128);
            auto maybe_b = decode_g2_element(sub2);
            if (std::holds_alternative<alt_bn128_error>(maybe_b)) {
                return std::get<alt_bn128_error>(maybe_b);
            }

            const auto& a = std::get<libff::alt_bn128_G1>(maybe_a);
            const auto& b = std::get<libff::alt_bn128_G2>(maybe_b);

            if (a.is_zero() || b.is_zero()) {
                continue;
            }

            accumulator = accumulator * alt_bn128_miller_loop(alt_bn128_precompute_G1(a), alt_bn128_precompute_G2(b));
            yield();
        }

        bool pair_result = false;
        if (alt_bn128_final_exponentiation(accumulator) == one) {
            pair_result = true;
        }

        return pair_result;
    }
}


#include<fastenvelope/obb.h>
#include <Eigen/Eigenvalues>
#include<fstream>
#include <iomanip>
#include<fastenvelope/FastEnvelope.h>// for is_triangle_degenerated()
namespace fastEnvelope {

	std::vector<obb> obb::build_obb_matrixs(const std::vector<Vector3>& face_vertices, const int p_face[8][3], const int c_face[6][3],
		const std::array<std::vector<int>, 8>& p_facepoint, const std::array<std::array<int, 4>, 6>& c_facepoint) {
		std::vector<Vector3> points;

		int dege;
		Vector3 normal;
		std::vector<obb> M;
		M.reserve(8);

		if (face_vertices.size() == polyhedron_point_number1) {
			//M.resize(polyhedron_face_number1);

			for (int j = 0; j < polyhedron_face_number1; j++) {
				//triangle to get normal: {envelope_vertices[i][p_face[j][0]],envelope_vertices[i][p_face[j][1]],envelope_vertices[i][p_face[j][2]]}
#ifdef DO_NOT_HAVE_DEGENERATED_FACES
				if (((face_vertices[p_face[j][0]] - face_vertices[p_face[j][1]]).cross(
					face_vertices[p_face[j][0]] - face_vertices[p_face[j][2]])).norm() < SCALAR_ZERO)
					normal = FastEnvelope::accurate_normal_vector(face_vertices[p_face[j][0]], face_vertices[p_face[j][1]],
						face_vertices[p_face[j][0]], face_vertices[p_face[j][2]]);
				else {
					normal = ((face_vertices[p_face[j][0]] - face_vertices[p_face[j][1]]).cross(face_vertices[p_face[j][0]] - face_vertices[p_face[j][2]])).normalized();
				}
#else
					//triangle to get normal: {envelope_vertices[i][p_face[j][0]],envelope_vertices[i][p_face[j][1]],envelope_vertices[i][p_face[j][2]]}

				dege = FastEnvelope::is_triangle_degenerated(face_vertices[p_face[j][0]], face_vertices[p_face[j][1]], face_vertices[p_face[j][2]]);
				if (dege == FastEnvelope::DEGENERATED_SEGMENT || dege == FastEnvelope::DEGENERATED_POINT) {
					std::cout << "need to fix here, face degeneration" << std::endl;
					exit(0);
				}

				if (dege == FastEnvelope::NERLY_DEGENERATED) {
					normal = FastEnvelope::accurate_normal_vector(face_vertices[p_face[j][0]], face_vertices[p_face[j][1]],
						face_vertices[p_face[j][0]], face_vertices[p_face[j][2]]);
				}
				if (dege == FastEnvelope::NOT_DEGENERATED) {
					normal = ((face_vertices[p_face[j][0]] - face_vertices[p_face[j][1]]).cross(face_vertices[p_face[j][0]] - face_vertices[p_face[j][2]])).normalized();
				}
#endif
				points.clear();
				points.resize(p_facepoint[j].size());
				for (int k = 0; k < p_facepoint[j].size(); k++) {
					points[k] = face_vertices[p_facepoint[j][k]];
				}

				M.emplace_back();
				build_obb(points, normal, OBB_OFFSET, M.back().Trans, M.back().invTrans);
			}

		}

		if (face_vertices.size() == polyhedron_point_number2) {
			M.resize(polyhedron_face_number2);

			for (int j = 0; j < polyhedron_face_number2; j++) {
#ifdef DO_NOT_HAVE_DEGENERATED_FACES
				if (((face_vertices[c_face[j][0]] - face_vertices[c_face[j][1]]).cross(
					face_vertices[c_face[j][0]] - face_vertices[c_face[j][2]])).norm() < SCALAR_ZERO)
					normal = FastEnvelope::accurate_normal_vector(face_vertices[c_face[j][0]], face_vertices[c_face[j][1]],
						face_vertices[c_face[j][0]], face_vertices[c_face[j][2]]);
				else {
					normal = (face_vertices[c_face[j][0]] - face_vertices[c_face[j][1]]).cross(face_vertices[c_face[j][0]] - face_vertices[c_face[j][2]]).normalized();
				}
#else

				//triangle to get normal: {envelope_vertices[i][p_face[j][0]],envelope_vertices[i][p_face[j][1]],envelope_vertices[i][p_face[j][2]]}

				dege = FastEnvelope::is_triangle_degenerated(face_vertices[c_face[j][0]], face_vertices[c_face[j][1]], face_vertices[c_face[j][2]]);
				if (dege == FastEnvelope::DEGENERATED_SEGMENT || dege == FastEnvelope::DEGENERATED_POINT) {
					std::cout << "need to fix here, face degeneration" << std::endl;
					exit(0);
				}

				if (dege == FastEnvelope::NERLY_DEGENERATED) {
					normal = FastEnvelope::accurate_normal_vector(face_vertices[c_face[j][0]], face_vertices[c_face[j][1]],
						face_vertices[c_face[j][0]], face_vertices[c_face[j][2]]);
				}
				if (dege == FastEnvelope::NOT_DEGENERATED) {
					normal = (face_vertices[c_face[j][0]] - face_vertices[c_face[j][1]]).cross(face_vertices[c_face[j][0]] - face_vertices[c_face[j][2]]).normalized();
				}

#endif
				points.clear();
				points.resize(c_facepoint[j].size());
				for (int k = 0; k < c_facepoint[j].size(); k++) {
					points[k] = face_vertices[c_facepoint[j][k]];
				}

				M.emplace_back();
				build_obb(points, normal, OBB_OFFSET, M.back().Trans, M.back().invTrans);
			}
		}
		return M;

	}

	bool obb::intersects(const obb& M2) const {
		return obb_intersection(Trans, invTrans, M2.Trans, M2.invTrans);
	}

	bool obb::intersects(const obb& M2, const obb& M3) const {

		//bool flag1=obb_intersection(Trans, invTrans, M2.Trans, M2.invTrans);
		bool flag1 = intersects(M2);
		if (flag1 == false) return false;

		//bool flag2 = obb_intersection(Trans, invTrans, M3.Trans, M3.invTrans);
		bool flag2 = intersects(M3);
		if (flag2 == false) return false;

		//bool flag3 = obb_intersection(M2.Trans, M2.invTrans, M3.Trans, M3.invTrans);
		bool flag3 = M2.intersects(M3);
		if (flag3 == false) return false;

		return true;
	}

	obb obb::build_triangle_obb_matrixs(const Vector3&t0, const Vector3&t1, const Vector3&t2) {
		Vector3 normal;
		obb res;
		if (((t0 - t1).cross(t0 - t2)).norm() < SCALAR_ZERO)
			normal = FastEnvelope::accurate_normal_vector(t0, t1,
				t0, t2);
		else {
			normal = ((t0 - t1).cross(t0 - t2)).normalized();
		}
		std::vector<Vector3> points(3);
		points[0] = t0; points[1] = t1; points[2] = t2;
		build_obb(points, normal, OBB_OFFSET, res.Trans, res.invTrans);
		return res;
	}

	void obb::build_obb(const std::vector<Vector3>& points, const Vector3& normal, const Scalar offset,
		Eigen::Matrix4d &M, Eigen::Matrix4d &invM)
	{
		Eigen::MatrixXd  PS(3, points.size()), PS1(3, points.size());
		Vector3 cent;
		cent <<
			0, 0, 0;

		for (int i = 0; i < points.size(); i++) {
			PS(0, i) = points[i][0];
			PS(1, i) = points[i][1];
			PS(2, i) = points[i][2];

		}
		cent[0] = PS.row(0).sum() / points.size();
		cent[1] = PS.row(1).sum() / points.size();
		cent[2] = PS.row(2).sum() / points.size();
		PS1.row(0) = PS.row(0) - Eigen::MatrixXd::Ones(1, points.size())*cent[0];
		PS1.row(1) = PS.row(1) - Eigen::MatrixXd::Ones(1, points.size())*cent[1];
		PS1.row(2) = PS.row(2) - Eigen::MatrixXd::Ones(1, points.size())*cent[2];
		Eigen::MatrixXd PCA = PS1 * PS1.transpose();
		Eigen::EigenSolver<Eigen::MatrixXd> eg(PCA);
		Eigen::VectorXcd ev = eg.eigenvalues();

		int maxid = 0;

		for (int i = 1; i < 3; i++) {

			if (ev[maxid].real() <= ev(i).real()) {
				maxid = i;
			}
		}
		assert(ev(maxid).real() > 0);
		Eigen::VectorXcd v = eg.eigenvectors().col(maxid);
		Vector3 y;
		y[0] = v[0].real(); y[1] = v[1].real(); y[2] = v[2].real();
		y = y.normalized();
		assert(normal.dot(y) < SCALAR_ZERO);

		Vector3 z = normal.cross(y);
		Matrix3 R;
		R.col(0) = normal;
		R.col(1) = y;
		R.col(2) = z;

		Matrix3 invR = R.transpose();//inv(R)=trans(R)
		Vector3 T0 = -1 * R*cent, T1;//R*cent+T0=[0,0,0]
		Eigen::Matrix4d Trans0, invTrans0, Trans1, invTrans1, Scalling, invScalling;

		Trans0 << R, T0,
			0, 0, 0, 1;
		invTrans0 << invR, cent,
			0, 0, 0, 1;
		Eigen::MatrixXd  PSG(4, points.size()), prog(4, points.size());
		PSG << PS,
			Eigen::MatrixXd::Ones(1, points.size());
		prog = Trans0 * PSG;//contains the projection of the points in new axis
		Vector3 min, max, mid, corner;
		for (int i = 0; i < 3; i++) {
			min[i] = prog.row(i).minCoeff();
			max[i] = prog.row(i).maxCoeff();
		}
		mid = (min + max) / 2;
		corner = max - mid + Vector3(offset, offset, offset);
		Eigen::MatrixXd procg(4, 1), centng(4, 1), centn(3, 1);
		procg << mid,
			1;
		centng = invTrans0 * procg;
		centn(0) = centng(0); centn(1) = centng(1); centn(2) = centng(2);
		T1 = -1 * R*centn;
		Trans1 << R, T1,
			0, 0, 0, 1;
		invTrans1 << invR, centn,
			0, 0, 0, 1;
		Scalling = Eigen::Matrix4d::Zero(4, 4);
		Scalling(0, 0) = corner(0);
		Scalling(1, 1) = corner(1);
		Scalling(2, 2) = corner(2);
		Scalling(3, 3) = 1;
		invScalling = Scalling.inverse();
		M = invScalling * Trans1;
		invM = invTrans1 * Scalling;
	}

	bool obb::obb_intersection(const Eigen::Matrix4d &M1, const Eigen::Matrix4d &invM1, const Eigen::Matrix4d &M2, const Eigen::Matrix4d &invM2) {

		Eigen::MatrixXd ub(8, 4);

		ub << -1, -1, -1, 1,
			1, -1, -1, 1,
			1, 1, -1, 1,
			-1, 1, -1, 1,
			-1, 1, 1, 1,
			-1, -1, 1, 1,
			1, -1, 1, 1,
			1, 1, 1, 1;
		Eigen::MatrixXd b21(4, 8), b12(4, 8);

		Scalar minx, miny, minz, maxx, maxy, maxz;

		b12 = M2 * invM1*ub.transpose();

		minx = b12.row(0).minCoeff();
		if (minx >= 1) return false;
		miny = b12.row(1).minCoeff();
		if (miny >= 1) return false;
		minz = b12.row(2).minCoeff();
		if (minz >= 1) return false;

		maxx = b12.row(0).maxCoeff();
		if (maxx <= -1) return false;
		maxy = b12.row(1).maxCoeff();
		if (maxy <= -1) return false;
		maxz = b12.row(2).maxCoeff();
		if (maxz <= -1) return false;

		b21 = M1 * invM2*ub.transpose();

		minx = b21.row(0).minCoeff();
		if (minx >= 1) return false;
		miny = b21.row(1).minCoeff();
		if (miny >= 1) return false;
		minz = b21.row(2).minCoeff();
		if (minz >= 1) return false;

		maxx = b21.row(0).maxCoeff();
		if (maxx <= -1) return false;
		maxy = b21.row(1).maxCoeff();
		if (maxy <= -1) return false;
		maxz = b21.row(2).maxCoeff();
		if (maxz <= -1) return false;

		return true;

	}
#include <time.h>
	void obb::test() {
		std::vector<Vector3> p, p1;
		Vector3 dis;
		srand(int(time(0)));
		dis = Vector3(rand() % 100, rand() % 100, rand() % 100) / 50;
		p.resize(20);
		for (int i = 0; i < p.size(); i++) {
			p[i] = Vector3(rand() % 100, rand() % 100, rand() % 100) / 100 + dis;
		}
		Eigen::MatrixXd  PS(3, p.size());

		for (int i = 0; i < p.size(); i++) {
			PS(0, i) = p[i][0];
			PS(1, i) = p[i][1];
			PS(2, i) = p[i][2];

		}

		Eigen::MatrixXd PCA = PS * PS.transpose();
		Eigen::EigenSolver<Eigen::MatrixXd> eg(PCA);
		Eigen::VectorXcd ev = eg.eigenvalues();

		int minid = 0;

		for (int i = 1; i < 3; i++) {

			if (ev[minid].real() >= ev(i).real()) {//find the minimal eigenvalue
				minid = i;
			}
		}

		Eigen::VectorXcd v = eg.eigenvectors().col(minid);
		Vector3 normal;
		normal[0] = v[0].real(); normal[1] = v[1].real(); normal[2] = v[2].real();
		Scalar offset = 0.1;
		Eigen::Matrix4d M, invM;
		build_obb(p, normal, offset, M, invM);
		std::ofstream fout;
		fout.open("D:\\vs\\fast_envelope\\obb\\points.txt");

		for (int i = 0; i < p.size(); i++) {

			fout << std::setprecision(17) << p[i][0] << " " << p[i][1] << " " << p[i][2] << std::endl;

		}
		fout.close();
		fout.open("D:\\vs\\fast_envelope\\obb\\matrixes.txt");

		for (int i = 0; i < 4; i++) {

			fout << std::setprecision(17) << M(i, 0) << " " << M(i, 1) << " " << M(i, 2) << " " << M(i, 3) << std::endl;

		}
		for (int i = 0; i < 4; i++) {

			fout << std::setprecision(17) << invM(i, 0) << " " << invM(i, 1) << " " << invM(i, 2) << " " << invM(i, 3) << std::endl;

		}

		fout.close();

		//////////////////////////

		//srand(int(time(0)));
		dis = Vector3(rand() % 100, rand() % 100, rand() % 100) / 100 * 3;
		p.clear();
		p.resize(20);
		for (int i = 0; i < p.size(); i++) {
			p[i] = Vector3(rand() % 100, rand() % 100, rand() % 100) / 100 + dis;
		}
		Eigen::MatrixXd  PS1(3, p.size());

		for (int i = 0; i < p.size(); i++) {
			PS1(0, i) = p[i][0];
			PS1(1, i) = p[i][1];
			PS1(2, i) = p[i][2];

		}

		Eigen::MatrixXd PCA1 = PS1 * PS1.transpose();
		Eigen::EigenSolver<Eigen::MatrixXd> eg1(PCA1);
		Eigen::VectorXcd ev1 = eg1.eigenvalues();

		minid = 0;

		for (int i = 1; i < 3; i++) {

			if (ev1[minid].real() >= ev1(i).real()) {//find the minimal eigenvalue
				minid = i;
			}
		}

		Eigen::VectorXcd v1 = eg1.eigenvectors().col(minid);

		normal[0] = v1[0].real(); normal[1] = v1[1].real(); normal[2] = v1[2].real();

		Eigen::Matrix4d M1, invM1;
		build_obb(p, normal, offset, M1, invM1);

		fout.open("D:\\vs\\fast_envelope\\obb\\points1.txt");

		for (int i = 0; i < p.size(); i++) {

			fout << std::setprecision(17) << p[i][0] << " " << p[i][1] << " " << p[i][2] << std::endl;

		}
		fout.close();
		fout.open("D:\\vs\\fast_envelope\\obb\\matrixes1.txt");

		for (int i = 0; i < 4; i++) {

			fout << std::setprecision(17) << M1(i, 0) << " " << M1(i, 1) << " " << M1(i, 2) << " " << M1(i, 3) << std::endl;

		}
		for (int i = 0; i < 4; i++) {

			fout << std::setprecision(17) << invM1(i, 0) << " " << invM1(i, 1) << " " << invM1(i, 2) << " " << invM1(i, 3) << std::endl;

		}

		fout.close();

		std::cout << "is intersected? " << obb_intersection(M, invM, M1, invM1) << std::endl;
	}
}


// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

/*
 *   File "lstm_b_tapenade_generated.c" is generated by Tapenade 3.14 (r7259) from this file.
 *   To reproduce such a generation you can use Tapenade CLI
 *   (can be downloaded from http://www-sop.inria.fr/tropics/tapenade/downloading.html)
 *
 *   After installing use the next command to generate a file:
 *
 *      tapenade -b -o lstm_tapenade -head "lstm_objective(loss)/(main_params extra_params)" lstm.c
 *
 *   This will produce a file "lstm_tapenade_b.c" which content will be the same as the content of the file "lstm_b_tapenade_generated.c",
 *   except one-line header. Moreover a log-file "lstm_tapenade_b.msg" will be produced.
 *
 *   NOTE: the code in "lstm_b_tapenade_generated.c" is wrong and won't work.
 *         REPAIRED SOURCE IS STORED IN THE FILE "lstm_b.c".
 *         You can either use diff tool or read "lstm_b.c" header to figure out what changes was performed to fix the code.
 *
 *   NOTE: you can also use Tapenade web server (http://tapenade.inria.fr:8080/tapenade/index.jsp)
 *         for generating but the result can be slightly different.
 */

#include "../adbench/lstm.h"

extern "C" {
#include "lstm.h"

// UTILS
// Sigmoid on scalar
double sigmoid(double x)
{
    return 1.0 / (1.0 + exp(-x));
}

// log(sum(exp(x), 2))
double logsumexp(double const* vect, int sz)
{
    double sum = 0.0;
    int i;

    for (i = 0; i < sz; i++)
    {
        sum += exp(vect[i]);
    }

    sum += 2;
    return log(sum);
}

// LSTM OBJECTIVE
// The LSTM model
void lstm_model(
    int hsize,
    double const* __restrict weight,
    double const* __restrict bias,
    double* __restrict hidden,
    double* __restrict cell,
    double const* __restrict input
)
{
    // TODO NOTE THIS
    //__builtin_assume(hsize > 0);

    double* gates = (double*)malloc(4 * hsize * sizeof(double));
    double* forget = &(gates[0]);
    double* ingate = &(gates[hsize]);
    double* outgate = &(gates[2 * hsize]);
    double* change = &(gates[3 * hsize]);

    int i;
    // caching input
    // hidden (needed)
    for (i = 0; i < hsize; i++)
    {
        forget[i] = sigmoid(input[i] * weight[i] + bias[i]);
        ingate[i] = sigmoid(hidden[i] * weight[hsize + i] + bias[hsize + i]);
        outgate[i] = sigmoid(input[i] * weight[2 * hsize + i] + bias[2 * hsize + i]);
        change[i] = tanh(hidden[i] * weight[3 * hsize + i] + bias[3 * hsize + i]);
    }

    // caching cell (needed)
    for (i = 0; i < hsize; i++)
    {
        cell[i] = cell[i] * forget[i] + ingate[i] * change[i];
    }

    for (i = 0; i < hsize; i++)
    {
        hidden[i] = outgate[i] * tanh(cell[i]);
    }

    free(gates);
}

// Predict LSTM output given an input
void lstm_predict(
    int l,
    int b,
    double const* __restrict w,
    double const* __restrict w2,
    double* __restrict s,
    double const* __restrict x,
    double* __restrict x2
)
{
    int i;
    for (i = 0; i < b; i++)
    {
        x2[i] = x[i] * w2[i];
    }

    double* xp = x2;
    for (i = 0; i <= 2 * l * b - 1; i += 2 * b)
    {
        lstm_model(b, &(w[i * 4]), &(w[(i + b) * 4]), &(s[i]), &(s[i + b]), xp);
        xp = &(s[i]);
    }

    for (i = 0; i < b; i++)
    {
        x2[i] = xp[i] * w2[b + i] + w2[2 * b + i];
    }
}

// LSTM objective (loss function)
void lstm_objective(
    int l,
    int c,
    int b,
    double const* __restrict main_params,
    double const* __restrict extra_params,
    double* __restrict state,
    double const* __restrict sequence,
    double* __restrict loss
)
{
    int i, t;
    double total = 0.0;
    int count = 0;
    const double* input = &(sequence[0]);
    double* ypred = (double*)malloc(b * sizeof(double));
    double* ynorm = (double*)malloc(b * sizeof(double));
    const double* ygold;
    double lse;

    __builtin_assume(b>0);
    for (t = 0; t <= (c - 1) * b - 1; t += b)
    {
        lstm_predict(l, b, main_params, extra_params, state, input, ypred);
        lse = logsumexp(ypred, b);
        for (i = 0; i < b; i++)
        {
            ynorm[i] = ypred[i] - lse;
        }

        ygold = &(sequence[t + b]);
        for (i = 0; i < b; i++)
        {
            total += ygold[i] * ynorm[i];
        }

        count += b;
        input = ygold;
    }

    *loss = -total / count;

    free(ypred);
    free(ynorm);
}

extern int enzyme_const;
extern int enzyme_dup;
extern int enzyme_dupnoneed;
void __enzyme_autodiff(...) noexcept;

// *      tapenade -b -o lstm_tapenade -head "lstm_objective(loss)/(main_params extra_params)" lstm.c

void dlstm_objective(
    int l,
    int c,
    int b,
    double const* main_params,
    double* dmain_params,
    double const* extra_params,
    double* dextra_params,
    double* state,
    double const* sequence,
    double* loss,
    double* dloss
)
{
    __enzyme_autodiff(lstm_objective,
        enzyme_const, l,
        enzyme_const, c,
        enzyme_const, b,
        enzyme_dup, main_params, dmain_params,
        enzyme_dup, extra_params, dextra_params,
        enzyme_const, state,
        enzyme_const, sequence,
        enzyme_dupnoneed, loss, dloss
    );
}

}

//! Tapenade
extern "C" {

#include <adBuffer.h>

/*
  Differentiation of sigmoid in reverse (adjoint) mode:
   gradient     of useful results: sigmoid
   with respect to varying inputs: x
*/
// UTILS
// Sigmoid on scalar
void sigmoid_b(double x, double *xb, double sigmoidb) {
    double temp;
    double sigmoid;
    temp = exp(-x) + 1.0;
    *xb = exp(-x)*sigmoidb/(temp*temp);
}

// UTILS
// Sigmoid on scalar
double sigmoid_nodiff(double x) {
    return 1.0/(1.0+exp(-x));
}

/*
  Differentiation of logsumexp in reverse (adjoint) mode:
   gradient     of useful results: logsumexp *vect
   with respect to varying inputs: *vect
   Plus diff mem management of: vect:in
*/
// log(sum(exp(x), 2))
void logsumexp_b(const double *vect, double *vectb, int sz, double logsumexpb)
{
    double sum = 0.0;
    double sumb = 0.0;
    int i;
    double logsumexp;
    for (i = 0; i < sz; ++i)
        sum = sum + exp(vect[i]);
    sum = sum + 2;
    sumb = logsumexpb/sum;
    for (i = sz-1; i > -1; --i)
        vectb[i] = vectb[i] + exp(vect[i])*sumb;
}

// log(sum(exp(x), 2))
double logsumexp_nodiff(const double *vect, int sz) {
    double sum = 0.0;
    int i;
    for (i = 0; i < sz; ++i)
        sum += exp(vect[i]);
    sum += 2;
    return log(sum);
}

/*
  Differentiation of lstm_model in reverse (adjoint) mode:
   gradient     of useful results: alloc(*gates) *cell *bias *hidden
                *weight *input
   with respect to varying inputs: alloc(*gates) *cell *bias *hidden
                *weight *input
   Plus diff mem management of: cell:in bias:in hidden:in weight:in
                input:in
*/
// LSTM OBJECTIVE
// The LSTM model
void lstm_model_b(int hsize, const double *weight, double *weightb, const
        double *bias, double *biasb, double *hidden, double *hiddenb, double *
        cell, double *cellb, const double *input, double *inputb) {
    double *gates;
    double *gatesb;
    double arg1;
    double arg1b;
    int ii1;
    double temp;
    double tempb;
    gatesb = (double *)malloc(4*hsize*sizeof(double));
    for (ii1 = 0; ii1 < 4*hsize; ++ii1)
        gatesb[ii1] = 0.0;
    gates = (double *)malloc(4*hsize*sizeof(double));
    double *forget = &(gates[0]);
    double *forgetb = &(gatesb[0]);
    double *ingate = &(gates[hsize]);
    double *ingateb = &(gatesb[hsize]);
    double *outgate = &(gates[2*hsize]);
    double *outgateb = &(gatesb[2*hsize]);
    double *change = &(gates[3*hsize]);
    double *changeb = &(gatesb[3*hsize]);
    int i;
    for (i = 0; i < hsize; ++i) {
        arg1 = input[i]*weight[i] + bias[i];
        forget[i] = sigmoid_nodiff(arg1);
        arg1 = hidden[i]*weight[hsize+i] + bias[hsize + i];
        ingate[i] = sigmoid_nodiff(arg1);
        arg1 = input[i]*weight[2*hsize+i] + bias[2*hsize + i];
        outgate[i] = sigmoid_nodiff(arg1);
        change[i] = tanh(hidden[i]*weight[3*hsize+i] + bias[3*hsize + i]);
    }
    for (i = 0; i < hsize; ++i) {
        pushReal8(cell[i]);
        cell[i] = cell[i]*forget[i] + ingate[i]*change[i];
    }
    for (i = hsize-1; i > -1; --i) {
        outgateb[i] = outgateb[i] + tanh(cell[i])*hiddenb[i];
        cellb[i] = cellb[i] + outgate[i]*(1.0-tanh(cell[i])*tanh(cell[i]))*
            hiddenb[i];
        hiddenb[i] = 0.0;
    }
    for (i = hsize-1; i > -1; --i) {
        popReal8(&(cell[i]));
        forgetb[i] = forgetb[i] + cell[i]*cellb[i];
        ingateb[i] = ingateb[i] + change[i]*cellb[i];
        changeb[i] = changeb[i] + ingate[i]*cellb[i];
        cellb[i] = forget[i]*cellb[i];
    }
    for (i = hsize-1; i > -1; --i) {
        temp = weight[3*hsize + i];
        tempb = (1.0-tanh(hidden[i]*temp+bias[3*hsize+i])*tanh(hidden[i]*temp+
            bias[3*hsize+i]))*changeb[i];
        hiddenb[i] = hiddenb[i] + temp*tempb;
        weightb[3*hsize + i] = weightb[3*hsize + i] + hidden[i]*tempb;
        biasb[3*hsize + i] = biasb[3*hsize + i] + tempb;
        changeb[i] = 0.0;
        arg1 = input[i]*weight[2*hsize+i] + bias[2*hsize + i];
        sigmoid_b(arg1, &arg1b, outgateb[i]);
        outgateb[i] = 0.0;
        inputb[i] = inputb[i] + weight[2*hsize+i]*arg1b;
        weightb[2*hsize + i] = weightb[2*hsize + i] + input[i]*arg1b;
        biasb[2*hsize + i] = biasb[2*hsize + i] + arg1b;
        arg1 = hidden[i]*weight[hsize+i] + bias[hsize + i];
        sigmoid_b(arg1, &arg1b, ingateb[i]);
        ingateb[i] = 0.0;
        hiddenb[i] = hiddenb[i] + weight[hsize+i]*arg1b;
        weightb[hsize + i] = weightb[hsize + i] + hidden[i]*arg1b;
        biasb[hsize + i] = biasb[hsize + i] + arg1b;
        arg1 = input[i]*weight[i] + bias[i];
        sigmoid_b(arg1, &arg1b, forgetb[i]);
        forgetb[i] = 0.0;
        inputb[i] = inputb[i] + weight[i]*arg1b;
        weightb[i] = weightb[i] + input[i]*arg1b;
        biasb[i] = biasb[i] + arg1b;
    }
    free(gates);
    free(gatesb);
}

// LSTM OBJECTIVE
// The LSTM model
void lstm_model_nodiff(int hsize, const double *weight, const double *bias,
        double *hidden, double *cell, const double *input) {
    double *gates;
    double arg1;
    gates = (double *)malloc(4*hsize*sizeof(double));
    double *forget = &(gates[0]);
    double *ingate = &(gates[hsize]);
    double *outgate = &(gates[2*hsize]);
    double *change = &(gates[3*hsize]);
    int i;
    for (i = 0; i < hsize; ++i) {
        arg1 = input[i]*weight[i] + bias[i];
        forget[i] = sigmoid_nodiff(arg1);
        arg1 = hidden[i]*weight[hsize+i] + bias[hsize + i];
        ingate[i] = sigmoid_nodiff(arg1);
        arg1 = input[i]*weight[2*hsize+i] + bias[2*hsize + i];
        outgate[i] = sigmoid_nodiff(arg1);
        change[i] = tanh(hidden[i]*weight[3*hsize+i] + bias[3*hsize + i]);
    }
    for (i = 0; i < hsize; ++i)
        cell[i] = cell[i]*forget[i] + ingate[i]*change[i];
    for (i = 0; i < hsize; ++i)
        hidden[i] = outgate[i]*tanh(cell[i]);
    free(gates);
}

/*
  Differentiation of lstm_predict in reverse (adjoint) mode:
   gradient     of useful results: alloc(*gates) *s *w *w2 *x2
   with respect to varying inputs: alloc(*gates) *s *w *w2 *x2
   Plus diff mem management of: s:in w:in w2:in x2:in
*/
// Predict LSTM output given an input
void lstm_predict_b(int l, int b, const double *w, double *wb, const double *
        w2, double *w2b, double *s, double *sb, const double *x, double *x2,
        double *x2b) {
    int i;
    double tmp;
    double tmpb;
    for (i = 0; i < b; ++i) {
        pushReal8(x2[i]);
        x2[i] = x[i]*w2[i];
    }
    double *xp = x2;
    double *xpb = x2b;
    for (i = 0; i <= 2*l*b-1; i += 2*b) {
        pushReal8Array(s + i, 2 * b); /* TFIX */
        lstm_model_nodiff(b, &(w[i*4]), &(w[(i+b)*4]), &(s[i]), &(s[i + b]),
                          xp);
        pushPointer8(xpb);
        xpb = &(sb[i]);
        pushPointer8(xp);
        xp = &(s[i]);
    }
    for (i = 0; i < b; ++i) {
        tmp = xp[i]*w2[b+i] + w2[2*b + i];
        pushReal8(x2[i]);
        x2[i] = tmp;
    }
    for (i = b-1; i > -1; --i) {
        popReal8(&(x2[i]));
        tmpb = x2b[i];
        x2b[i] = 0.0;
        xpb[i] = xpb[i] + w2[b+i]*tmpb;
        w2b[b + i] = w2b[b + i] + xp[i]*tmpb;
        w2b[2*b + i] = w2b[2*b + i] + tmpb;
    }
    for (i = 2*l*b-(2*l*b-1)%(2*b)-1; i >= 0; i += -(2*b)) { /* TFIX */
        popPointer8((void **)&xp);
        popPointer8((void **)&xpb);
        popReal8Array(s + i, 2 * b); /* TFIX */
        lstm_model_b(b, &(w[i*4]), &(wb[i*4]), &(w[(i+b)*4]), &(wb[(i+b)*4]),
                     &(s[i]), &(sb[i]), &(s[i + b]), &(sb[i + b]), xp, xpb);
    }
    for (i = b-1; i > -1; --i) {
        popReal8(&(x2[i]));
        w2b[i] = w2b[i] + x[i]*x2b[i];
        x2b[i] = 0.0;
    }
}

// Predict LSTM output given an input
void lstm_predict_nodiff(int l, int b, const double *w, const double *w2,
        double *s, const double *x, double *x2) {
    int i;
    for (i = 0; i < b; ++i)
        x2[i] = x[i]*w2[i];
    double *xp = x2;
    for (i = 0; i <= 2*l*b-1; i += 2*b) {
        lstm_model_nodiff(b, &(w[i*4]), &(w[(i+b)*4]), &(s[i]), &(s[i + b]),
                          xp);
        xp = &(s[i]);
    }
    for (i = 0; i < b; ++i)
        x2[i] = xp[i]*w2[b+i] + w2[2*b + i];
}

/*
  Differentiation of lstm_objective in reverse (adjoint) mode:
   gradient     of useful results: *loss
   with respect to varying inputs: *main_params *extra_params
                *loss
   RW status of diff variables: *main_params:out *extra_params:out
                *loss:in-out
   Plus diff mem management of: extra_params:in loss:in
*/
// LSTM objective (loss function)
void lstm_objective_b(int l, int c, int b, const double *main_params, double *
        main_paramsb, const double *extra_params, double *extra_paramsb,
        double *state, const double *sequence, double *loss, double *lossb) {
    int i, t;
    double total = 0.0;
    double totalb = 0.0;
    int count = 0;
    const double *input = &(sequence[0]);
    double *ypred;
    double *ypredb;
    int ii1;
    int branch;
    double* stateb = (double*)malloc(2 * l * b * sizeof(double)); /* TFIX */
    ypredb = (double *)malloc(b*sizeof(double));
    for (ii1 = 0; ii1 < b; ++ii1)
        ypredb[ii1] = 0.0;
    ypred = (double *)malloc(b*sizeof(double));
    double *ynorm;
    double *ynormb;
    ynormb = (double *)malloc(b*sizeof(double));
    for (ii1 = 0; ii1 < b; ++ii1)
        ynormb[ii1] = 0.0;
    ynorm = (double *)malloc(b*sizeof(double));
    const double* ygold = NULL; /* TFIX */
    double lse;
    double lseb;
    for (t = 0; t <= (c-1)*b-1; t += b) {
        if (ypred) {
            pushReal8Array(ypred, b); /* TFIX */
            pushControl1b(1);
        } else
            pushControl1b(0);
        pushReal8Array(state, 2 * b * l); /* TFIX */
        lstm_predict_nodiff(l, b, main_params, extra_params, state, input,
                            ypred);
        pushPointer8((void*)ygold);
        ygold = &(sequence[t + b]);
        count = count + b;
        pushPointer8((void*)input);
        input = ygold;
    }
    totalb = -(*lossb/count);
    *lossb = 0.0;
    for (ii1 = 0; ii1 < 8 * l * b; ii1++) /* TFIX */
        main_paramsb[ii1] = 0.0;
    for (ii1 = 0; ii1 < 3 * b; ii1++) /* TFIX */
        extra_paramsb[ii1] = 0.0;
    for (t = 0; t < 2 * l * b; t++) /* TFIX */
        stateb[t] = 0.0;
    for (t = (c-1)*b-((c-1)*b-1)%b-1; t >= 0; t += -b) { /* TFIX */
        popPointer8((void **)&input);
        for (i = b-1; i > -1; --i)
            ynormb[i] = ynormb[i] + ygold[i]*totalb;
        popPointer8((void **)&ygold);
        lseb = 0.0;
        for (i = b-1; i > -1; --i) {
            ypredb[i] = ypredb[i] + ynormb[i];
            lseb = lseb - ynormb[i];
            ynormb[i] = 0.0;
        }
        logsumexp_b(ypred, ypredb, b, lseb);
        popReal8Array(state, 2 * b * l); /* TFIX */
        popControl1b(&branch);
        if (branch == 1)
            popReal8Array(ypred, b); /* TFIX */
        lstm_predict_b(l, b, main_params, main_paramsb, extra_params,
                       extra_paramsb, state, stateb, input, ypred, ypredb);
    }
    free(ynorm);
    free(ynormb);
    free(ypred);
    free(ypredb);
    free(stateb); /* TFIX */ // Added to dispose memory allocated in repaired code
}

}

#if 1
//! Adept
#include <adept_source.h>
#include <adept.h>
#include <adept_arrays.h>
using adept::adouble;
using adept::aVector;

namespace adeptTest {
// Sigmoid on scalar
template<typename T>
T sigmoid(T x) {
    return adouble(1) / (adouble(1) + exp(-x));
}

// log(sum(exp(x), 2))
template<typename T>
T logsumexp(const T* vect, int sz) {
    T sum = 0.0;
    for (int i = 0; i < sz; ++i)
        sum += exp(vect[i]);
    sum += adouble(2);
    return log(sum);
}

// LSTM OBJECTIVE

// The LSTM model
template<typename T>
void lstm_model(
    int hsize,
    T* weight,
    T* bias,
    T* hidden,
    T* cell,
    T* input
)
{

    T* gates = new T[4*hsize];
    T* forget = &(gates[0]);
    T* ingate = &(gates[hsize]);
    T* outgate = &(gates[2 * hsize]);
    T* change = &(gates[3 * hsize]);

    int i;
    // caching input
    // hidden (needed)
    for (i = 0; i < hsize; i++)
    {
        forget[i] = sigmoid<adouble>(input[i] * weight[i] + bias[i]);
        ingate[i] = sigmoid<adouble>(hidden[i] * weight[hsize + i] + bias[hsize + i]);
        outgate[i] = sigmoid<adouble>(input[i] * weight[2 * hsize + i] + bias[2 * hsize + i]);
        change[i] = tanh(hidden[i] * weight[3 * hsize + i] + bias[3 * hsize + i]);
    }

    // caching cell (needed)
    for (i = 0; i < hsize; i++)
    {
        cell[i] = cell[i] * forget[i] + ingate[i] * change[i];
    }

    for (i = 0; i < hsize; i++)
    {
        hidden[i] = outgate[i] * tanh(cell[i]);
    }

    delete[] gates;
}

// Predict LSTM output given an input
template<typename T>
void lstm_predict(
    int l,
    int b,
    T* w,
    T* w2,
    T* s,
    T* x,
    T* x2
)
{
    int i;
    for (i = 0; i < b; i++)
    {
        x2[i] = x[i] * w2[i];
    }

    T* xp = x2;
    for (i = 0; i <= 2 * l * b - 1; i += 2 * b)
    {
        lstm_model(b, &(w[i * 4]), &(w[(i + b) * 4]), &(s[i]), &(s[i + b]), xp);
        xp = &(s[i]);
    }

    for (i = 0; i < b; i++)
    {
        x2[i] = xp[i] * w2[b + i] + w2[2 * b + i];
    }
}

// LSTM objective (loss function)
template<typename T>
void lstm_objective(
    int l,
    int c,
    int b,
    T * __restrict main_params,
    T * __restrict extra_params,
    T* __restrict state,
    T * __restrict sequence,
    T* __restrict loss
)
{
    int i, t;
    T total = 0.0;
    int count = 0;
    T* input = &(sequence[0]);
    T* ypred = new T[b];
    T* ynorm = new T[b];
    T* ygold;
    T lse;

    __builtin_assume(b>0);
    for (t = 0; t <= (c - 1) * b - 1; t += b)
    {
        lstm_predict(l, b, main_params, extra_params, state, input, ypred);
        lse = logsumexp(ypred, b);
        for (i = 0; i < b; i++)
        {
            ynorm[i] = ypred[i] - lse;
        }

        ygold = &(sequence[t + b]);
        for (i = 0; i < b; i++)
        {
            total += ygold[i] * ynorm[i];
        }

        count += b;
        input = ygold;
    }

    *loss = -total / adouble(count);

    delete[] ypred;
    delete[] ynorm;
}
};

// Note ADBench did not have an adept impl
void adept_dlstm_objective(int l, int c, int b, const double *main_params, double *
        main_paramsb, const double *extra_params, double *extra_paramsb,
        double *state, const double *sequence, double *loss, double *lossb) {

    int main_sz = 2 * l * 4 * b;
    int extra_sz = 3 * b;
    int state_sz = 2 * l * b;
    int seq_sz = c* b;

  adept::Stack stack;

  adouble *amain = new adouble[main_sz];
  adouble *aextra = new adouble[extra_sz];
  adouble *astate = new adouble[state_sz];
  adouble *aseq = new adouble[seq_sz];

      adept::set_values(amain, main_sz, main_params);
      adept::set_values(aextra, extra_sz, extra_params);
      adept::set_values(astate, state_sz, state);
      adept::set_values(aseq, seq_sz, sequence);

      adouble aloss;

      stack.new_recording();
      adouble aerr;

      adeptTest::lstm_objective(l, c, b, amain, aextra, astate, aseq, &aloss);
      aloss.set_gradient(1.); // only one J row here

      stack.compute_adjoint();

      adept::get_gradients(amain, main_sz, main_paramsb);
      adept::get_gradients(aextra, extra_sz, extra_paramsb);

}
#endif


// File: train_adagrad.cc
// Author: Karl Moritz Hermann (mail@karlmoritz.com)
// Created: 01-01-2013
// Last Update: Thu 03 Oct 2013 11:48:28 AM BST

// STL
#include <iostream>
#include <cmath>
#include <algorithm>
#include <random>

// Boost
#include <boost/program_options/variables_map.hpp>
#include <boost/program_options/parsers.hpp>

// L-BFGS
#include <lbfgs.h>

// Local
#include "train_adagrad.h"
#include "train_update.h"
#include "recursive_autoencoder.h"
#include "utils.h"
#include "fast_math.h"

using namespace std;
namespace bpo = boost::program_options;

int train_adagrad(Model &model, int iterations, float eta, int batches, Real lambda)
{
  Real* vars = nullptr;
  int number_vars = 0;

  setVarsAndNumber(vars,number_vars,model);
  WeightArrayType theta(vars,number_vars);

  Real* Gt_d = new Real[number_vars]();
  // Real* Ginv_d = new Real[number_vars]();
  WeightArrayType Gt(Gt_d,number_vars);
  // WeightArrayType Ginv(Ginv_d,number_vars);
  // Gt.setOnes(); // initialize to ones..

  Real* gradient = new Real[number_vars]();
  lbfgsfloatval_t error = 0;

  // Remove L2 regularization as AdaGrad uses L1 instead
  model.calc_L2 = false;

  int size = model.corpus.size();
  int num_batches = min(batches,size/2);
  int batchsize = max((size/num_batches),2);
  // eta = eta / num_batches;
  cout << "Batch size: " << batchsize << "  eta " << eta << endl;

  Real update;
  Real l1_reg;

  for (auto iteration = 0; iteration < iterations; ++iteration)
  {
    cout << "Iteration " << iteration << endl;
    std::random_shuffle ( model.indexes.begin(), model.indexes.end() );
    for (auto batch = 0; batch < num_batches; ++batch)
    {
      model.from = batch*batchsize;
      model.to = min((batch+1)*batchsize,size);
      error = computeCostAndGrad(model,nullptr,gradient,number_vars);

      // seeing that I need to iterate ...
      for (int i = 0; i < number_vars; ++i) {
        Gt_d[i] += gradient[i]*gradient[i];
        // Update weight: ( eta / \sqrt(Sum square gradients) ) * gradient
        if (Gt_d[i] != 0) {
        update = vars[i] - ((eta / (sqrt(Gt_d[i]))) * gradient[i]);
        l1_reg = (eta / (sqrt(Gt_d[i]))) * lambda;
        vars[i] = signum(update) * max(0.0, abs(update) - l1_reg);
        }
      }
    }

    if (iteration % model.rae.config.dump_freq == 0)
    {
      printf("Dumping model ...\n");
      dumpModel(model,iteration);
    }
  }

  delete [] gradient;
  // delete [] Ginv_d;
  delete [] Gt_d;
  return 0;
}


#pragma once
#include <unordered_set>
#include <numeric>
#include "TopicModel.hpp"
#include <Eigen/Dense>
#include "../Utils/Utils.hpp"
#include "../Utils/math.h"
#include "../Utils/sample.hpp"

/*
Implementation of LDA using Collapsed Variational Bayes zero-order estimation by bab2min

* Blei, D. M., Ng, A. Y., & Jordan, M. I. (2003). Latent dirichlet allocation. Journal of machine Learning research, 3(Jan), 993-1022.

Term Weighting Scheme is based on following paper:
* Wilson, A. T., & Chew, P. A. (2010, June). Term weighting schemes for latent dirichlet allocation. In human language technologies: The 2010 annual conference of the North American Chapter of the Association for Computational Linguistics (pp. 465-473). Association for Computational Linguistics.

*/

#define GETTER(name, type, field) type get##name() const override { return field; }
namespace tomoto
{
	struct DocumentLDACVB0 : public DocumentBase
	{
	public:
		using DocumentBase::DocumentBase;

		Eigen::MatrixXf Zs;
		Eigen::VectorXf numByTopic;

		DEFINE_SERIALIZER_AFTER_BASE(DocumentBase, Zs);

		template<typename _TopicModel> void update(Float* ptr, const _TopicModel& mdl);

		int32_t getSumWordWeight() const
		{
			return this->words.size();
		}
	};

	struct ModelStateLDACVB0
	{
		Eigen::VectorXf zLikelihood;
		Eigen::VectorXf numByTopic;
		Eigen::MatrixXf numByTopicWord;

		DEFINE_SERIALIZER(numByTopic, numByTopicWord);
	};

	class ILDACVB0Model : public ITopicModel
	{
	public:
		using DefaultDocType = DocumentLDACVB0;
		static ILDACVB0Model* create(size_t _K = 1, Float _alpha = 0.1, Float _eta = 0.01, size_t _rg = std::random_device{}());

		virtual size_t addDoc(const std::vector<std::string>& words) = 0;
		virtual std::unique_ptr<DocumentBase> makeDoc(const std::vector<std::string>& words) const = 0;
		TermWeight getTermWeight() const { return TermWeight::one; };
		virtual size_t getOptimInterval() const = 0;
		virtual void setOptimInterval(size_t) = 0;
		virtual void setBurnInIteration(size_t) {}
		virtual std::vector<size_t> getCountByTopic() const = 0;
		virtual size_t getK() const = 0;
		virtual Float getAlpha() const = 0;
		virtual Float getEta() const = 0;

		virtual std::vector<Float> getWordPrior(const std::string& word) const { return {}; }
		virtual void setWordPrior(const std::string& word, const std::vector<Float>& priors) {}
	};

	template<typename _Interface = ILDACVB0Model,
		typename _Derived = void,
		typename _DocType = DocumentLDACVB0,
		typename _ModelState = ModelStateLDACVB0>
	class LDACVB0Model : public TopicModel<0, _Interface,
		typename std::conditional<std::is_same<_Derived, void>::value, LDACVB0Model<>, _Derived>::type,
		_DocType, _ModelState>
	{
	protected:
		using DerivedClass = typename std::conditional<std::is_same<_Derived, void>::value, LDACVB0Model, _Derived>::type;
		using BaseClass = TopicModel<0, _Interface, DerivedClass, _DocType, _ModelState>;
		friend BaseClass;

		static constexpr const char TWID[] = "one\0";
		static constexpr static constexpr char TMID[] = "LDA\0";

		Float alpha;
		Vector alphas;
		Float eta;
		Tid K;
		size_t optimInterval = 50;

		template<typename _List>
		static Float calcDigammaSum(_List list, size_t len, Float alpha)
		{
			auto listExpr = Vector::NullaryExpr(len, list);
			auto dAlpha = math::digammaT(alpha);
			return (math::digammaApprox(listExpr.array() + alpha) - dAlpha).sum();
		}

		void optimizeParameters(ThreadPool& pool, _ModelState* localData)
		{
			const auto K = this->K;
			for (size_t i = 0; i < 5; ++i)
			{
				Float denom = calcDigammaSum([&](size_t i) { return this->docs[i].getSumWordWeight(); }, this->docs.size(), alphas.sum());
				for (size_t k = 0; k < K; ++k)
				{
					Float nom = calcDigammaSum([&](size_t i) { return this->docs[i].numByTopic[k]; }, this->docs.size(), alphas(k));
					alphas(k) = std::max(nom / denom * alphas(k), 1e-5f);
				}
			}
		}

		const Eigen::VectorXf& getZLikelihoods(_ModelState& ld, const _DocType& doc, size_t docId, size_t vid) const
		{
			const size_t V = this->realV;
			assert(vid < V);
			auto& zLikelihood = ld.zLikelihood;
			zLikelihood = (doc.numByTopic.array().template cast<Float>() + alphas.array())
				* (ld.numByTopicWord.col(vid).array().template cast<Float>() + eta)
				/ (ld.numByTopic.array().template cast<Float>() + V * eta);
			zLikelihood /= zLikelihood.sum() + 1e-10;
			return zLikelihood;
		}

		template<int _Inc, typename _Vec>
		inline void addWordTo(_ModelState& ld, _DocType& doc, uint32_t pid, Vid vid, _Vec tDist) const
		{
			assert(vid < this->realV);
			constexpr bool _dec = _Inc < 0;
			doc.numByTopic += _Inc * tDist;
			if (_dec) doc.numByTopic = doc.numByTopic.cwiseMax(0);
			ld.numByTopic += _Inc * tDist;
			if (_dec) ld.numByTopic = ld.numByTopic.cwiseMax(0);
			ld.numByTopicWord.col(vid) += _Inc * tDist;
			if (_dec) ld.numByTopicWord.col(vid) = ld.numByTopicWord.col(vid).cwiseMax(0);
		}

		template<ParallelScheme _ps, bool _infer, typename _ExtraDocData>
		void sampleDocument(_DocType& doc, const _ExtraDocData& edd, size_t docId, _ModelState& ld, _RandGen& rgs, size_t iterationCnt, size_t partitionId = 0) const
		{
			for (size_t w = 0; w < doc.words.size(); ++w)
			{
				if (doc.words[w] >= this->realV) continue;
				addWordTo<-1>(ld, doc, w, doc.words[w], doc.Zs.col(w));
				doc.Zs.col(w) = static_cast<const DerivedClass*>(this)->getZLikelihoods(ld, doc, docId, doc.words[w]);
				addWordTo<1>(ld, doc, w, doc.words[w], doc.Zs.col(w));
			}
		}

		template<typename _DocIter, typename _ExtraDocData>
		void updatePartition(ThreadPool& pool, _ModelState* localData, _DocIter first, _DocIter last, _ExtraDocData& edd)
		{
		}

		template<ParallelScheme _ps>
		void trainOne(ThreadPool& pool, _ModelState* localData, _RandGen* rgs)
		{
			std::vector<std::future<void>> res;
			const size_t chStride = std::min(pool.getNumWorkers() * 8, this->docs.size());
			for (size_t ch = 0; ch < chStride; ++ch)
			{
				res.emplace_back(pool.enqueue([&, this, ch, chStride](size_t threadId)
				{
					forShuffled((this->docs.size() - 1 - ch) / chStride + 1, rgs[threadId](), [&, this](size_t id)
					{
						static_cast<DerivedClass*>(this)->template sampleDocument<ParallelScheme::copy_merge>(
							this->docs[id * chStride + ch], 0, id * chStride + ch,
							localData[threadId], rgs[threadId], this->globalStep);
					});
				}));
			}
			for (auto& r : res) r.get();
			static_cast<DerivedClass*>(this)->updateGlobalInfo(pool, localData);
			static_cast<DerivedClass*>(this)->mergeState(pool, this->globalState, this->tState, localData);
			if (this->globalStep >= 250 && optimInterval && (this->globalStep + 1) % optimInterval == 0)
			{
				static_cast<DerivedClass*>(this)->optimizeParameters(pool, localData);
			}
		}

		void updateGlobalInfo(ThreadPool& pool, _ModelState* localData)
		{
			std::vector<std::future<void>> res;

			this->globalState.numByTopic.setZero();
			this->globalState.numByTopicWord.setZero();
			for (auto& doc : this->docs)
			{
				doc.numByTopic = doc.Zs.rowwise().sum();
				this->globalState.numByTopic += doc.numByTopic;
				for (size_t i = 0; i < doc.words.size(); ++i)
				{
					this->globalState.numByTopicWord.col(doc.words[i]) += doc.Zs.col(i);
				}
			}

			for (size_t i = 0; i < pool.getNumWorkers(); ++i)
			{
				res.emplace_back(pool.enqueue([&, i](size_t threadId)
				{
					localData[i] = this->globalState;
				}));
			}
			for (auto& r : res) r.get();
		}

		void mergeState(ThreadPool& pool, _ModelState& globalState, _ModelState& tState, _ModelState* localData) const
		{
		}

		template<typename _DocIter>
		double getLLDocs(_DocIter _first, _DocIter _last) const
		{
			double ll = 0;
			// doc-topic distribution
			ll += (math::lgammaT(K*alpha) - math::lgammaT(alpha)*K) * std::distance(_first, _last);
			for (; _first != _last; ++_first)
			{
				auto& doc = *_first;
				ll -= math::lgammaT(doc.getSumWordWeight() + K * alpha);
				for (Tid k = 0; k < K; ++k)
				{
					ll += math::lgammaT(doc.numByTopic[k] + alpha);
				}
			}
			return ll;
		}

		double getLLRest(const _ModelState& ld) const
		{
			double ll = 0;
			const size_t V = this->realV;
			// topic-word distribution
			// it has the very-small-value problem
			ll += (math::lgammaT(V*eta) - math::lgammaT(eta)*V) * K;
			for (Tid k = 0; k < K; ++k)
			{
				ll -= math::lgammaT(ld.numByTopic[k] + V * eta);
				for (Vid v = 0; v < V; ++v)
				{
					ll += math::lgammaT(ld.numByTopicWord(k, v) + eta);
				}
			}
			return ll;
		}

		double getLL() const
		{
			return static_cast<const DerivedClass*>(this)->template getLLDocs<>(this->docs.begin(), this->docs.end())
				+ static_cast<const DerivedClass*>(this)->getLLRest(this->globalState);
		}

		void prepareShared()
		{
		}

		void prepareDoc(_DocType& doc, Float* topicDocPtr, size_t wordSize) const
		{
			doc.numByTopic = Eigen::VectorXf::Zero(K);
			doc.Zs = Eigen::MatrixXf::Zero(K, wordSize);
		}

		void initGlobalState(bool initDocs)
		{
			const size_t V = this->realV;
			this->globalState.zLikelihood = Vector::Zero(K);
			if (initDocs)
			{
				this->globalState.numByTopic = Vector::Zero(K);
				this->globalState.numByTopicWord = Matrix::Zero(K, V);
			}
		}

		struct Generator
		{
			std::uniform_int_distribution<Tid> theta;
		};

		Generator makeGeneratorForInit(const _DocType*) const
		{
			return Generator{ std::uniform_int_distribution<Tid>{0, (Tid)(K - 1)} };
		}

		template<bool _Infer>
		void updateStateWithDoc(Generator& g, _ModelState& ld, _RandGen& rgs, _DocType& doc, size_t i) const
		{
			doc.Zs.col(i).setZero();
			doc.Zs(g.theta(rgs), i) = 1;
			addWordTo<1>(ld, doc, i, doc.words[i], doc.Zs.col(i));
		}

		template<bool _Infer, typename _Generator>
		void initializeDocState(_DocType& doc, Float* topicDocPtr, _Generator& g, _ModelState& ld, _RandGen& rgs) const
		{
			std::vector<uint32_t> tf(this->realV);
			static_cast<const DerivedClass*>(this)->prepareDoc(doc, topicDocPtr, doc.words.size());

			for (size_t i = 0; i < doc.words.size(); ++i)
			{
				if (doc.words[i] >= this->realV) continue;
				static_cast<const DerivedClass*>(this)->template updateStateWithDoc<_Infer>(g, ld, rgs, doc, i);
			}
		}

		std::vector<uint64_t> _getTopicsCount() const
		{
			Eigen::VectorXf cnt = Eigen::VectorXf::Zero(K);
			for (auto& doc : this->docs)
			{
				cnt += doc.Zs.rowwise().sum();
			}

			return { cnt.data(), cnt.data() + K };
		}

		template<ParallelScheme _ps>
		size_t estimateMaxThreads() const
		{
			if (_ps == ParallelScheme::partition)
			{
				return this->realV / 4;
			}
			if (_ps == ParallelScheme::copy_merge)
			{
				return this->docs.size() / 2;
			}
			return (size_t)-1;
		}

		DEFINE_SERIALIZER(alpha, eta, K);

	public:
		LDACVB0Model(size_t _K = 1, Float _alpha = 0.1, Float _eta = 0.01, size_t _rg = std::random_device{}())
			: BaseClass(_rg), K(_K), alpha(_alpha), eta(_eta)
		{
			alphas = Vector::Constant(K, alpha);
		}
		GETTER(K, size_t, K);
		GETTER(Alpha, Float, alpha);
		GETTER(Eta, Float, eta);
		GETTER(OptimInterval, size_t, optimInterval);

		void setOptimInterval(size_t _optimInterval) override
		{
			optimInterval = _optimInterval;
		}

		size_t addDoc(const std::vector<std::string>& words) override
		{
			return this->_addDoc(this->_makeDoc(words));
		}

		std::unique_ptr<DocumentBase> makeDoc(const std::vector<std::string>& words) const override
		{
			return std::make_unique<_DocType>(as_mutable(this)->template _makeDoc<true>(words));
		}

		void updateDocs()
		{
			for (auto& doc : this->docs)
			{
				doc.template update<>(nullptr, *static_cast<DerivedClass*>(this));
			}
		}

		void prepare(bool initDocs = true, size_t minWordCnt = 0, size_t minWordDf = 0, size_t removeTopN = 0) override
		{
			if (initDocs) this->removeStopwords(minWordCnt, minWordDf, removeTopN);
			static_cast<DerivedClass*>(this)->updateWeakArray();
			static_cast<DerivedClass*>(this)->initGlobalState(initDocs);

			if (initDocs)
			{
				auto generator = static_cast<DerivedClass*>(this)->makeGeneratorForInit(nullptr);
				for (auto& doc : this->docs)
				{
					initializeDocState<false>(doc, nullptr, generator, this->globalState, this->rg);
				}
			}
			else
			{
				static_cast<DerivedClass*>(this)->updateDocs();
			}
			static_cast<DerivedClass*>(this)->prepareShared();
		}

		std::vector<size_t> getCountByTopic() const override
		{
			return static_cast<const DerivedClass*>(this)->_getTopicsCount();
		}

		std::vector<Float> getTopicsByDoc(const _DocType& doc) const
		{
			std::vector<Float> ret(K);
			Float sum = doc.getSumWordWeight() + K * alpha;
			transform(doc.numByTopic.data(), doc.numByTopic.data() + K, ret.begin(), [sum, this](size_t n)
			{
				return (n + alpha) / sum;
			});
			return ret;
		}

		std::vector<Float> _getWidsByTopic(Tid tid, bool normalize = true) const
		{
			assert(tid < K);
			const size_t V = this->realV;
			std::vector<Float> ret(V);
			Float sum = this->globalState.numByTopic[tid] + V * eta;
			auto r = this->globalState.numByTopicWord.row(tid);
			for (size_t v = 0; v < V; ++v)
			{
				ret[v] = (r[v] + eta) / sum;
			}
			return ret;
		}

		template<bool _Together, ParallelScheme _ps, typename _Iter>
		std::vector<double> _infer(_Iter docFirst, _Iter docLast, size_t maxIter, Float tolerance, size_t numWorkers) const
		{
			return {};
		}
	};

	template<typename _TopicModel>
	void DocumentLDACVB0::update(Float * ptr, const _TopicModel & mdl)
	{
		numByTopic = Eigen::VectorXf::Zero(mdl.getK());
		for (size_t i = 0; i < Zs.cols(); ++i)
		{
			numByTopic += Zs.col(i);
		}
	}

	inline ILDACVB0Model* ILDACVB0Model::create(size_t _K, Float _alpha, Float _eta, const _RandGen& _rg)
	{
		return new LDACVB0Model<>(_K, _alpha, _eta, _rg);
	}

}


// Copyright 2019, Collabora, Ltd.
// SPDX-License-Identifier: BSL-1.0
/*!
 * @file
 * @brief  Base implementations for math library.
 * @author Jakob Bornecrantz <jakob@collabora.com>
 * @author Ryan Pavlik <ryan.pavlik@collabora.com>
 * @ingroup aux_math
 */

#include <Eigen/Core>
#include <Eigen/Geometry>

#include <assert.h>

#include "math/m_api.h"
#include "math/m_eigen_interop.h"

/*
 *
 * Copy helpers.
 *
 */

static inline Eigen::Quaternionf
copy(const struct xrt_quat& q)
{
	// Eigen constructor order is different from XRT, OpenHMD and OpenXR!
	//  Eigen: `float w, x, y, z`.
	// OpenXR: `float x, y, z, w`.
	return Eigen::Quaternionf(q.w, q.x, q.y, q.z);
}

static inline Eigen::Quaternionf
copy(const struct xrt_quat* q)
{
	return copy(*q);
}

static inline Eigen::Vector3f
copy(const struct xrt_vec3& v)
{
	return Eigen::Vector3f(v.x, v.y, v.z);
}

static inline Eigen::Vector3f
copy(const struct xrt_vec3* v)
{
	return copy(*v);
}

/*
 *
 * Exported vector functions.
 *
 */
void
math_vec3_accum(const struct xrt_vec3* additional, struct xrt_vec3* inAndOut)
{
	assert(additional != NULL);
	assert(inAndOut != NULL);

	map_vec3(*inAndOut) += map_vec3(*additional);
}

/*
 *
 * Exported quaternion functions.
 *
 */

void
math_quat_rotate(const struct xrt_quat* left,
                 const struct xrt_quat* right,
                 struct xrt_quat* result)
{
	assert(left != NULL);
	assert(right != NULL);
	assert(result != NULL);

	auto l = copy(left);
	auto r = copy(right);

	auto q = l * r;

	map_quat(*result) = q;
}

void
math_quat_rotate_vec3(const struct xrt_quat* left,
                      const struct xrt_vec3* right,
                      struct xrt_vec3* result)
{
	assert(left != NULL);
	assert(right != NULL);
	assert(result != NULL);

	auto l = copy(left);
	auto r = copy(right);

	auto v = l * r;

	map_vec3(*result) = v;
}

/*
 *
 * Exported pose functions.
 *
 */

bool
math_pose_validate(const struct xrt_pose* pose)
{
	assert(pose != NULL);

	const float FLOAT_EPSILON = Eigen::NumTraits<float>::epsilon();
	auto norm = orientation(*pose).squaredNorm();
	if (norm > 1.0f + FLOAT_EPSILON || norm < 1.0f - FLOAT_EPSILON) {
		return false;
	}

	// Technically not yet a required check, but easier to stop problems
	// now than once denormalized numbers pollute the rest of our state.
	// see https://gitlab.khronos.org/openxr/openxr/issues/922
	if (!orientation(*pose).coeffs().allFinite()) {
		return false;
	}
	if (!position(*pose).allFinite()) {
		return false;
	}
	return true;
}

void
math_pose_invert(const struct xrt_pose* pose, struct xrt_pose* outPose)
{
	assert(pose != NULL);
	assert(outPose != NULL);

	// store results to temporary locals so we can do this "in-place"
	// (pose == outPose) if desired.
	Eigen::Vector3f newPosition = -position(*pose);
	// Conjugate legal here since pose must be normalized/unit length.
	Eigen::Quaternionf newOrientation = orientation(*pose).conjugate();

	position(*outPose) = newPosition;
	orientation(*outPose) = newOrientation;
}

/*!
 * Return the result of transforming a point by a pose/transform.
 */
static inline Eigen::Vector3f
transform_point(const xrt_pose& transform, const xrt_vec3& point)
{
	return orientation(transform) * map_vec3(point) + position(transform);
}

/*!
 * Return the result of transforming a pose by a pose/transform.
 */
static inline xrt_pose
transform_pose(const xrt_pose& transform, const xrt_pose& pose)
{
	xrt_pose ret;
	position(ret) = transform_point(transform, pose.position);
	orientation(ret) = orientation(transform) * orientation(pose);
	return ret;
}

void
math_pose_transform(const struct xrt_pose* transform,
                    const struct xrt_pose* pose,
                    struct xrt_pose* outPose)
{
	assert(pose != NULL);
	assert(transform != NULL);
	assert(outPose != NULL);

	xrt_pose newPose = transform_pose(*transform, *pose);
	memcpy(outPose, &newPose, sizeof(xrt_pose));
}

void
math_pose_openxr_locate(const struct xrt_pose* space_pose,
                        const struct xrt_pose* relative_pose,
                        const struct xrt_pose* base_space_pose,
                        struct xrt_pose* result)
{
	assert(space_pose != NULL);
	assert(relative_pose != NULL);
	assert(base_space_pose != NULL);
	assert(result != NULL);

	// Compilers are slighty better optimizing
	// if we copy the arguments in one go.
	const auto bsp = *base_space_pose;
	const auto rel = *relative_pose;
	const auto spc = *space_pose;
	struct xrt_pose pose;

	// Apply the invert of the base space to identity.
	math_pose_invert(&bsp, &pose);

	// Apply the pure pose from the space relation.
	math_pose_transform(&pose, &rel, &pose);

	// Apply the space pose.
	math_pose_transform(&pose, &spc, &pose);

	*result = pose;
}

/*!
 * Return the result of rotating a derivative vector by a matrix.
 *
 * This is a differential transform.
 */
static inline Eigen::Vector3f
rotate_deriv(Eigen::Matrix3f const& rotation,
             const xrt_vec3& derivativeVector,
             Eigen::Matrix3f const& rotationInverse)
{
	return ((rotation * map_vec3(derivativeVector)).transpose() *
	        rotationInverse)
	    .transpose();
}

#ifndef XRT_DOXYGEN

#define MAKE_REL_FLAG_CHECK(NAME, MASK)                                        \
	static inline bool NAME(xrt_space_relation_flags flags)                \
	{                                                                      \
		return ((flags & (MASK)) != 0);                                \
	}

MAKE_REL_FLAG_CHECK(has_some_pose_component,
                    XRT_SPACE_RELATION_POSITION_VALID_BIT |
                        XRT_SPACE_RELATION_ORIENTATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_position, XRT_SPACE_RELATION_POSITION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_orientation, XRT_SPACE_RELATION_ORIENTATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_lin_vel, XRT_SPACE_RELATION_LINEAR_VELOCITY_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_ang_vel, XRT_SPACE_RELATION_ANGULAR_VELOCITY_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_lin_acc,
                    XRT_SPACE_RELATION_LINEAR_ACCELERATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_ang_acc,
                    XRT_SPACE_RELATION_ANGULAR_ACCELERATION_VALID_BIT)
MAKE_REL_FLAG_CHECK(has_some_derivative,
                    XRT_SPACE_RELATION_LINEAR_VELOCITY_VALID_BIT |
                        XRT_SPACE_RELATION_ANGULAR_VELOCITY_VALID_BIT |
                        XRT_SPACE_RELATION_LINEAR_ACCELERATION_VALID_BIT |
                        XRT_SPACE_RELATION_ANGULAR_ACCELERATION_VALID_BIT)

#undef MAKE_REL_FLAG_CHECK

#endif // !XRT_DOXYGEN

/*!
 * Apply a transform to a space relation.
 */
static inline void
transform_accumulate_pose(const xrt_pose& transform,
                          xrt_space_relation& relation,
                          bool do_translation = true,
                          bool do_rotation = true)
{
	assert(do_translation || do_rotation);

	// Save the quat in case we are self-transforming.
	Eigen::Quaternionf quat = orientation(transform);

	auto flags = relation.relation_flags;
	// so code looks similar
	auto in_out_relation = &relation;

	// transform (rotate and translate) the pose, if applicable.
	if (has_some_pose_component(flags)) {
		// Zero out transform parts we don't want to use,
		// because math_pose_transform doesn't take flags.
		xrt_pose transform_copy = transform;
		if (!do_translation) {
			position(transform_copy) = Eigen::Vector3f::Zero();
		}
		if (!do_rotation) {
			orientation(transform_copy) =
			    Eigen::Quaternionf::Identity();
		}

		math_pose_transform(&in_out_relation->pose, &transform,
		                    &in_out_relation->pose);
	}

	if (do_rotation && has_some_derivative(flags)) {

		// prepare matrices required for rotating derivatives from the
		// saved quat.
		Eigen::Matrix3f rot = quat.toRotationMatrix();
		Eigen::Matrix3f rotInverse = rot.inverse();

		// Rotate derivatives, if applicable.
		if (has_lin_vel(flags)) {
			map_vec3(in_out_relation->linear_velocity) =
			    rotate_deriv(rot, in_out_relation->linear_velocity,
			                 rotInverse);
		}

		if (has_ang_vel(flags)) {
			map_vec3(in_out_relation->angular_velocity) =
			    rotate_deriv(rot, in_out_relation->angular_velocity,
			                 rotInverse);
		}

		if (has_lin_acc(flags)) {
			map_vec3(in_out_relation->linear_acceleration) =
			    rotate_deriv(rot,
			                 in_out_relation->linear_acceleration,
			                 rotInverse);
		}

		if (has_ang_acc(flags)) {
			map_vec3(in_out_relation->angular_acceleration) =
			    rotate_deriv(rot,
			                 in_out_relation->angular_acceleration,
			                 rotInverse);
		}
	}
}

static const struct xrt_space_relation BLANK_RELATION = {
    XRT_SPACE_RELATION_BITMASK_ALL,
    {{0.0f, 0.0f, 0.0f, 1.0f}, {0.0f, 0.0f, 0.0f}},
    {0, 0, 0},
    {0, 0, 0},
    {0, 0, 0},
    {0, 0, 0},
};

void
math_relation_reset(struct xrt_space_relation* out)
{
	*out = BLANK_RELATION;
}

void
math_relation_accumulate_transform(const struct xrt_pose* transform,
                                   struct xrt_space_relation* in_out_relation)
{
	assert(transform != nullptr);
	assert(in_out_relation != nullptr);

	// No modifying the validity flags here.
	transform_accumulate_pose(*transform, *in_out_relation);
}

void
math_relation_accumulate_relation(
    const struct xrt_space_relation* additional_relation,
    struct xrt_space_relation* in_out_relation)
{
	assert(additional_relation != NULL);
	assert(in_out_relation != NULL);

	// Update the flags.
	xrt_space_relation_flags flags = (enum xrt_space_relation_flags)(
	    in_out_relation->relation_flags &
	    additional_relation->relation_flags);
	in_out_relation->relation_flags = flags;

	if (has_some_pose_component(flags)) {
		// First, just do the pose part (including rotating
		// derivatives, if applicable).
		transform_accumulate_pose(additional_relation->pose,
		                          *in_out_relation, has_position(flags),
		                          has_orientation(flags));
	}

	// Then, accumulate the derivatives, if required.
	if (has_lin_vel(flags)) {
		map_vec3(in_out_relation->linear_velocity) +=
		    map_vec3(additional_relation->linear_velocity);
	}

	if (has_ang_vel(flags)) {
		map_vec3(in_out_relation->angular_velocity) +=
		    map_vec3(additional_relation->angular_velocity);
	}

	if (has_lin_acc(flags)) {
		map_vec3(in_out_relation->linear_acceleration) +=
		    map_vec3(additional_relation->linear_acceleration);
	}

	if (has_ang_acc(flags)) {
		map_vec3(in_out_relation->angular_acceleration) +=
		    map_vec3(additional_relation->angular_acceleration);
	}
}

void
math_relation_openxr_locate(const struct xrt_pose* space_pose,
                            const struct xrt_space_relation* relative_relation,
                            const struct xrt_pose* base_space_pose,
                            struct xrt_space_relation* result)
{
	assert(space_pose != NULL);
	assert(relative_relation != NULL);
	assert(base_space_pose != NULL);
	assert(result != NULL);

	// Compilers are slighty better optimizing
	// if we copy the arguments in one go.
	const auto bsp = *base_space_pose;
	const auto spc = *space_pose;
	struct xrt_space_relation accumulating_relation = BLANK_RELATION;

	// Apply the invert of the base space to identity.
	math_pose_invert(&bsp, &accumulating_relation.pose);

	// Apply the pure relation between spaces.
	math_relation_accumulate_relation(relative_relation,
	                                  &accumulating_relation);

	// Apply the space pose.
	math_relation_accumulate_transform(&spc, &accumulating_relation);

	*result = accumulating_relation;
}


#include <vector>
#include <string>
#include <map>

#include "rate_model.h"

#include <armadillo>
using namespace arma;

inline int signof(double d) {return d >= 0 ? 1 : -1;}
inline double roundto(double in) {return floor(in*(1000)+0.5)/(1000);}

RateModel::RateModel(int _nstates):Q(_nstates,_nstates),labels(),Q_mask(),
    lasteigval(_nstates,_nstates),lasteigvec(_nstates,_nstates),
    eigval(_nstates,_nstates),eigvec(_nstates,_nstates),
    lasteigval_simple(_nstates,_nstates),lasteigvec_simple(_nstates,_nstates),
    eigval_simple(_nstates,_nstates),eigvec_simple(_nstates,_nstates),nstates(_nstates) {

    setup_Q();
    sameQ = false;
    lasteigval.fill(0);
    lasteigvec.fill(0);
    eigval.fill(0);
    eigvec.fill(0);
    lasteigval_simple.fill(0);
    lasteigvec_simple.fill(0);
    eigval_simple.fill(0);
    eigvec_simple.fill(0);
}

void RateModel::set_Q_cell(int from, int to, double num) {
    Q(from,to) = num;
    sameQ = false;
}

void RateModel::set_Q_diag() {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i != j) {
                su += Q(i, j);
            }
        }
        Q(i, i) = 0-su;
    }
    sameQ = false;
}

void RateModel::setup_Q() {
    Q.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i != j) {
                Q(i, j) = 1./nstates;
            } else {
                Q(i, j) = -(1./nstates * (nstates-1));
            }
        }
    }
    sameQ = false;
}

void RateModel::setup_Q(vector<vector<double> > & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ[i][j];
        }
    }
    for (unsigned int i=0; i < Q.n_rows; i++) {
        colvec a = (sum(Q,1));
        Q(i, i) = -(a(i)-Q(i, i));
    }
    sameQ = false;
}

void RateModel::setup_Q(mat & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ(i, j);
        }
    }
    set_Q_diag();
    sameQ = false;
}

void RateModel::set_n_qs(int number) {
    for (int i=0; i < number; i++) {
        mat tm(nstates,nstates);
        Qs.push_back(tm);
    }
}

void RateModel::set_Q_which(mat & inQ,int which) {
    for (unsigned int i=0; i < Qs[which].n_rows; i++) {
        for (unsigned int j=0; j < Qs[which].n_cols; j++) {
            Qs[which](i, j) = inQ(i, j);
        }
    }
    for (unsigned int i=0; i < Qs[which].n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < Qs[which].n_cols; j++) {
            if (i != j) {
                su += Qs[which](i, j);
            }
        }
        Qs[which](i, i) = 0-su;
    }
    sameQ = false;
}

void RateModel::set_Q(mat & inQ) {
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            Q(i, j) = inQ(i, j);
        }
    }
    sameQ = false;
}

mat & RateModel::get_Q() {
    return Q;
}

cx_mat RateModel::setup_P(double bl,bool store_p_matrices) {
    //sameQ = false;
    eigvec.fill(0);
    eigval.fill(0);
    bool isImag = get_eigenvec_eigenval_from_Q(&eigval, &eigvec); // isImag is not used?
    //cout << eigval << endl;
    //cout << eigvec << endl;
    for (int i=0; i < nstates; i++) {
        eigval(i, i) = exp(eigval(i, i) * bl);
    }
    cx_mat C_inv = inv(eigvec);
    cx_mat P = eigvec * eigval * C_inv;
    neg_p = false;
    for (unsigned int i=0; i < P.n_rows; i++) {
        for (unsigned int j=0; j < P.n_cols; j++) {
            if (real(P(i, j))<0)
            neg_p = true;
        }
    }
    if (store_p_matrices == true) {
        stored_p_matrices[bl] = P;
    }
    return P;
}

void RateModel::setup_P_simple(mat & p,double bl,bool store_p_matrices) {
//    sameQ = false;
    eigvec_simple.fill(0);
    eigval_simple.fill(0);
    get_eigenvec_eigenval_from_Q_simple(&eigval_simple, &eigvec_simple);
    //cout << eigval << endl;
    //cout << eigvec << endl;
    for (int i=0; i < nstates; i++) {
        eigval_simple(i, i) = exp(eigval_simple(i, i) * bl);
    }
    mat C_inv = inv(eigvec_simple);
    p = eigvec_simple * eigval_simple * C_inv;
/*    if (store_p_matrices == true) {
    stored_p_matrices[bl] = P;
    }*/
    //return P;
}

void RateModel::get_eigenvec_eigenval_from_Q_simple(mat * eigval, mat * eigvec) {
    if (sameQ == true) {
        for (unsigned int i=0; i < Q.n_rows; i++) {
            for (unsigned int j=0; j < Q.n_cols; j++) {
                (*eigval)(i, j) = lasteigval_simple(i, j);
                (*eigvec)(i, j) = lasteigvec_simple(i, j);
            }
        }
        return;
    }
    mat tQ(nstates,nstates); tQ.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            tQ(i, j) = Q(i, j);
        }
    }
    colvec eigva;
    mat eigve;
    eig_sym(eigva,eigve,tQ);
    //bool isImag = false; // not used
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i == j) {
                (*eigval)(i, j) = eigva(i);
                lasteigval_simple(i, j) = eigva(i);
            } else {
                (*eigval)(i, j) = 0;
                lasteigval_simple(i, j) = 0;
            }
            (*eigvec)(i, j) = eigve(i, j);
            lasteigvec_simple(i, j) = eigve(i, j);
        }
    }
    return;
}

/*
 * this should be used to caluculate the eigenvalues and eigenvectors
 * as U * Q * U-1 -- eigen decomposition
 *
 * this should use the armadillo library
 */
bool RateModel::get_eigenvec_eigenval_from_Q(cx_mat * eigval, cx_mat * eigvec) {
    if (sameQ == true) {
        for (unsigned int i=0; i < Q.n_rows; i++) {
            for (unsigned int j=0; j < Q.n_cols; j++) {
                (*eigval)(i, j) = lasteigval(i, j);
                (*eigvec)(i, j) = lasteigvec(i, j);
            }
        }
        return lastImag;
    }
    mat tQ(nstates,nstates); tQ.fill(0);
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            tQ(i, j) = Q(i, j);
        }
    }
    cx_colvec eigva;
    cx_mat eigve;
    eig_gen(eigva,eigve,tQ);
    bool isImag = false;
    for (unsigned int i=0; i < Q.n_rows; i++) {
        for (unsigned int j=0; j < Q.n_cols; j++) {
            if (i == j) {
                (*eigval)(i, j) = eigva(i);
                lasteigval(i, j) = eigva(i);
            } else {
                (*eigval)(i, j) = 0;
                lasteigval(i, j) = 0;
            }
            (*eigvec)(i, j) = eigve(i, j);
            lasteigvec(i, j) = eigve(i, j);
            if (imag((*eigvec)(i, j)) > 0 || imag((*eigval)(i, j))) {
                isImag = true;
            }
        }
    }
    //lasteigval = eigval;
    //lasteigvec = eigvec;
    lastImag = isImag;
    return isImag;
}

//taking out fortran
//
/*
extern"C" {
    void wrapalldmexpv_(int * n,int* m,double * t,double* v,double * w,double* tol,double* anorm,double* wsp,int * lwsp,int* iwsp,int *liwsp, int * itrace,int *iflag,int *ia, int *ja, double *a, int *nz, double * res);
    void wrapsingledmexpv_(int * n,int* m,double * t,double* v,double * w,double* tol,double* anorm,double* wsp,int * lwsp,int* iwsp,int *liwsp, int * itrace,int *iflag,int *ia, int *ja, double *a, int *nz, double * res);
    void wrapdgpadm_(int * ideg,int * m,double * t,double * H,int * ldh,double * wsp,int * lwsp,int * ipiv,int * iexph,int *ns,int *iflag );
}

void RateModel::setup_fortran_P_whichQ(int which, mat & P, double t) {
    Q = Qs[which];
    setup_fortran_P(P,t,false);
}
*/
//taking out fortran
/*
 * runs the basic padm fortran expokit full matrix exp
 *
void RateModel::setup_fortran_P(mat & P, double t, bool store_p_matrices) {
    //
    //  return P, the matrix of dist-to-dist transition probabilities,
    //  from the model's rate matrix (Q) over a time duration (t)
    //
    int ideg = 6;
    int m = Q.n_rows; // square so you only need the rows
    int ldh = m;
    double tol = 1;
    int iflag = 0;
    int lwsp = 4*m*m+6+1;
    double * wsp = new double[lwsp];
    int * ipiv = new int[m];
    int iexph = 0;
    int ns = 0;
    double * H = new double [m*m];
    convert_matrix_to_single_row_for_fortran(Q, t, H);
    wrapdgpadm_(&ideg, &m, &tol, H, &ldh, wsp, &lwsp, ipiv, &iexph, &ns, &iflag);

    for (int i=0; i < m; i++) {
        for (int j=0; j < m; j++) {
            P(i, j) = wsp[iexph + (j-1) * m + (i-1) + m];
        }
    }
    delete [] wsp;
    delete [] ipiv;
    delete [] H;
    for (int i=0; i < nstates; i++) {
        double sum = 0.0;
        for (int j=0; j < nstates; j++) {
            sum += P(i, j);
        }
        for (int j=0; j < nstates; j++) {
            P(i, j) = (P(i, j)/sum);
        }
    }

}
*/
void RateModel::set_sameQ(bool s) {
    sameQ = s;
}

void update_simple_goldman_yang_q(mat * inm, double K, double w, mat & bigpibf,mat &bigpiK, mat & bigpiw) {
    double s = 0;
    for (int i=0; i < 61; i++) {
        for (int j=0; j < 61; j++) {
            if (bigpibf(i, j) == 0) {
                (*inm)(i, j) = 0;
            } else {
                (*inm)(i, j) = 1/61.;
                if (bigpiK(i, j) != 0) {
                    (*inm)(i, j) *= K;
                }
                if (bigpiw(i, j) != 0) {
                    (*inm)(i, j) *= w;
                }
            }
            if (i == j) {
                (*inm)(i, j) = 0;
            }
            s += (*inm)(i, j);
        }
    }
    (*inm) = (*inm)/s;
    for (unsigned int i=0; i < (*inm).n_rows; i++) {
        double su = 0;
        for (unsigned int j=0; j < (*inm).n_cols; j++) {
            if (i != j) {
            su += (*inm)(i, j);
            }
        }
        (*inm)(i, i) = 0-su;
    }
    (*inm) = (*inm)/(1/61.);
    //(*inm) = trans((*inm));
}

bool test_transition(char a, char b) {
    bool ret = false;
    if ((a == 'A' &&  b == 'G') || (a == 'C' &&  b == 'T') || (a == 'G' &&  b == 'A') || (a == 'T' &&  b == 'C')) {
        ret = true;
    }
    return ret;
}

void generate_bigpibf_K_w(mat * bf, mat * K, mat * w,map<string, string> & codon_dict,
    map<string, vector<int> > & codon_index, vector<string> & codon_list) {
    for (int i=0; i < 61; i++) {
        for (int j=0; j < 61; j++) {
            int diff = 0;
            bool transit = false;
            bool nonsyn = false;
            for (int m=0; m < 3; m++) {
                if (codon_list[i][m] != codon_list[j][m]) {
                    diff += 1;
                }
                transit = test_transition(codon_list[i][m],codon_list[j][m]);
            }
            if (diff > 1) {
                (*bf)(i, j) = 0;
            } else {
                if (codon_dict[codon_list[i]] != codon_dict[codon_list[j]]) {
                    nonsyn = true;
                }
                (*bf)(i, j) = 1;
                if (transit) {
                    (*K)(i, j) = 1;
                }
                if (nonsyn) {
                    (*w)(i, j) = 1;
                }
            }
        }
    }
}

//take out fortran
/*
void convert_matrix_to_single_row_for_fortran(mat & inmatrix, double t, double * H) {
    int count = 0;
    for (unsigned int i=0; i < inmatrix.n_cols; i++) {
        for (unsigned int j=0; j < inmatrix.n_cols; j++) {
            H[i+(j*inmatrix.n_cols)] = inmatrix(i, j)*t;
            count += 1;
        }
    }
}
*/


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2012 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2012 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2012 Mateusz Loskot, London, UK.

// This file was modified by Oracle on 2014, 2016, 2017.
// Modifications copyright (c) 2014-2017 Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_CORE_SRS_HPP
#define BOOST_GEOMETRY_CORE_SRS_HPP

#include <cstddef>

#include <boost/static_assert.hpp>

#include <boost/geometry/core/radius.hpp>
#include <boost/geometry/core/tag.hpp>
#include <boost/geometry/core/tags.hpp>

namespace boost { namespace geometry
{

namespace srs
{

/*!
    \brief Defines spheroid radius values for use in geographical CS calculations
    \note See http://en.wikipedia.org/wiki/Figure_of_the_Earth
          and http://en.wikipedia.org/wiki/World_Geodetic_System#A_new_World_Geodetic_System:_WGS84
*/
template <typename RadiusType>
class spheroid
{
public:
    spheroid(RadiusType const& a, RadiusType const& b)
        : m_a(a)
        , m_b(b)
    {}

    spheroid()
        : m_a(RadiusType(6378137.0))
        , m_b(RadiusType(6356752.3142451793))
    {}

    template <std::size_t I>
    RadiusType get_radius() const
    {
        BOOST_STATIC_ASSERT(I < 3);

        return I < 2 ? m_a : m_b;
    }

    template <std::size_t I>
    void set_radius(RadiusType const& radius)
    {
        BOOST_STATIC_ASSERT(I < 3);

        (I < 2 ? m_a : m_b) = radius;
    }

private:
    RadiusType m_a, m_b; // equatorial radius, polar radius
};

} // namespace srs

// Traits specializations for spheroid
#ifndef DOXYGEN_NO_TRAITS_SPECIALIZATIONS
namespace traits
{

template <typename RadiusType>
struct tag< srs::spheroid<RadiusType> >
{
    typedef srs_spheroid_tag type;
};

template <typename RadiusType>
struct radius_type< srs::spheroid<RadiusType> >
{
    typedef RadiusType type;
};

template <typename RadiusType, std::size_t Dimension>
struct radius_access<srs::spheroid<RadiusType>, Dimension>
{
    typedef srs::spheroid<RadiusType> spheroid_type;

    static inline RadiusType get(spheroid_type const& s)
    {
        return s.template get_radius<Dimension>();
    }

    static inline void set(spheroid_type& s, RadiusType const& value)
    {
        s.template set_radius<Dimension>(value);
    }
};

} // namespace traits
#endif // DOXYGEN_NO_TRAITS_SPECIALIZATIONS

namespace srs
{

/*!
    \brief Defines sphere radius value for use in spherical CS calculations
*/
template <typename RadiusType>
class sphere
{
public:
    explicit sphere(RadiusType const& r)
        : m_r(r)
    {}

    sphere()
        : m_r(RadiusType((2.0 * 6378137.0 + 6356752.3142451793) / 3.0))
    {}

    template <std::size_t I>
    RadiusType get_radius() const
    {
        BOOST_STATIC_ASSERT(I < 3);

        return m_r;
    }

    template <std::size_t I>
    void set_radius(RadiusType const& radius)
    {
        BOOST_STATIC_ASSERT(I < 3);

        m_r = radius;
    }

private:
    RadiusType m_r; // radius
};

} // namespace srs

// Traits specializations for sphere
#ifndef DOXYGEN_NO_TRAITS_SPECIALIZATIONS
namespace traits
{

template <typename RadiusType>
struct tag< srs::sphere<RadiusType> >
{
    typedef srs_sphere_tag type;
};

template <typename RadiusType>
struct radius_type< srs::sphere<RadiusType> >
{
    typedef RadiusType type;
};

template <typename RadiusType, std::size_t Dimension>
struct radius_access<srs::sphere<RadiusType>, Dimension>
{
    typedef srs::sphere<RadiusType> sphere_type;

    static inline RadiusType get(sphere_type const& s)
    {
        return s.template get_radius<Dimension>();
    }

    static inline void set(sphere_type& s, RadiusType const& value)
    {
        s.template set_radius<Dimension>(value);
    }
};

} // namespace traits
#endif // DOXYGEN_NO_TRAITS_SPECIALIZATIONS

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_CORE_SRS_HPP


// Copyright 2019 AES WBC Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdarg.h>

#include <NTL/mat_GF2.h>

#include "aes_private.h"

namespace {

void err_quit(const char *fmt, ...) {
  va_list ap;
  char buf[1024];

  va_start(ap, fmt);
  vsprintf(buf, fmt, ap);
  strcat(buf, "\n");
  fputs(buf, stderr);
  fflush(stderr);
  va_end(ap);

  exit(1);
}

void read_key(const char *in, uint8_t* key, size_t size) {
  if (strlen(in) != size << 1)
    err_quit("Invalid key (should be a valid %d-bits hexadecimal string)",
        (size == 16) ? 128 : ((size == 24) ? 192 : 256));
  for (size_t i = 0; i < size; i++) {
    sscanf(in + i * 2, "%2hhx", key + i);
  }
}

template<typename T>
inline NTL::vec_GF2 from_scalar(T in);

template<>
inline NTL::vec_GF2 from_scalar(uint8_t in) {
  NTL::vec_GF2 result;
  result.SetLength(8);
  for (int i = 0; i < 8; i++) {
    result[7 - i] = ((in >> i) & 1);
  }
  return result;
}

template<>
inline NTL::vec_GF2 from_scalar(uint32_t in) {
  NTL::vec_GF2 result;
  result.SetLength(32);
  for (int i = 0; i < 32; i++) {
    result[31 - i] = ((in >> i) & 1);
  }
  return result;
}

template<typename T>
inline T to_scalar(const NTL::vec_GF2& in);

template<>
inline uint8_t to_scalar(const NTL::vec_GF2& in) {
  uint8_t result = 0;
  for (int i = 0; i < 2; i++) {
    long i0 = NTL::rep(in[i*4+0]), i1 = NTL::rep(in[i*4+1]),
         i2 = NTL::rep(in[i*4+2]), i3 = NTL::rep(in[i*4+3]);
    result = (result << 4) | (i0 << 3) | (i1 << 2) | (i2 << 1) | (i3 << 0);
  }
  return result;
}

template<>
inline uint32_t to_scalar(const NTL::vec_GF2& in) {
  uint32_t result = 0;
  for (int i = 0; i < 8; i++) {
    long i0 = NTL::rep(in[i*4+0]), i1 = NTL::rep(in[i*4+1]),
         i2 = NTL::rep(in[i*4+2]), i3 = NTL::rep(in[i*4+3]);
    result = (result << 4) | (i0 << 3) | (i1 << 2) | (i2 << 1) | (i3 << 0);
  }
  return result;
}

template<typename T>
inline T mul(const NTL::mat_GF2& mat, T x) {
  return to_scalar<T>(mat * from_scalar<T>(x));
}

NTL::mat_GF2 GenerateGF2RandomMatrix(int dimension) {
  NTL::mat_GF2 mat(NTL::INIT_SIZE, dimension, dimension);
  for (int i = 0; i < dimension; i++) {
    for (int j = 0; j < dimension; j++) {
      mat[i][j] = NTL::random_GF2();
    }
  }
  return mat;
}

NTL::mat_GF2 GenerateRandomGF2InvertibleMatrix(int dimension) {
  for (;;) {
    NTL::mat_GF2 result = GenerateGF2RandomMatrix(dimension);
    if (NTL::determinant(result) != 0)
      return result;
  }
}

// Calculate the T-boxes, which is a combination of the AddRoundKeyAfterShift
// and the SubBytes functions.
void CalculateTboxes(const uint32_t roundKey[],
    uint8_t Tboxes[][16][256], int Nr) {
  for (int r = 0; r < Nr; r++) {
    for (int x = 0; x < 256; x++) {
      uint8_t state[16] = {
        (uint8_t)x, (uint8_t)x, (uint8_t)x, (uint8_t)x,
        (uint8_t)x, (uint8_t)x, (uint8_t)x, (uint8_t)x,
        (uint8_t)x, (uint8_t)x, (uint8_t)x, (uint8_t)x,
        (uint8_t)x, (uint8_t)x, (uint8_t)x, (uint8_t)x
      };
      AddRoundKeyAfterShift(state, &roundKey[r*4]);
      SubBytes(state);
      if (r == Nr-1) {
        AddRoundKey(state, &roundKey[4*Nr]);
      }
      for (int i = 0; i < 16; i++) {
        Tboxes[r][i][x] = state[i];
      }
    }
  }
}

void CalculateTy(uint8_t Ty[4][256][4]) {
  for (int x = 0; x < 256; x++) {
    Ty[0][x][0] = gf_mul[x][0];
    Ty[0][x][1] = gf_mul[x][1];
    Ty[0][x][2] = x;
    Ty[0][x][3] = x;

    Ty[1][x][0] = x;
    Ty[1][x][1] = gf_mul[x][0];
    Ty[1][x][2] = gf_mul[x][1];
    Ty[1][x][3] = x;

    Ty[2][x][0] = x;
    Ty[2][x][1] = x;
    Ty[2][x][2] = gf_mul[x][0];
    Ty[2][x][3] = gf_mul[x][1];

    Ty[3][x][0] = gf_mul[x][1];
    Ty[3][x][1] = x;
    Ty[3][x][2] = x;
    Ty[3][x][3] = gf_mul[x][0];
  }
}

void CalculateTyBoxes(uint32_t roundKey[],
    uint32_t Tyboxes[][16][256], uint8_t TboxesLast[16][256],
    uint32_t MBL[][16][256], bool enableL, bool enableMB, int Nr) {
  uint8_t Tboxes[Nr][16][256];
  uint8_t Ty[4][256][4];

  CalculateTboxes(roundKey, Tboxes, Nr);
  CalculateTy(Ty);

  for (int r = 0; r < Nr-1; r++) {
    for (int x = 0; x < 256; x++) {
      for (int j = 0; j < 4; j++) {
        for (int i = 0; i < 4; i++) {
          uint32_t v0 = Ty[0][Tboxes[r][j*4 + i][x]][i],
                   v1 = Ty[1][Tboxes[r][j*4 + i][x]][i],
                   v2 = Ty[2][Tboxes[r][j*4 + i][x]][i],
                   v3 = Ty[3][Tboxes[r][j*4 + i][x]][i];
          Tyboxes[r][j*4 + i][x] = (v0 << 24) | (v1 << 16) | (v2 << 8) | v3;
          MBL[r][j*4 + i][x] = x << ((3 - i) << 3);
        }
      }
    }
  }

  for (int x = 0; x < 256; x++) {
    for (int i = 0; i < 16; i++) {
      TboxesLast[i][x] = Tboxes[Nr-1][i][x];
    }
  }

  if (enableMB) {
    NTL::mat_GF2 MB[Nr-1][4];
    for (int r = 0; r < Nr-1; r++) {
      for (int i = 0; i < 4; i++) {
        MB[r][i] = GenerateRandomGF2InvertibleMatrix(32);
      }
    }

    // When applying MB and inv(MB), the operation is quite easy; there is no
    // need to safeguard the existing table, as it is a simple substitution.
    for (int r = 0; r < Nr-1; r++) {
      for (int x = 0; x < 256; x++) {
        for (int i = 0; i < 16; i++) {
          Tyboxes[r][i][x] = mul<uint32_t>(MB[r][i >> 2], Tyboxes[r][i][x]);
          MBL[r][i][x] = mul<uint32_t>(NTL::inv(MB[r][i >> 2]), MBL[r][i][x]);
        }
      }
    }
  }

  if (enableL) {
    NTL::mat_GF2 L[Nr-1][16];
    for (int r = 0; r < Nr-1; r++) {
      for (int i = 0; i < 16; i++) {
        L[r][i] = GenerateRandomGF2InvertibleMatrix(8);
      }
    }

    // When applying L and inv(L), things get a little tricky. As it involves
    // non-linear substitutions, the original table has to be copied before
    // being updated.
    for (int r = 0; r < Nr-1; r++) {

      if (r > 0) {
        // Rounds 1 to Nr-1 are reversed here.
        for (int i = 0; i < 16; i++) {
          uint32_t oldTyboxes[256];
          for (int x = 0; x < 256; x++)
            oldTyboxes[x] = Tyboxes[r][i][x];
          for (int x = 0; x < 256; x++)
            Tyboxes[r][i][x] = oldTyboxes[mul<uint8_t>(NTL::inv(L[r-1][i]), x)];
        }
      }

      // Apply the L transformation at each round.
      for (int j = 0; j < 4; ++j) {
        for (int x = 0; x < 256; x++) {
          uint32_t out0 = MBL[r][j*4 + 0][x];
          uint32_t out1 = MBL[r][j*4 + 1][x];
          uint32_t out2 = MBL[r][j*4 + 2][x];
          uint32_t out3 = MBL[r][j*4 + 3][x];

          MBL[r][j*4 + 0][x] = (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 0]], out0 >> 24) << 24)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 1]], out0 >> 16) << 16)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 2]], out0 >>  8) <<  8)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 3]], out0 >>  0) <<  0);

          MBL[r][j*4 + 1][x] = (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 0]], out1 >> 24) << 24)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 1]], out1 >> 16) << 16)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 2]], out1 >>  8) <<  8)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 3]], out1 >>  0) <<  0);

          MBL[r][j*4 + 2][x] = (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 0]], out2 >> 24) << 24)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 1]], out2 >> 16) << 16)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 2]], out2 >>  8) <<  8)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 3]], out2 >>  0) <<  0);

          MBL[r][j*4 + 3][x] = (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 0]], out3 >> 24) << 24)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 1]], out3 >> 16) << 16)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 2]], out3 >>  8) <<  8)
                             | (mul<uint8_t>(L[r][InvShiftRowsTab[j*4 + 3]], out3 >>  0) <<  0);
        }
      }
    }

    // The last and final round 9 is reversed here.
    for (int i = 0; i < 16; i++) {
      uint8_t oldTboxesLast[256];
      for (int x = 0; x < 256; x++)
        oldTboxesLast[x] = TboxesLast[i][x];
      for (int x = 0; x < 256; x++)
        TboxesLast[i][x] = oldTboxesLast[mul<uint8_t>(NTL::inv(L[Nr-2][i]), x)];
    }
  }
}

void GenerateXorTable(FILE* out, int Nr) {
  uint8_t Xor[Nr-1][96][16][16];
  for (int r = 0; r < Nr-1; r++)
    for (int n = 0; n < 96; n++)
      for (int i = 0; i < 16; i++)
        for (int j = 0; j < 16; j++)
          Xor[r][n][i][j] = i ^ j;

  fprintf(out, "constexpr uint8_t Xor[%d][96][16][16] = {\n", Nr-1);
  for (int r = 0; r < Nr-1; r++) {
    fprintf(out, "  {\n");
    for (int n = 0; n < 96; n++) {
      fprintf(out, "    {\n");
      for (int i = 0; i < 16; i++) {
        fprintf(out, "      { ");
        for (int j = 0; j < 16; j++)
          fprintf(out, "0x%02x, ", Xor[r][n][i][j]);
        fprintf(out, "},\n");
      }
      fprintf(out, "    },\n");
    }
    fprintf(out, "  },\n");
  }
  fprintf(out, "};\n\n");
}

void GenerateEncryptingTables(FILE* out, uint32_t* roundKey, int Nr) {
  uint32_t Tyboxes[Nr-1][16][256];
  uint8_t TboxesLast[16][256];
  uint32_t MBL[Nr-1][16][256];

  CalculateTyBoxes(roundKey, Tyboxes, TboxesLast, MBL, true, true, Nr);

  fprintf(out, "constexpr uint32_t Tyboxes[%d][16][256] = {\n", Nr-1);
  for (int r = 0; r < Nr-1; r++) {
    fprintf(out, "  {\n");
    for (int i = 0; i < 16; i++) {
      fprintf(out, "    {\n");
      for (int x = 0; x < 256; x++) {
        if ((x % 8) == 0) {
          fprintf(out, "      ");
        }
        fprintf(out, "0x%08x,", Tyboxes[r][i][x]);
        if (x > 0 && (x % 8) == 7) {
          fprintf(out, "\n");
        } else {
          fprintf(out, " ");
        }
      }
      fprintf(out, "    },\n");
    }
    fprintf(out, "  },\n");
  }
  fprintf(out, "};\n\n");

  fprintf(out, "constexpr uint8_t TboxesLast[16][256] = {\n");
  for (int i = 0; i < 16; i++) {
    fprintf(out, "  {\n");
    for (int x = 0; x < 256; x++) {
      if (x % 16 == 0) {
        fprintf(out, "    ");
      }
      fprintf(out, "0x%02x, ", TboxesLast[i][x]);
      if (x % 16 == 15) {
        fprintf(out, "\n");
      }
    }
    fprintf(out, "  },\n");
  }
  fprintf(out, "};\n\n");

  fprintf(out, "constexpr uint32_t MBL[%d][16][256] = {\n", Nr-1);
  for (int r = 0; r < Nr-1; r++) {
    fprintf(out, "  {\n");
    for (int i = 0; i < 16; i++) {
      fprintf(out, "    {\n");
      for (int x = 0; x < 256; x++) {
        if ((x % 8) == 0) {
          fprintf(out, "      ");
        }
        fprintf(out, "0x%08x,", MBL[r][i][x]);
        if (x > 0 && (x % 8) == 7) {
          fprintf(out, "\n");
        } else {
          fprintf(out, " ");
        }
      }
      fprintf(out, "    },\n");
    }
    fprintf(out, "  },\n");
  }
  fprintf(out, "};\n\n");
}

void GenerateTables(const char* hexKey, int Nk, int Nr) {
  uint8_t key[Nk*4];
  uint32_t roundKey[(Nr+1)*4];

  read_key(hexKey, key, Nk*4);
  ExpandKeys(key, roundKey, Nk, Nr);

  FILE* out = fopen("aes_whitebox_tables.cc", "w");

  fprintf(out,
      "// This file is generated, do not edit.\n"
      "\n"
      "namespace {\n"
      "\n"
      "constexpr int Nr = %d;\n"
      "\n", Nr);

  GenerateXorTable(out, Nr);
  GenerateEncryptingTables(out, roundKey, Nr);

  fprintf(out, "}  // namespace");

  fflush(out);
  fclose(out);
}

void syntax() {
  err_quit("Syntax: aes_whitebox_gen <aes128|aes192|aes256> <hex-key>");
}

}  // namespace

int main(int argc, char* argv[]) {
  int Nk = 0, Nr = 0;

  if (argc != 3) {
    syntax();
  } else if (strcmp(argv[1], "aes128") == 0) {
    Nk = 4, Nr = 10;
  } else if (strcmp(argv[1], "aes192") == 0) {
    Nk = 6, Nr = 12;
  } else if (strcmp(argv[1], "aes256") == 0) {
    Nk = 8, Nr = 14;
  } else if (strcmp(argv[1], "aes512") == 0) {
    Nk = 16, Nr = 22;
  } else if (strcmp(argv[1], "aes1024") == 0) {
    Nk = 32, Nr = 38;
  } else if (strcmp(argv[1], "aes2048") == 0) {
    Nk = 64, Nr = 70;
  } else if (strcmp(argv[1], "aes4096") == 0) {
    Nk = 128, Nr = 134;
  } else {
    syntax();
  }

  GenerateTables(argv[2], Nk, Nr);
  return 0;
}