/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "tfest-precomp.h"	// Precompiled headers
//
#include <mrpt/config.h>

#include "se2_l2_internal.h"

#if MRPT_ARCH_INTEL_COMPATIBLE

#include <mrpt/core/SSE_types.h>
#include <mrpt/core/exceptions.h>

using namespace mrpt::tfest;

internal::se2_l2_impl_return_t<float> internal::se2_l2_impl_SSE2(
	const TMatchingPairList& in_correspondences)
{
	// SSE vectorized version:
	const size_t N = in_correspondences.size();
	ASSERT_(N >= 2);
	const float N_inv = 1.0f / N;  // For efficiency, keep this value.

	// Ensure correct types:
	static_assert(sizeof(TMatchingPair::global.x) == sizeof(float));
	static_assert(sizeof(TMatchingPair::local.x) == sizeof(float));

	internal::se2_l2_impl_return_t<float> ret;

	__m128 sum_a_xyz = _mm_setzero_ps();  // All 4 zeros (0.0f)
	__m128 sum_b_xyz = _mm_setzero_ps();  // All 4 zeros (0.0f)

	//   [ f0     f1      f2      f3  ]
	//    xa*xb  ya*yb   xa*yb  xb*ya
	__m128 sum_ab_xyz = _mm_setzero_ps();  // All 4 zeros (0.0f)

	for (const auto& in_correspondence : in_correspondences)
	{
		// Get the pair of points in the correspondence:
		//   a_xyyx = [   xa     ay   |   xa    ya ]
		//   b_xyyx = [   xb     yb   |   yb    xb ]
		//      (product)
		//            [  xa*xb  ya*yb   xa*yb  xb*ya
		//                LO0    LO1     HI2    HI3
		// Note: _MM_SHUFFLE(hi3,hi2,lo1,lo0)
		const __m128 a_xyz =
			_mm_loadu_ps(&in_correspondence.global.x);	// *Unaligned* load
		const __m128 b_xyz =
			_mm_loadu_ps(&in_correspondence.local.x);  // *Unaligned* load

		const auto a_xyxy =
			_mm_shuffle_ps(a_xyz, a_xyz, _MM_SHUFFLE(1, 0, 1, 0));
		const auto b_xyyx =
			_mm_shuffle_ps(b_xyz, b_xyz, _MM_SHUFFLE(0, 1, 1, 0));

		// Compute the terms:
		sum_a_xyz = _mm_add_ps(sum_a_xyz, a_xyz);
		sum_b_xyz = _mm_add_ps(sum_b_xyz, b_xyz);

		//   [ f0     f1      f2      f3  ]
		//    xa*xb  ya*yb   xa*yb  xb*ya
		sum_ab_xyz = _mm_add_ps(sum_ab_xyz, _mm_mul_ps(a_xyxy, b_xyyx));
	}

	alignas(MRPT_MAX_STATIC_ALIGN_BYTES) float sums_a[4], sums_b[4];
	_mm_store_ps(sums_a, sum_a_xyz);
	_mm_store_ps(sums_b, sum_b_xyz);

	float SumXa = sums_a[0];
	float SumYa = sums_a[1];
	float SumXb = sums_b[0];
	float SumYb = sums_b[1];

	// Compute all four means:
	const __m128 Ninv_4val =
		_mm_set1_ps(N_inv);	 // load 4 copies of the same value
	sum_a_xyz = _mm_mul_ps(sum_a_xyz, Ninv_4val);
	sum_b_xyz = _mm_mul_ps(sum_b_xyz, Ninv_4val);

	// means_a[0]: mean_x_a
	// means_a[1]: mean_y_a
	// means_b[0]: mean_x_b
	// means_b[1]: mean_y_b
	alignas(MRPT_MAX_STATIC_ALIGN_BYTES) float means_a[4], means_b[4];
	_mm_store_ps(means_a, sum_a_xyz);
	_mm_store_ps(means_b, sum_b_xyz);

	ret.mean_x_a = means_a[0];
	ret.mean_y_a = means_a[1];
	ret.mean_x_b = means_b[0];
	ret.mean_y_b = means_b[1];

	//      Sxx   Syy     Sxy    Syx
	//    xa*xb  ya*yb   xa*yb  xb*ya
	alignas(MRPT_MAX_STATIC_ALIGN_BYTES) float cross_sums[4];
	_mm_store_ps(cross_sums, sum_ab_xyz);

	float Sxx = cross_sums[0];
	float Syy = cross_sums[1];
	float Sxy = cross_sums[2];
	float Syx = cross_sums[3];

	// Auxiliary variables Ax,Ay:
	ret.Ax = N * (Sxx + Syy) - SumXa * SumXb - SumYa * SumYb;
	ret.Ay = SumXa * SumYb + N * (Syx - Sxy) - SumXb * SumYa;

	return ret;
}

#endif	// MRPT_ARCH_INTEL_COMPATIBLE
