/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#pragma once

#include <functional>

namespace __nvhpc_std {

/**
 * Implementation to be called by both overwriting overloads and updating overloads
 * `__is_zero_out_C` sets if C will be zeroed out before calling the back-end or not
 */
template< class _exec_space
        , class _Scalar
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t _numRows_A
        , ::std::size_t _numCols_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_C
        , class _SizeType_C
        , ::std::size_t _numRows_C
        , ::std::size_t _numCols_C
        , class _Layout_C
        , class _Accessor_C
        , class _Triangle
        >
void __symmetric_matrix_rank_k_update_impl( __nvhpc_exec<_exec_space> __exec
                                          , _Scalar __alpha
                                          , std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A> __A
                                          , std::mdspan<_ElementType_C, std::extents<_SizeType_C, _numRows_C, _numCols_C>, _Layout_C, _Accessor_C> __C
                                          , _Triangle __t
                                          , bool __is_zero_out_C
                                          )
{
    constexpr __types_supported = __data_types_supported( __stdblas_data_type<_ElementType_A>, __stdblas_output_data_type<_ElementType_C> );
    constexpr     __A_supported = __input_supported <_Layout_A, _Accessor_A>();
    constexpr     __C_supported = __output_supported<_Layout_C, _Accessor_C>();

#ifndef STDBLAS_FALLBACK_UNSUPPORTED_CASES
    __STDBLAS_STATIC_ASSERT_TYPES ( __types_supported      );
    __STDBLAS_STATIC_ASSERT_INPUT (     __A_supported, __A );
    __STDBLAS_STATIC_ASSERT_OUTPUT(     __C_supported, __C );
#endif

    if constexpr ( __types_supported && __A_supported && __C_supported )
    {
        using    __A_t = typename std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A>;
        using __valA_t = typename __A_t::value_type;
        using __ptrA_t = typename std::unique_ptr<__valA_t, std::function<void(__valA_t*)>>;

        char   __op_A;
        bool __conj_A;
        std::tuple< __ptrA_t, __A_t > __work_A{ __ptrA_t(), __A };

        __extract_ops_no_C( __A, false /* __is_operate_on_transposed */, &__op_A, &__conj_A );

        // TODO check for precision loss
        _Scalar const __my_alpha( __extract_scaling_factor( __A ) );

        if constexpr ( std::linalg::is_complex<__valA_t>::value ) { if ( __conj_A ) { __work_A = __create_conjugate( __exec, __A ); } }

        __symmetric_matrix_rank_k_update_impl( __exec
                                             , _ElementType_C{ __alpha * __my_alpha * __my_alpha }
                                             , _ElementType_C{ __is_zero_out_C ? 0. : 1. }
                                             , ( __conj_A ? std::get<1>( __work_A ) : __A )
                                             , __C
                                             , __op_A
                                             , std::linalg::__get_fill_mode( __C, __t )
                                             );
    }
    else
    {
#ifdef STDBLAS_VERBOSE
        __STDBLAS_COMPILE_TIME_FALLBACK_MESSAGE( symmetric_matrix_rank_k_update );
#endif
        std::linalg::symmetric_matrix_rank_k_update( std::execution::seq, __alpha, __A, __C, __t );
    }
}

/**
 * Overwriting overload: C := alpha * A * A^T
 */
_NVHPC_MDSPAN_TEMPLATE_REQUIRES(
          class _exec_space
        , class _Scalar
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t _numRows_A
        , ::std::size_t _numCols_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_C
        , class _SizeType_C
        , ::std::size_t _numRows_C
        , ::std::size_t _numCols_C
        , class _Layout_C
        , class _Accessor_C
        , class _Triangle
        , /* requires */
        ( !std::linalg::impl::is_mdspan_v< _Scalar > )
)
void symmetric_matrix_rank_k_update( __nvhpc_exec<_exec_space> __exec
                                   , _Scalar __alpha
                                   , std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A> __A
                                   , std::mdspan<_ElementType_C, std::extents<_SizeType_C, _numRows_C, _numCols_C>, _Layout_C, _Accessor_C> __C
                                   , _Triangle __t
                                   )
{
    __symmetric_matrix_rank_k_update_impl( __exec, __alpha, __A, __C, __t, true /* zero out C */ );
}

/**
 * Overwriting overload w/o alpha: C := A * A^T
 */
template< class _exec_space
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t _numRows_A
        , ::std::size_t _numCols_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_C
        , class _SizeType_C
        , ::std::size_t _numRows_C
        , ::std::size_t _numCols_C
        , class _Layout_C
        , class _Accessor_C
        , class _Triangle
        >
void symmetric_matrix_rank_k_update( __nvhpc_exec<_exec_space> __exec
                                   , std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A> __A
                                   , std::mdspan<_ElementType_C, std::extents<_SizeType_C, _numRows_C, _numCols_C>, _Layout_C, _Accessor_C> __C
                                   , _Triangle __t
                                   )
{
    __symmetric_matrix_rank_k_update_impl( __exec, _ElementType_A(1.), __A, __C, __t, true /* zero out C */ );
}

/**
 * Updating overload: C := E + alpha * A * A^T
 */
template< class _exec_space
        , class _Scalar
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t _numRows_A
        , ::std::size_t _numCols_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_E
        , class _Extents_E
        , class _Layout_E
        , class _Accessor_E
        , class _ElementType_C
        , class _SizeType_C
        , ::std::size_t _numRows_C
        , ::std::size_t _numCols_C
        , class _Layout_C
        , class _Accessor_C
        , class _Triangle
        >
void symmetric_matrix_rank_k_update( __nvhpc_exec<_exec_space> __exec
                                   , _Scalar __alpha
                                   , std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A> __A
                                   , std::mdspan<_ElementType_E, _Extents_E                                        , _Layout_E, _Accessor_E> __E
                                   , std::mdspan<_ElementType_C, std::extents<_SizeType_C, _numRows_C, _numCols_C>, _Layout_C, _Accessor_C> __C
                                   , _Triangle __t
                                   )
{
    // Copy upper/lower triangle of E to C
    copy_matrix( __exec, std::linalg::__get_fill_mode( __E, __t ), __E, __C );

    __symmetric_matrix_rank_k_update_impl( __exec, __alpha, __A, __C, __t, false /* zero out C */ );
}

/**
 * Updating overload without alpha: C := E + A * A^T
 */
template< class _exec_space
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t _numRows_A
        , ::std::size_t _numCols_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_E
        , class _Extents_E
        , class _Layout_E
        , class _Accessor_E
        , class _ElementType_C
        , class _SizeType_C
        , ::std::size_t _numRows_C
        , ::std::size_t _numCols_C
        , class _Layout_C
        , class _Accessor_C
        , class _Triangle
        >
void symmetric_matrix_rank_k_update( __nvhpc_exec<_exec_space> __exec
                                   , std::mdspan<_ElementType_A, std::extents<_SizeType_A, _numRows_A, _numCols_A>, _Layout_A, _Accessor_A> __A
                                   , std::mdspan<_ElementType_E, _Extents_E                                        , _Layout_E, _Accessor_E> __E
                                   , std::mdspan<_ElementType_C, std::extents<_SizeType_C, _numRows_C, _numCols_C>, _Layout_C, _Accessor_C> __C
                                   , _Triangle __t
                                   )
{
    // Copy upper/lower triangle of E to C
    copy_matrix( __exec, std::linalg::__get_fill_mode( __E, __t ), __E, __C );

    __symmetric_matrix_rank_k_update_impl( __exec, _ElementType_A(1.), __A, __C, __t, false /* zero out C */ );
}

}
