"""
Extract scalar values from .npz result files.

The extraction paths correspond to the hierarchical structure returned by
`load_results()` (see io.py), where metrics and observables are JSON-decoded
and stored under the 'results' key.

Usage:

.. code-block:: sh

    llg3d.extract run.npz params/np results/metrics/total_time

This will extract the number of processes and total simulation time from the run.npz file.
Paths are separated by '/' and navigate through nested dictionaries and object attributes.
"""

import argparse
from pathlib import Path

import numpy as np

from llg3d.io import RunResults, load_results


def _navigate(value: object, levels: list[str]):
    """
    Recursively navigate through nested structures using keys.

    Tries subscript access (dicts/arrays) first, then attribute access (objects).

    Args:
        value: The current value to navigate.
        levels: List of keys/attributes to navigate through.

    Returns:
        The final value after navigating through all levels.
    """
    if not levels:
        return value

    k = levels[0]
    try:
        # Try dict/array subscript access
        next_value = value[k]  # type: ignore
    except (TypeError, KeyError, IndexError):
        # Fallback to attribute access
        next_value = getattr(value, k)

    return _navigate(next_value, levels[1:])

ExtractedValue = float | int | str | bool | np.integer | np.floating

def extract_values(npz_file: Path, *keys: str) -> list[ExtractedValue]:
    """
    Extract scalar values from a .npz file.

    Args:
        npz_file: Path to the .npz result file
        *keys: tuple of keys of the scalar values to extract
            (slash-separated for nested keys)

    Returns:
        The list of extracted scalar values

    Raises:
        ValueError: If one of the extracted values is not a scalar
    """
    results: RunResults = load_results(npz_file)
    values: list[ExtractedValue] = []
    for key in keys:
        # Split path into levels
        levels: list[str] = key.split("/")
        raw_value = _navigate(results, levels)

        # Handle numpy scalars and arrays with shape ()
        if isinstance(raw_value, np.ndarray):
            if raw_value.shape != ():
                raise ValueError(f"Extracted value for key '{key}' is not a scalar.")
            value = raw_value.item()  # Convert numpy scalar to Python scalar
        else:
            value = raw_value

        if not isinstance(value, (float, int, str, bool, np.integer, np.floating)):
            raise ValueError(f"Extracted value for key '{key}' is not a scalar.")
        values.append(value)
    return values


def main():  # pragma: no cover
    """Parse command line arguments and print simulation info."""
    parser = argparse.ArgumentParser(
        description="Extract scalar values from .npz result files.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("filename", help="Path to the .npz result file", type=Path)
    parser.add_argument(
        "keys",
        nargs="+",
        help="Key(s) of the scalar value(s) to extract (slash-separated for nested keys)",
    )
    args = parser.parse_args()

    values = extract_values(args.filename, *args.keys)
    print(" ".join(map(str, values)))
