from dnv_bladed_results import ResultsApi, Run, Group, UsageExamples
import dnv_bladed_results as dnvres
import numpy as np
import psutil
import gc
import os

############   Bladed Results API: Memory management   ############

# Demonstrates how to manage memory effectively when processing large data sets.

# The Results API provides functions for managing memory, which may be useful when operating on large data sets, or when using 32-bit Python. 
# In most cases, 64-bit Python with sufficient RAM will not require these functions.
# 64-bit Python is recommended for improved performance and larger data handling.

# The available memory management options are demonstrated below, and may be used in conjunction with each other.

# Key points:
# - Use ResultsApi.clear_runs() to clear all cached data.
# - Use group.clear_buffer() and run.clear_group_buffers() to clear buffered group data.
# - Use variable.clear_data() to free variable data.

# See the Results API glossary available in the Bladed Knowledge Centre for definitions of Results API entities.

# IMPORTANT
# NumPy arrays returned by the Results API are non-owning views into internal C++ buffers.
# Calling ResultsApi.clear_runs() or destroying result objects invalidates all previously returned arrays.
# Accessing them afterwards results in undefined behavior.

# This example requires the psutil library, available via pip.  It has been tested with psutil >= 6.0.0.


def run_script():
    r"""
    Example script that runs when this Python module is called as the main function.
    
    Demonstrates how to manage memory by:
    - Clearing the run cache.
    - Clearing group buffers.
    - Clearing variable data arrays.
    """

    # Example showing how to disable/enable run caching and clear the run cache
    run_cache_management()

    # Example showing how to clear group buffers to reduce memory usage
    group_buffer_management()

    # Example showing how to clear variable data to reduce memory usage
    variable_1d_data_management()
    variable_2d_data_management()

    # Test demonstrating stable memory footprint when runs, groups, and variables are repeatedly cached and cleared
    memory_stability_test()

    print()


##########
#  Runs  #
##########

def run_cache_management():
    r"""
    By default, the Results API caches `Run` objects returned by `ResultsApi.get_run` and `ResultsApi.get_runs*`.
    This can be disabled via `ResultsApi.CacheSettings.run_caching_enabled = False`.

    Caching is generally recommended, as `Run` objects are small and caching reduces disk access.

    The cache can be cleared with `ResultsApi.clear_runs()`, which frees all memory associated with runs, including buffered groups and variables.
    """
    print("\n*** Run cache management ***\n")

    #########################
    #  Run caching enabled  #
    #########################

    ResultsApi.clear_runs()
    original_cache_setting = ResultsApi.CacheSettings.run_caching_enabled

    try:
        # Enable run caching
        ResultsApi.CacheSettings.run_caching_enabled = True
        print("Run caching is enabled.")

        # Get runs
        runs_root_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo")
        runs = ResultsApi.get_runs(runs_root_directory)

        print(f"Request returned {runs.size} runs.")

        assert runs.size > 0, "Expected some runs"
        assert ResultsApi._cached_runs_count == runs.size, "Expected cached runs when run caching is enabled"

        print(f"The run cache contains {ResultsApi._cached_runs_count} runs.")

        # Clear the cache
        ResultsApi.clear_runs()
        assert ResultsApi._cached_runs_count == 0, "Expected no cached runs following cache clearance"

        print("The run cache has been cleared and is now empty.")

        # Remove local references so Run objects can be destroyed
        del runs
        gc.collect()

        ##########################
        #  Run caching disabled  #
        ##########################

        # Disable run caching
        ResultsApi.CacheSettings.run_caching_enabled = False
        print("\nRun caching is disabled.")

        # Get demo runs again
        runs = ResultsApi.get_runs(runs_root_directory)

        print(f"Request returned {runs.size} runs.")

        assert runs.size > 0, "Expected some runs"
        assert ResultsApi._cached_runs_count == 0, "Expected no cached runs when run caching is disabled"

        print("The run cache is empty.")

        # Remove local references so Run objects can be destroyed
        del runs
        gc.collect()

    finally:
        # Revert to default run caching setting
        ResultsApi.CacheSettings.run_caching_enabled = original_cache_setting


############
#  Groups  #
############

def group_buffer_management():
    r"""
    When a variable's data is requested by calling `variable.get_data*`, its group file (.$nn) is loaded into memory and retained for the rest of the session.
    This speeds up subsequent access to variables in the same group by reducing disk reads which can be slow (espcially over a network).
    
    Group files can be very large - hundreds of megabytes or more - depending on the calculation.

    Individual group buffers can be cleared with `group.clear_buffer()`, as shown in this example.
    When not using groups and requesting variables straight from runs, it is possible to clear all underlying group buffers in a run with `run.clear_group_buffers()`.

    Buffering can be disabled via `ResultsApi.CacheSettings.buffered_read_enabled = False`, but it is generally recommended to keep it enabled and manually clear buffers
    when they are no longer needed.

    Note: All variables are unloaded and their data cleared by `ResultsApi.clear_runs()`.
    """
    print("\n*** Group buffer management ***\n")

    ResultsApi.clear_runs()

    run_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo/powprod5MW")
    run = ResultsApi.get_run(run_directory, "powprod5MW")
    total_buffer_clearance_kb = 0

    # Get all groups in run
    groups = run.get_groups()
    group: Group
    for group in groups:
        # Load the first variable in each group
        first_variable_name_in_group = group.get_variable_names()[0]
        if (group.is_one_dimensional):
            variable = group.get_variable_1d(first_variable_name_in_group)
        else:
            variable = group.get_variable_2d(first_variable_name_in_group)
        
        # A variable has been returned, but no data is yet in memory
        assert not group._is_buffered, "Expected non-buffered group"
        assert not variable._has_cached_data, "Expected unpopulated variable data array"

        # Get data
        data = variable.get_data()
        print(f"Loaded variable '{variable.name}' from group '{group.name}'.")

        # The group buffer is in memory and the variable data array has been populated
        assert group._is_buffered, "Expected buffered group"
        assert variable._has_cached_data, "Expected populated variable data array"

        buffer_size_kb = round(group._buffer_size / 1024, 2)

        # We do not intend to read any other variables from the current group - clear its buffer to free memory
        group.clear_buffer()

        # Important: It is not necessary to 

        # Group buffer has been cleared to save memory
        assert not group._is_buffered, "Expected empty group buffer"
        
        # Group buffer clearance does not affect existing variable data, which is still usable:
        assert variable._has_cached_data, "Expected populated variable data array"
        assert len(data) > 0, "Expected populated variable data array"

        print(f"Group buffer has been cleared, freeing {buffer_size_kb} KB.\n")
        total_buffer_clearance_kb = total_buffer_clearance_kb + buffer_size_kb

    print(f"Total memory saving: {round(total_buffer_clearance_kb, 2)} KB.")


###################
#  Variable data  #
###################

def variable_1d_data_management():
    r"""
    Data returned by `variable.get_data()` is kept in memory until cleared with `variable.clear_data()`.
    Clearing removes only the data; metadata (e.g. name, statistics) remains available.

    Clearing helps to reduce memory usage when processing many variables in sequence.

    Note: All variables are unloaded and their data cleared by `ResultsApi.clear_runs()`.
    """
    print("\n*** 1D variable data management ***\n")

    ResultsApi.clear_runs()

    run_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo/powprod5MW")
    run = ResultsApi.get_run(run_directory, "powprod5MW")

    # 1D Variable - get all data
    variable_1d = run.get_variable_1d("brake speed")
    assert not variable_1d._has_cached_data, "Expected unpopulated 1D variable data array"

    # Populate 1D variable data
    data_1d = variable_1d.get_data()
    assert variable_1d._has_cached_data, "Expected populated 1D variable data array"

    print(f"Loaded 1D variable '{variable_1d.name}'.")

    # Get data size
    bytes_per_value = 4
    if (variable_1d.data_type == dnvres.VARIABLE_DATA_TYPE_FLOAT64):
        bytes_per_value = 8
    data_size_kb = round((variable_1d.data_point_count * bytes_per_value) / 1024, 2)

    # Clear 1D variable data
    # Important: the following call invalidates the local data_1d variable, which is a non-owning view of the underlying storage
    variable_1d.clear_data()
    
    # It is recommended to delete the local reference to avoid accidental use after free
    del data_1d

    assert not variable_1d._has_cached_data, "Expected unpopulated 1D variable data array"
    print(f"Variable data array has been cleared, saving {data_size_kb} KB.")


def variable_2d_data_management():
    r"""
    Data returned by `variable.get_data*` is kept in memory until cleared with `variable.clear_data()`.
    Clearing removes only the data; metadata (e.g. name, statistics) remains available.

    Clearing helps to reduce memory usage when processing many variables in sequence.

    Note: All variables are unloaded and their data cleared by `ResultsApi.clear_runs()`.
    """
    print("\n*** 2D variable data management ***\n")

    ResultsApi.clear_runs()

    run_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo/powprod5MW")
    run = ResultsApi.get_run(run_directory, "powprod5MW")

    # 2D Variable - get all data
    variable_2d = run.get_variable_2d("blade 1 mx (principal elastic axes)")
    assert not variable_2d._has_cached_data, "Expected unpopulated 2D variable data array"

    # Populate 2D variable data
    # For large datasets, the functions get_data_at_index and get_data_at_value are recommended to reduce memory consumption
    data_2d = variable_2d.get_data()
    assert variable_2d._has_cached_data, "Expected populated 2D variable data array"

    print(f"Loaded 2D variable '{variable_2d.name}' - all data.")

    # Get data size
    bytes_per_value = 4
    if (variable_2d.data_type == dnvres.VARIABLE_DATA_TYPE_FLOAT64):
        bytes_per_value = 8

    secondary_ind_var = variable_2d.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_SECONDARY)
    num_secondary_axis_values = secondary_ind_var.number_of_values
    data_size_kb = round((variable_2d.data_point_count * num_secondary_axis_values * bytes_per_value) / 1024, 2)

    # Clear 2D variable data
    variable_2d.clear_data()

    # It is recommended to delete the local reference to avoid accidental use after free
    del data_2d

    assert not variable_2d._has_cached_data, "Expected unpopulated 2D variable data array."
    print(f"Variable data array has been cleared, saving {data_size_kb} KB.")


    # 2D Variable - get data for individual secondary axis values
    secondary_axis_values = secondary_ind_var.get_values_as_number()
    for i in range(num_secondary_axis_values):
        assert not variable_2d._has_cached_data, "Expected unpopulated 2D variable data array"
        
        # Populate 2D variable data
        data_2d_subset = variable_2d.get_data_at_index(i)
        assert variable_2d._has_cached_data, "Expected populated 2D variable data array"
        
        print(f"\nLoaded 2D variable '{variable_2d.name}' @ {secondary_axis_values[i]}{secondary_ind_var.si_unit}")

        # Get data size
        data_size_kb = round((variable_2d.data_point_count * bytes_per_value) / 1024, 2)

        # Clear 2D variable data
        # Important: the following call invalidates the local data_2d_subset variable, which is a non-owning view of the underlying storage
        variable_2d.clear_data()

        # It is recommended to delete the local reference to avoid accidental use after free
        del data_2d_subset

        assert not variable_2d._has_cached_data, "Expected unpopulated 2D variable data array"
        print(f"Variable data array has been cleared, saving {data_size_kb} KB.")


######################
#  Memory stability  #
######################

def memory_stability_test():
    r"""
    Demonstrates memory stability when repeatedly caching and clearing runs, groups, and variables.
    """
    print("\n*** Memory stability test ***\n")
    
    ResultsApi.clear_runs()

    runs_root_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo")
    process = psutil.Process()

    for i in range(20):

        # Get demo runs
        runs = ResultsApi.get_runs(runs_root_directory)

        # Iterate the returned runs, getting all groups, all variables from each group, and all data for each variable
        numGroups = 0
        numVariables = 0
        run: Run
        for run in runs:
            if (run.is_calculation_type_supported):
                groups = run.get_groups()
                numGroups = numGroups + groups.size
                group: Group
                for group in groups:
                    variable_names = group.get_variable_names()
                    for variable_name in variable_names:
                        if (group.is_one_dimensional):
                            variable = group.get_variable_1d(variable_name)
                            variable.get_data()
                        else:
                            variable = group.get_variable_2d(variable_name)
                            variable.get_data()
                        numVariables = numVariables + 1

        print(f"Iteration {i + 1}: found {runs.size} runs; {numGroups} groups; {numVariables} variables")
        assert ResultsApi._cached_runs_count == 18, "Expected number of cached runs"

        # Record memory consumption before clearing runs
        memory_before_clear = round(process.memory_info().rss / 1024 ** 2, 2)
        print(f"Process memory before clear: {memory_before_clear} MB")

        # Track peak memory footprint over time
        if (i == 0):
            min_peak_memory = memory_before_clear
            max_peak_memory = memory_before_clear
        else:
            if (memory_before_clear < min_peak_memory):
                min_peak_memory = memory_before_clear
            if (memory_before_clear > max_peak_memory):
                max_peak_memory = memory_before_clear

        # Clear runs - this also clears group buffers and variables
        ResultsApi.clear_runs()

        assert ResultsApi._cached_runs_count == 0, "Expected no cached runs following cache clearance"

        # Explicitly remove references to loaded runs and trigger garbage collection (loop variables remain bound after the loop exits)
        del runs
        del run
        del group
        del variable
        gc.collect()

        # Record memory consumption after clearing runs
        memory_after_clear = round(process.memory_info().rss / 1024 ** 2, 2)
        print(f"Process memory after clear : {memory_after_clear} MB")

        # Typically we expect the allocator return memory to the OS, but it may pool it internally for future use (hence <= comparison)
        assert memory_after_clear <= memory_before_clear, "Expected cache clearance to free memory"

        print()

    # Check that peak memory consumption does not exceed the expected limit
    memory_delta_mb = max_peak_memory - min_peak_memory
    tolerance = 5.0
    assert memory_delta_mb < tolerance, f"Peak memory footprint increased by more than {tolerance} MB which is not expected"

    print("Memory footprint after clear and peak memory footprint remained within expected limits.")
    

if __name__ == "__main__":
    run_script()