from dnv_bladed_results import ResultsApi, Group, IndependentVariable, OutputGroup1D_Float32, OutputGroup2D_Float32, UsageExamples
import dnv_bladed_results as dnvres
import numpy as np
import os
import shutil
from pathlib import Path

############   Bladed Results API: Create filtered output group   ############

# Demonstrates how to overwrite an existing group from a subset of the group's variables.
# This can reduce group size by retaining only the required variables.


def run_script():
    r"""
    Example script that runs when this Python module is called as the main function.
    
    - Copies existing 1D (hub loads) and 2D (blade loads, tower loads) groups from example run.
    - Reads a subset of variables from each group.
    - Creates new output groups using variable subset and overwrites the copied groups.
    - Performs roundtrip validation to check the written groups contain the expected values.
    """

    output_directory = "./Output"
    source_run_directory = os.path.join(UsageExamples.__path__[0], "Runs/demo/powprod5MW")
    source_run_name = "powprod5MW"

    # Write 1D hub loads output group from variable subset
    write_hub_loads_subset(source_run_directory, source_run_name, output_directory)

    # Write 2D blade loads output group from variable subset
    write_blade_loads_subset(source_run_directory, source_run_name, output_directory)

    # Write 2D tower loads output group from variable subset
    write_tower_loads_subset(source_run_directory, source_run_name, output_directory)

    print()

    # Clear the cache
    ResultsApi.clear_runs()


def write_hub_loads_subset(source_run_directory, source_run_name, output_directory):

    # Overwrite source run
    output_run_name = source_run_name
    group_number = 22

    # Copy source run from demo data
    prepare_inputs(source_run_directory, source_run_name, output_directory, output_run_name, group_number)

    # Get the run and group
    run = ResultsApi.get_run(output_directory, output_run_name)
    source_group = run.get_group("hub loads: rotating gl coordinates")

    # Set variable names to include in subset
    subset_variable_names = ["rotating hub mx", "rotating hub fx"]

    # Create and write the subset group
    output_group = write_1d_group_subset(source_group, output_directory, output_run_name, subset_variable_names)

    # Print summary
    percent_saved = file_size_diff(source_run_directory, source_run_name, output_directory, output_run_name, group_number)
    print_summary(source_group, output_group, percent_saved)


def write_blade_loads_subset(source_run_directory, source_run_name, output_directory):

    # Overwrite source run
    output_run_name = source_run_name
    group_number = 15

    # Copy source run from demo data
    prepare_inputs(source_run_directory, source_run_name, output_directory, output_run_name, group_number)

    # Get the run and group
    run = ResultsApi.get_run(output_directory, output_run_name)
    source_group = run.get_group("blade 1 loads: principal elastic axes")

    # Get the number of blade stations from the secondary independent variable
    secondary_independent_variable = source_group.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_SECONDARY)
    num_stations = secondary_independent_variable.number_of_values

    # Set the indices of the blade root and tip
    subset_blade_station_indices = [0, num_stations - 1]

    # Set variable names to include in subset
    subset_variable_names = ["blade 1 mx (principal elastic axes)", "blade 1 fx (principal elastic axes)"]

    # Create and write the subset group
    output_group = write_2d_group_subset(source_group, output_directory, output_run_name, subset_variable_names, subset_blade_station_indices)

    # Print summary
    percent_saved = file_size_diff(source_run_directory, source_run_name, output_directory, output_run_name, group_number)
    print_summary(source_group, output_group, percent_saved, num_stations, len(subset_blade_station_indices))


def write_tower_loads_subset(source_run_directory, source_run_name, output_directory):

    # Overwrite source run
    output_run_name = source_run_name
    group_number = 25

    # Copy source run from demo data
    prepare_inputs(source_run_directory, source_run_name, output_directory, output_run_name, group_number)

    # Get the run and group
    run = ResultsApi.get_run(output_directory, output_run_name)
    source_group = run.get_group("tower member loads - local coordinates")

    # Get the number of member ends from the secondary independent variable
    secondary_independent_variable = source_group.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_SECONDARY)
    num_member_ends = secondary_independent_variable.number_of_values

    # Set the indices of both ends for the first and last members
    subset_member_end_indices = [0, 1, num_member_ends - 2, num_member_ends - 1]

    # Set variable names to include in subset
    subset_variable_names = ["tower mx", "tower fx"]

    # Create and write the subset group
    output_group = write_2d_group_subset(source_group, output_directory, output_run_name, subset_variable_names, subset_member_end_indices)

    # Print summary
    percent_saved = file_size_diff(source_run_directory, source_run_name, output_directory, output_run_name, group_number)
    print_summary(source_group, output_group, percent_saved, num_member_ends, len(subset_member_end_indices))


def write_1d_group_subset(
    source_group: Group,
    output_directory: str,
    output_run_name: str,
    subset_variable_names: list[str]) -> OutputGroup1D_Float32:
    """
    Generic function to create a subset of a 1D group from an existing group.

    Parameters
    ----------
    source_group : Group
        The source 1D group to subset.
    output_directory : str
        Output directory for the subset group.
    output_run_name : str
        File name for the subset group.
    subset_variable_names : list[str]
        List of variable names to extract.
    """

    # Reuse primary independent variable (time axis)
    primary_independent_variable = source_group.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_PRIMARY)

    dest_group_header_path = f"{output_directory}/{output_run_name}.%{source_group.number}"

    # Create new output group
    output_group = OutputGroup1D_Float32(
        dest_group_header_path,
        source_group.name,
        source_group.calculation_short_name,
        source_group.config,
        primary_independent_variable)

    # Extract and write variables
    data_copies = {}
    for var_name in subset_variable_names:
        variable = source_group.get_variable_1d(var_name)
        data = variable.get_data()
        output_group.add_variable_with_data(variable.name, variable.si_unit, data)
        data_copies[var_name] = np.copy(data)
    
    # Write the subset output group
    success = ResultsApi.write_output(output_group)
    assert success, f"Failed to create subset group: {dest_group_header_path}"

    # Validate data integrity
    for var_name in subset_variable_names:
        prove_1d_roundtrip(data_copies[var_name], var_name, output_directory, output_run_name)
        
    return output_group


def write_2d_group_subset(
    source_group: Group,
    output_directory: str,
    output_run_name: str,
    subset_variable_names: list[str],
    subset_label_indices: list[int]) -> OutputGroup2D_Float32:
    """
    Generic function to create a subset of a 2D group from an existing group.

    Parameters
    ----------
    source_group : Group
        The source 2D group to subset.
    output_directory : str
        Output directory for the subset group.
    output_run_name : str
        File name for the subset group.
    subset_variable_names : list[str]
        List of variable names to extract.
    subset_label_indices : list[int]
        Indices of the secondary independent variable to extract.
    """

    # Reuse primary independent variable (time axis)
    primary_independent_variable = source_group.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_PRIMARY)

    # Get secondary independent variable, and create subset from requested indices
    secondary_independent_variable = source_group.get_independent_variable(dnvres.INDEPENDENT_VARIABLE_ID_SECONDARY)
    if (secondary_independent_variable.has_numeric_values):
        all_labels = secondary_independent_variable.get_values_as_number()
    else:
        all_labels = secondary_independent_variable.get_values_as_string()
    
    subset_labels = [all_labels[i] for i in subset_label_indices]
    secondary_independent_variable_subset = IndependentVariable(secondary_independent_variable.name, secondary_independent_variable.si_unit, subset_labels)

    dest_group_header_path = f"{output_directory}/{output_run_name}.%{source_group.number}"

    # Create new output group
    output_group = OutputGroup2D_Float32(
        dest_group_header_path,
        source_group.name,
        source_group.calculation_short_name,
        source_group.config,
        primary_independent_variable,
        secondary_independent_variable_subset)

    # Helper function to extract and stack data
    def get_stacked_data(variable_2d, indices):
        rows = [variable_2d.get_data_at_index(idx) for idx in indices]
        return np.vstack(rows)

    # Extract and write variables
    data_copies = {}
    for var_name in subset_variable_names:
        variable = source_group.get_variable_2d(var_name)
        stacked_data = get_stacked_data(variable, subset_label_indices)
        output_group.add_variable_with_data(variable.name, variable.si_unit, stacked_data)
        data_copies[var_name] = np.copy(stacked_data)
    
    # Write the subset output group
    success = ResultsApi.write_output(output_group)
    assert success, f"Failed to create subset group: {dest_group_header_path}"

    # Validate data integrity
    for var_name in subset_variable_names:
        prove_2d_roundtrip(data_copies[var_name], var_name, output_directory, output_run_name)
        
    return output_group


def prove_1d_roundtrip(data_original, user_variable_name, output_directory, output_run_name):

    ResultsApi.clear_runs()
    run = ResultsApi.get_run(output_directory, output_run_name)
    assert run.calculation_type == dnvres.CALCULATION_TYPE_POWER_PRODUCTION_SIMULATION

    data_roundtrip = run.get_variable_1d(user_variable_name).get_data()
    assert len(data_roundtrip) == len(data_original)

    point_index = 0
    for val in data_original:
        assert (np.float32(val)) == data_roundtrip[point_index], f"Expected value equality on data roundtrip: array index = {point_index}"
        point_index = point_index + 1


def prove_2d_roundtrip(data_original, user_variable_name, output_directory, output_run_name):

    ResultsApi.clear_runs()
    num_secondary_axis_vals = len(data_original)

    run = ResultsApi.get_run(output_directory, output_run_name)
    assert run.calculation_type == dnvres.CALCULATION_TYPE_POWER_PRODUCTION_SIMULATION

    data_roundtrip = run.get_variable_2d(user_variable_name).get_data()
    assert len(data_roundtrip) == num_secondary_axis_vals

    point_index = 0
    for secondary_axis_value_index in range(0, num_secondary_axis_vals, 1):
        assert len(data_roundtrip[secondary_axis_value_index]) == len(data_original[secondary_axis_value_index])
        for val in data_original[secondary_axis_value_index]:
            assert (np.float32(val)) == (data_roundtrip[secondary_axis_value_index][point_index]), f"Expected value equality on data roundtrip: secondary axis index = {secondary_axis_value_index} point index = {point_index}"
            point_index = point_index + 1
        point_index = 0


def print_summary(source_group: Group, output_group, percent_saved, source_num_labels = 0, output_num_labels = 0):

    source_labels = ""
    output_labels = ""
    if source_num_labels > 0:
        source_labels = f"{source_num_labels} secondary independent variable values; "
    if output_num_labels > 0:
        output_labels = f"{output_num_labels} secondary independent variable values; "

    print(f"\nSuccessfully created group '{output_group.name}' from a subset of an existing group: {output_group.path}" +
          f"\nOriginal group   : {source_group.number_of_variables} variables; " +
          f"{source_labels}" +
          f"{source_group.data_point_count} points" +
          f"\nNew group        : {output_group.number_of_variables} variables; " + 
          f"{output_labels}" +
          f"{output_group.data_point_count} points" +
          f"\nDisk space saved : {percent_saved} %")


def file_size(path: Path) -> int:
    if path.exists():
        return path.stat().st_size
    return 0


def file_size_diff(source_run_directory, source_run_name, output_directory, output_run_name, group_number) -> float:

    source_group_header_path = Path(source_run_directory) / f"{source_run_name}.%{group_number}"
    source_group_data_path = Path(source_run_directory) / f"{source_run_name}.${group_number}"

    dest_group_header_path = Path(output_directory) / f"{output_run_name}.%{group_number}"
    dest_group_data_path = Path(output_directory) / f"{output_run_name}.${group_number}"

    source_total = file_size(source_group_header_path) + file_size(source_group_data_path)
    dest_total = file_size(dest_group_header_path) + file_size(dest_group_data_path)

    if source_total == 0:
        return 0.0

    # Calculate percentage saving
    return round((1 - dest_total / source_total) * 100, 2)


def prepare_inputs(source_run_directory, source_run_name, output_directory, output_run_name, group_number):
    
    source_group_header_path = f"{source_run_directory}/{source_run_name}.%{group_number}"
    source_group_data_path = f"{source_run_directory}/{source_run_name}.${group_number}"
    source_group_termination_path = f"{source_run_directory}/{source_run_name}.$TE"

    dest_group_header_path = f"{output_directory}/{output_run_name}.%{group_number}"
    dest_group_data_path = f"{output_directory}/{output_run_name}.${group_number}"
    dest_group_termination_path = f"{output_directory}/{output_run_name}.$TE"

    shutil.copy(source_group_header_path, dest_group_header_path)
    shutil.copy(source_group_data_path, dest_group_data_path)
    shutil.copy(source_group_termination_path, dest_group_termination_path)


if __name__ == "__main__":
    run_script()
