from __future__ import annotations
from pathlib import Path
from damask_parse.readers import read_HDF5_file
from damask_parse.utils import generate_viz


def process_HDF5(
    damask_hdf5_file: Path | str,
    damask_post_processing: list[dict],
    VE_response_data: dict,
    damask_viz: list[dict] | dict | bool | None,
    remove_damask_hdf5: bool,
):
    """
    Operate on and extract data from an HDF5 file generated by a DAMASK run.

    Supported post-processing operations
    ------------------------------------
    - `add_absolute`
    - `add_calculation`
    - `add_stress_Cauchy`
    - `add_determinant`
    - `add_deviator`
    - `add_eigenvalue`
    - `add_eigenvector`
    - `add_IPF_color`
    - `add_maximum_shear`
    - `add_equivalent_Mises`
    - `add_norm`
    - `add_stress_second_Piola_Kirchhoff`
    - `add_pole`
    - `add_rotation`
    - `add_spherical`
    - `add_strain`
    - `add_stretch_tensor`
    - `add_curl`
    - `add_divergence`
    - `add_gradient`

    Parameters
    ----------
    damask_hdf5_path : Path or str
        Path to the HDF5 file.
    damask_post_processing
        List of methods to invoke on the DADF5 object. This is a list of dicts with the
        following keys:
            name : str
                The name of the DADF5 method.
            args : dict
                Parameter names and their values to pass to the DADF5 method. This
                assumes all DADF5 method parameters are of positional-or-keyword type.
            opts : dict, optional
                Additional options.
    VE_response_data : dict
        Contents are as follows:

        volume_data : list of dict, optional
            List of data to extract from the entire volume. This is a list of dict
            with following keys:
                field_name: str
                    Name of the data field to extract
                increments: list of dict, optional
                    List of increment specifications to extract data from. Values
                    refer to increments in the simulation. Default to all. This is
                    a list of dict one of the following sets of keys:
                        values: list of int
                            List of incremnts to extract
                        ----OR----
                        start: int
                            First increment to extract
                        stop: int
                            Final incremnt to extract (inclusive)
                        step: int
                            Step between increments to extract
                out_name: str, optional
                    Name of the data
                transforms: list of dict, optional
                    If specified this is a list of dicts, each with the following keys:
                        sum_along_axes : int, optional
                            If specified, take the sum the array along this axis.
                        mean_along_axes: int, optional
                            If specified, take the mean average of the array along this
                            axis.

        phase_data : list of dict, optional
            List of data to extract from a single phase. This is a list of dict
            with following keys:
                field_name: str
                    Name of the data field to extract
                phase_name : str
                    Name of phase to
                increments: list of dict, optional
                    List of increment specifications to extract data from. Values
                    refer to increments in the simulation. Default to all. This is
                    a list of dict one of the following sets of keys:
                        values: list of int
                            List of incremnts to extract
                        ----OR----
                        start: int
                            First increment to extract
                        stop: int
                            Final incremnt to extract (inclusive)
                        step: int
                            Step between increments to extract
                out_name: str, optional
                    Name of the data
                transforms: list of dict, optional
                    If specified this is a list of dicts, each with the following keys:
                        sum_along_axes : int, optional
                            If specified, take the sum the array along this axis.
                        mean_along_axes: int, optional
                            If specified, take the mean average of the array along this
                            axis.

        field_data : list of dict, optional
            List of field data to extract. Only extracts for the first constituent
            of each material point. This is a list of dict with following keys:
                field_name: str
                    Name of the data field to extract
                increments: list of dict, optional
                    List of increment specifications to extract data from. Values
                    refer to increments in the simulation. Default to all. This is
                    a list of dict one of the following sets of keys:
                        values: list of int
                            List of incremnts to extract
                        ----OR----
                        start: int
                            First increment to extract
                        stop: int
                            Final incremnt to extract (inclusive)
                        step: int
                            Step between increments to extract
            Special field_name keys exist, 'grain' and 'phase'. Use 'u_n' or 'u_p'
            for displacement.

        grain_data : list of dict, optional
            List of grain data to extract. Only extracts for the first constituent
            of each material point. This is a list of dict with following keys:
                field_name: str
                    Name of the data field to extract
                increments: list of dict, optional
                    List of increment specifications to extract data from. Values
                    refer to increments in the simulation. Default to all. This is
                    a list of dict one of the following sets of keys:
                        values: list of int
                            List of incremnts to extract
                        ----OR----
                        start: int
                            First increment to extract
                        stop: int
                            Final incremnt to extract (inclusive)
                        step: int
                            Step between increments to extract
    damask_viz
        Visualisation control data.

    Returns
    -------
    volume_element_response : dict
        Dict with keys determined by the `field_data` dict.
    """
    VE_response = read_HDF5_file(
        hdf5_path=damask_hdf5_file,
        operations=damask_post_processing,
        **VE_response_data,
    )
    generate_viz(hdf5_path=damask_hdf5_file, viz_spec=damask_viz, parsed_outs=VE_response)

    # remove damask HDF5 file (TODO: needs to be a better way to do this in hpcflow)
    if remove_damask_hdf5:
        damask_hdf5_file.unlink()

    return VE_response
