"""Plot x-profiles of magnetization from a result file."""

import argparse
from pathlib import Path

from matplotlib import pyplot as plt

from ..grid import Grid
from ..io import RunResults, load_results
from .extract import extract_values


def parse_slice(slice_str: str) -> slice:
    """
    Parse a string representation of a slice into a slice object.

    Handles formats like 'start:stop:step', 'start:stop', ':stop', '::step', etc.

    Args:
        slice_str: String representation of the slice.

    Returns:
        A slice object corresponding to the input string.
    """
    parts = slice_str.split(":")
    # Convert parts to int or None
    args = [int(p) if p else None for p in parts]
    return slice(*args)


def plot_x_profiles(
    file: Path,
    image_filepath: Path | str | None = None,
    show: bool = False,
    time_slice: str = ":",
    m_index: int = 0,
):
    """
    Plot the results from the given files.

    Args:
        file: Path to the result file.
        image_filepath: Path to the output image file.
            if None, the image will not be saved.
        show: display the graph in a graphical window.
        time_slice: String representing the time slice to plot (e.g. ":-5", "::10").
        m_index: Index of the magnetization component to plot (0, 1, or 2).
    """
    fig, ax = plt.subplots()
    result: RunResults = load_results(file)
    grid_params = extract_values(
        file, "params/Jx", "params/Jy", "params/Jz", "params/dx"
    )
    grid = Grid(*grid_params)  # type: ignore
    x_coords = grid.get_x_coords(local=False)
    x_profiles = result.get_record("x_profiles")

    # Parse and apply slice
    sl = parse_slice(time_slice)
    times = x_profiles["t"][sl]
    m_profiles = x_profiles[f"m{m_index}"][sl]

    for i, time in enumerate(times):
        ax.plot(x_coords, m_profiles[i], label=f"t = {time:.3e} s")
    ax.set_xlabel("x")
    ax.set_ylabel(rf"$m_{{{m_index}}}$")
    ax.grid()
    ax.legend()
    ax.set_title(rf"Longitudinal profiles of magnetization $m_{{{m_index + 1}}}$")

    if show:
        plt.show()

    if image_filepath is not None:
        fig.savefig(image_filepath, dpi=300)
        print(f"Written to {image_filepath}")


def get_cli_args(
    description: str | None, default_image_filepath: Path
) -> argparse.Namespace:
    """
    Parse command-line arguments for post-processing scripts.

    Args:
        description: Description of the script for the help message.
        default_image_filepath: Default path to save the output image.

    Returns:
        Parsed command-line arguments.
    """
    parser = argparse.ArgumentParser(
        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("file", type=Path, help="Path to the result file.")
    parser.add_argument(
        "-i",
        "--image_filepath",
        type=Path,
        default=default_image_filepath,
        help="Path to save the image",
    )
    parser.add_argument(
        "-s",
        "--show",
        action="store_true",
        default=False,
        help="Display the plot (omit to disable display in non-interactive runs).",
    )
    parser.add_argument(
        "-t",
        "--time",
        type=str,
        default=":",
        help="Slice for time steps (e.g., '-5:' for last 5, '::10' for every 10th).",
    )
    parser.add_argument(
        "-m",
        "--m_index",
        type=int,
        choices=[1, 2, 3],
        default=1,
        help="Index of the magnetization component to plot (1, 2, or 3).",
    )
    return parser.parse_args()


def main():  # pragma: no cover
    """Parse CLI arguments and call the plot function."""
    args = get_cli_args(
        description=__doc__,
        default_image_filepath=Path("x_profile.png"),
    )

    plot_x_profiles(
        args.file,
        image_filepath=args.image_filepath,
        show=args.show,
        time_slice=args.time,
        m_index=args.m_index,
    )
