/*
 * 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

namespace __nvhpc_std {

/**
 * Equivalent of lacpy in LAPACK, copies the upper/lower triangular part of a matrix
 * This is only for internal uses (by matrix_rank_k_update, for example). It's not part of the linear algebra APIs.
 * Fallback is not supported here.
 */
template< class _exec_space
        , class _ElementType_A
        , class _SizeType_A
        , ::std::size_t ... _ext_A
        , class _Layout_A
        , class _Accessor_A
        , class _ElementType_B
        , class _SizeType_B
        , ::std::size_t ... _ext_B
        , class _Layout_B
        , class _Accessor_B
        >
void copy_matrix( __nvhpc_exec< _exec_space > __exec
                , const char __uplo
                , std::mdspan< _ElementType_A, std::extents<_SizeType_A, _ext_A ...>, _Layout_A, _Accessor_A > __A
                , std::mdspan< _ElementType_B, std::extents<_SizeType_B, _ext_B ...>, _Layout_B, _Accessor_B > __B
                )
{
    static_assert (__A.rank() == 2);
    static_assert (__B.rank() == 2);

    constexpr bool __types_supported = __data_types_supported( __stdblas_data_type<_ElementType_A>, __stdblas_output_data_type<_ElementType_B> );
    constexpr bool     __A_supported = __input_supported < _Layout_A, _Accessor_A >();
    constexpr bool     __B_supported = __output_supported< _Layout_B, _Accessor_B >();

    __STDBLAS_STATIC_ASSERT_TYPES ( __types_supported      );
    __STDBLAS_STATIC_ASSERT_INPUT (     __A_supported, __A );
    __STDBLAS_STATIC_ASSERT_OUTPUT(     __B_supported, __B );

    using __A_t = decltype( __A );

    bool const __is_transposed = !std::linalg::__is_column_major( __B );
    char   __op_A = 'N';
    bool __conj_A = false;

    __extract_ops( __A, __is_transposed, &__op_A, &__conj_A );
    bool const __transp_A = ( __op_A != 'N' );

    constexpr bool __is_scaled_A = __is_scaled< __A_t >();

    if ( __transp_A || __conj_A || __is_scaled_A )
    { // Need to create a work array to store the input matrix elements with in-place operations applied
        using __valA_t = typename __A_t::value_type;
        using __ptrA_t = typename std::unique_ptr<__valA_t, std::function<void(__valA_t*)>>;
        char const __my_uplo  = ( !__transp_A ? __uplo : ( __uplo == 'U' ? 'L' : 'U' ) );

        std::tuple< __ptrA_t, __A_t > __work_A{ __ptrA_t(), __A };
        if ( __transp_A )
        {
           __work_A = __create_transpose( __exec, __op_A, __A );

           if constexpr ( __is_scaled_A )
           {
               __valA_t __alpha = __valA_t{ __extract_scaling_factor( __A ) };
               __scale_impl( __exec, __alpha, std::get<1>(__work_A) );
           }
        }
        else
        {
            __work_A = __create_array_like( __exec, __A );
            __copy_impl( __exec, __A, std::get<1>(__work_A) );
        }
        __copy_matrix_impl( __exec, __my_uplo, std::get<1>(__work_A), __B );
    }
    else
    { // Directly call matrix_copy with the input matrix
        __copy_matrix_impl( __exec, __uplo, __A, __B );
    }
}

}
