# *****************************************************************************

# Copyright (c) 2025 AISS Group at Harbin Institute of Technology. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

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

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# *****************************************************************************



import gc

import time



import asnumpy as ap

import numpy as np

from loguru import logger



from utils import calculate_stable_metric, create_arrays





def bench_add(add_func, m1, m2, warmup: int, iterations: int, is_npu: bool = False) -> list:

    """

    Benchmark function for add operation.

    Memory-optimized: explicitly delete intermediate variables to avoid OOM.

    """

    # 1. Warmup phase

    for _ in range(warmup):

        res = add_func(m1, m2)

        del res  # Release immediately



    # 2. Benchmark phase

    times = []

    for _ in range(iterations):

        start = time.perf_counter()



        # Execute computation

        res = add_func(m1, m2)

        end = time.perf_counter()

        times.append(end - start)



        # Key optimization: explicitly delete result objects to prevent memory accumulation

        del res



    # Force garbage collection

    if is_npu:

        gc.collect()



    return times





def run_test_case(

    shape: tuple[int, ...],

    dtype: np.dtype = np.float32,

    warmup: int = 40,

    iterations: int = 400,

) -> dict[str, float]:

    """Run a single test case"""

    logger.info(f"{'=' * 50}")

    logger.info(f"Test shape: {shape}")



    m1_asnp, m2_asnp, m1_np, m2_np = create_arrays(shape, dtype)



    try:

        # --- Benchmark AsNumpy ---

        asnp_times = bench_add(

            ap.add,

            m1_asnp, m2_asnp,

            warmup=warmup,

            iterations=iterations,

            is_npu=True

        )



        # --- Benchmark NumPy ---

        np_times = bench_add(

            np.add,

            m1_np, m2_np,

            warmup=warmup,

            iterations=iterations,

            is_npu=False

        )



        # Calculate statistics

        metric_asnp = calculate_stable_metric(asnp_times)

        metric_np = calculate_stable_metric(np_times)

        speedup = metric_np / metric_asnp if metric_asnp > 0 else 0



        # Verify result consistency

        result_asnp = ap.add(m1_asnp, m2_asnp).to_numpy()

        result_np = np.add(m1_np, m2_np)



        max_diff = np.max(np.abs(result_asnp - result_np))

        max_val = np.max(np.abs(result_np))

        rel_diff = max_diff / max_val if max_val > 0 else max_diff



        if rel_diff < 1e-4:

            logger.info(f"Verification passed: results are consistent (relative diff: {rel_diff:.2e})")

        else:

            logger.warning(f"Results differ (max relative diff: {rel_diff:.2e})")



        return {

            'shape': shape,

            'asnumpy_metric': metric_asnp,

            'numpy_metric': metric_np,

            'speedup': speedup,

            'relative_diff': rel_diff

        }



    finally:

        # Explicitly clean up large objects for current shape

        del m1_asnp, m2_asnp, m1_np, m2_np

        gc.collect()





if __name__ == "__main__":

    logger.info("=" * 70)

    logger.info("README example code performance benchmark")

    logger.info("Test operation: add (element-wise addition)")

    logger.info("Statistics strategy: after warmup, take mid-segment fastest speed (exclude slowest 10%)")

    logger.info("=" * 70)



    # Test configuration

    shapes = [

        (500, 500),            # Medium scale

        (1000, 1000),          # Large scale test

        (2000, 2000),          # Larger scale test

        (3000, 3000),          # Extra large scale test

    ]

    dtype = np.dtype(np.float32)



    # Adjusted parameters: reduced iterations to fit NPU memory limits

    warmup_iterations = 40

    test_iterations = 400



    logger.info("\nConfiguration:")

    logger.info(f"  Data type: {dtype}")

    logger.info(f"  Warmup iterations: {warmup_iterations}")

    logger.info(f"  Test iterations: {test_iterations}")

    logger.info("  Statistics method: sort, exclude slowest 10%, take minimum")

    logger.info(f"\n{'=' * 70}\n")



    results = []

    for shape in shapes:

        try:

            result = run_test_case(shape, dtype, warmup_iterations, test_iterations)

            results.append(result)

        except Exception as e:

            logger.error(f"Test failed: {e}")

            import traceback

            traceback.print_exc()



    # Output summary results

    print(f"{'Shape':<15} | {'Data Size':<12} | {'AsNumpy':<12} | {'NumPy':<12} | {'Speedup':<10}")

    print(f"{'':15} | {'':12} | {'(ms)':<12} | {'(ms)':<12} | {'':10}")

    print("-" * 85)



    for result in results:

        shape_str = str(result['shape'])

        data_size = np.prod(result['shape'])

        data_size_str = f"{data_size:,}"

        asnp_time = f"{result['asnumpy_metric'] * 1000:.4f}"

        np_time = f"{result['numpy_metric'] * 1000:.4f}"

        speedup_str = f"{result['speedup']:.2f}x"



        print(f"{shape_str:<15} | {data_size_str:<12} | {asnp_time:<12} | {np_time:<12} | {speedup_str}")



    print("-" * 85)



    # Statistics

    if results:

        avg_speedup = sum(r['speedup'] for r in results) / len(results)

        max_speedup = max(r['speedup'] for r in results)



        logger.info(f"\n{'=' * 70}")

        logger.info("Performance statistics:")

        logger.info(f"  Average speedup: {avg_speedup:.2f}x")

        logger.info(f"  Maximum speedup: {max_speedup:.2f}x")

        logger.info(f"{'=' * 70}")



    logger.info("\nNotes:")

    logger.info("  • Using float32 data type for NPU compatibility")

    logger.info("  • Iterations adjusted to ensure memory safety")

    logger.info("  • Uses 'mid-segment fastest speed' algorithm to showcase NPU's true compute capability")

    logger.info("=" * 70)