amrex_tools.py

"""
AMReX-agnostic tooling extracted from pele_tools.

This module hosts generic utilities for inputs parsing and run directory setup.
Pele-specific knowledge and repositories remain in utils/pele_tools.py.
"""

from __future__ import annotations

from datetime import datetime
import json
from pathlib import Path
from typing import Any

from langchain.tools import tool


def parse_amrex_inputs(filepath: str, source_info: dict | None = None) -> dict[str, Any]:
    """
    Parse an AMReX inputs file into a nested dict with metadata.

    Parameters
    ----------
    filepath : str
        Path to the inputs file.
    source_info : Optional[Dict], optional
        Source metadata to include in the output.

    Returns
    -------
    Dict[str, Any]
        Parsed inputs and metadata.
    """
    inputs: dict[str, Any] = {}
    filepath = Path(filepath)

    metadata = {
        "inputs_file": filepath,
        "inputs_name": filepath.name,
        "inputs_path": filepath.absolute(),
        "requested_path": str(filepath),
        "parsed_at": datetime.now().isoformat(),
    }

    try:
        metadata["relative_path"] = filepath.relative_to(Path.cwd())
    except ValueError:
        metadata["relative_path"] = filepath.absolute()

    if source_info:
        metadata["source"] = source_info.get("source", "unknown")
        metadata["example_name"] = source_info.get("example_name")
        metadata["repo_path"] = source_info.get("repo_path")
        metadata["version"] = source_info.get("version")
        metadata["checksum"] = source_info.get("checksum")
    else:
        metadata["source"] = "local"

    with open(filepath) as handle:
        for line in handle:
            line = line.split("#")[0].split("!")[0].strip()
            if not line or "=" not in line:
                continue

            key, value = line.split("=", 1)
            key = key.strip()
            value = value.strip()

            if "." in key:
                section, param = key.rsplit(".", 1)
                if section not in inputs:
                    inputs[section] = {}
                inputs[section][param] = value
            else:
                inputs[key] = value

    inputs["_metadata"] = metadata
    return inputs


@tool
def validate_amrex_inputs(inputs_file: str) -> str:
    """
    Validate inputs file with generic AMReX checks.

    This intentionally avoids solver-specific assumptions. It only validates:
    - File exists
    - Inputs can be parsed

    Parameters
    ----------
        inputs_file: Path to inputs file

    Returns
    -------
        str: Validation results as JSON {"errors": [...], "warnings": [...]}
    """
    errors: list[str] = []
    warnings: list[str] = []

    inputs_path = Path(inputs_file)
    if not inputs_path.exists():
        return json.dumps({
            "errors": [f"File not found: {inputs_file}"],
            "warnings": []
        }, indent=2)

    try:
        parse_amrex_inputs(inputs_file)
    except Exception as exc:
        return json.dumps({
            "errors": [f"Parse error: {exc}"],
            "warnings": []
        }, indent=2)

    return json.dumps({"errors": errors, "warnings": warnings}, indent=2)


def dict_to_amrex_inputs(
    inputs_dict: dict[str, Any],
    output_path: str,
    base_inputs_dict: dict[str, Any] | None = None,
) -> Path:
    """
    Convert dict to AMReX inputs file.

    ParmParse uses the last occurrence of a parameter if it appears multiple times.
    This function comments out overridden base values and writes modifications last.

    Parameters
    ----------
    inputs_dict : Dict[str, Any]
        Updated inputs mapping.
    output_path : str
        Path to write the inputs file.
    base_inputs_dict : Optional[Dict[str, Any]], optional
        Base inputs to overlay.

    Returns
    -------
    Path
        Path to the written inputs file.
    """
    lines = [
        "# Generated by AMReX Assistant",
        "# NOTE: If a parameter appears multiple times, LAST value is used (ParmParse)",
        "",
    ]

    overridden = set()
    if base_inputs_dict:
        for section in inputs_dict:
            if section in base_inputs_dict and isinstance(inputs_dict[section], dict):
                for param in inputs_dict[section]:
                    if param in base_inputs_dict[section]:
                        overridden.add(f"{section}.{param}")

    if base_inputs_dict:
        lines.extend(
            [
                "# ========================================",
                "# Base parameters (from reference input)",
                "# ========================================",
                "",
            ]
        )
        for section in sorted(base_inputs_dict.keys()):
            params = base_inputs_dict[section]
            if isinstance(params, dict):
                lines.append(f"[{section}]")
                for param, value in params.items():
                    param_key = f"{section}.{param}"
                    if param_key in overridden:
                        lines.append(f"# {section}.{param} = {value}  # OVERRIDDEN")
                    else:
                        lines.append(f"{section}.{param} = {value}")
                lines.append("")
            else:
                lines.append(f"{section} = {params}")

    lines.extend(
        [
            "# ========================================",
            "# Modifications",
            "# ========================================",
            "",
        ]
    )

    for section in sorted(inputs_dict.keys()):
        params = inputs_dict[section]
        if isinstance(params, dict):
            lines.append(f"[{section}]")
            for param, value in params.items():
                param_key = f"{section}.{param}"
                if param_key in overridden:
                    lines.append(f"{section}.{param} = {value}  # OVERRIDE")
                else:
                    lines.append(f"{section}.{param} = {value}")
            lines.append("")
        else:
            lines.append(f"{section} = {params}")

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as handle:
        handle.write("\n".join(lines))

    return output_path


@tool
def setup_run_directory(
    base_name: str = "run",
    base_dir: str | None = None,
    config: dict | None = None,
) -> str:
    """
    Create run directory with timestamp.

    Uses config.output_dir when provided.

    Parameters
    ----------
    base_name : str, optional
        Base name for the run directory.
    base_dir : Optional[str], optional
        Root directory to create the run directory in.
    config : Optional[Dict], optional
        Optional config with output_dir override.

    Returns
    -------
    str
        Absolute path to the created run directory.
    """
    if config and base_dir is None:
        output_dir = config.get("output_dir")
        if output_dir:
            base_dir = str(output_dir)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir_name = f"run_{base_name}_{timestamp}"

    run_dir = Path(base_dir) / run_dir_name if base_dir else Path.cwd() / run_dir_name

    run_dir.mkdir(parents=True, exist_ok=True)
    return str(run_dir.resolve())


@tool
def find_inputs_file(
    directory: str,
    prefer_git_oldest: bool = True,
    config: dict | None = None,
) -> str | None:
    """
    Find inputs file in a directory using heuristics.

    Parameters
    ----------
    directory : str
        Directory to search.
    prefer_git_oldest : bool, optional
        Prefer the oldest file by git history when available.
    config : Optional[Dict], optional
        Optional config with inputs selection hints.

    Returns
    -------
    Optional[str]
        Selected inputs file path or ``None``.
    """
    import subprocess

    directory = Path(directory)
    if not directory.exists():
        return None

    candidates = []
    for pattern in ["inputs*", "*.inp"]:
        candidates.extend(directory.glob(pattern))

    candidates = [c for c in candidates if c.is_file()]
    if not candidates:
        return None
    if len(candidates) == 1:
        return str(candidates[0])

    try:
        result = subprocess.run(
            ["git", "rev-parse", "--git-dir"],
            cwd=directory,
            capture_output=True,
            timeout=2,
        )

        if result.returncode == 0:
            file_scores: dict[Path, dict[str, int]] = {}
            for candidate in candidates:
                rel_path = candidate.relative_to(directory)

                first_commit = subprocess.run(
                    ["git", "log", "--follow", "--format=%at", "--", str(rel_path)],
                    cwd=directory,
                    capture_output=True,
                    text=True,
                    timeout=2,
                )

                commit_count = subprocess.run(
                    ["git", "log", "--follow", "--oneline", "--", str(rel_path)],
                    cwd=directory,
                    capture_output=True,
                    text=True,
                    timeout=2,
                )

                if first_commit.returncode == 0 and first_commit.stdout.strip():
                    timestamps = first_commit.stdout.strip().split("\n")
                    oldest_timestamp = int(timestamps[-1]) if timestamps else 0
                    n_commits = (
                        len(commit_count.stdout.strip().split("\n"))
                        if commit_count.returncode == 0
                        else 0
                    )

                    file_scores[candidate] = {
                        "oldest_timestamp": oldest_timestamp,
                        "n_commits": n_commits,
                    }

            if file_scores:
                if prefer_git_oldest:
                    best = min(
                        file_scores.items(),
                        key=lambda x: x[1]["oldest_timestamp"],
                    )
                else:
                    best = max(
                        file_scores.items(),
                        key=lambda x: x[1]["n_commits"],
                    )
                return str(best[0])

    except (subprocess.TimeoutExpired, Exception):
        pass

    candidates_sorted = sorted(
        candidates,
        key=lambda x: (x.name != "inputs", len(x.name), x.name),
    )

    return str(candidates_sorted[0])


@tool
def copy_to_rundir(
    run_dir: str,
    executable_path: str,
    inputs_path: str | None = None,
    inputs_dir: str | None = None,
    config: dict | None = None,
) -> dict[str, str]:
    """
    Copy executable and inputs to run directory.

    Parameters
    ----------
    run_dir : str
        Run directory path.
    executable_path : str
        Executable path to copy.
    inputs_path : Optional[str], optional
        Inputs file path to copy.
    inputs_dir : Optional[str], optional
        Directory containing inputs candidates.
    config : Optional[Dict], optional
        Optional config with inputs selection hints.

    Returns
    -------
    Dict[str, str]
        Paths to copied artifacts keyed by type.
    """
    import shutil

    run_dir = Path(run_dir)
    files: dict[str, str] = {}

    baseline_dir = None
    if executable_path:
        exe_path = Path(executable_path)
        if exe_path.exists():
            exe_dest = run_dir / exe_path.name
            shutil.copy2(exe_path, exe_dest)
            exe_dest.chmod(0o755)
            files["executable"] = str(exe_dest)
            baseline_dir = exe_path.parent
        else:
            raise FileNotFoundError(f"Executable not found: {executable_path}")

    if inputs_path is None and inputs_dir is not None:
        inputs_path = find_inputs_file.invoke({"directory": inputs_dir})
        if inputs_path:
            print(f"[INFO] Auto-selected inputs: {Path(inputs_path).name}")

    if inputs_path:
        inputs_src = Path(inputs_path)
        if inputs_src.exists():
            inputs_dest = run_dir / "inputs"
            if inputs_src.resolve() != inputs_dest.resolve():
                shutil.copy2(inputs_src, inputs_dest)
            files["inputs"] = str(inputs_dest)

            if baseline_dir is None and inputs_dir is not None:
                baseline_dir = Path(inputs_dir)
            if baseline_dir:
                aux_files = _copy_auxiliary_files(baseline_dir, run_dir, inputs_src)
                files.update(aux_files)
        else:
            raise FileNotFoundError(f"Inputs not found: {inputs_path}")
    else:
        raise ValueError("Must provide either inputs_path or inputs_dir")

    grid = "?"
    max_level = "?"

    if inputs_src.exists():
        try:
            params = parse_amrex_inputs(str(inputs_src))
            grid = params.get("amr", {}).get("n_cell", "?")
            max_level = params.get("amr", {}).get("max_level", "?")
        except Exception:
            pass

    readme = run_dir / "README.md"
    exe_name = Path(executable_path).name if executable_path else "?"
    readme_content = f"""# AMReX Simulation Run

**Created:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

## Configuration
- Grid: {grid}
- Max level: {max_level}
- Executable: {exe_name}

## Files
- `inputs` - AMReX configuration
- `{exe_name}` - Compiled executable
- `run_local.sh` - Local run script (generated for local runs)
- `submit.sh` - SLURM batch script (generated for cluster runs)
- `run.out` - Simulation output (batch runs only)
- `plt*` - AMReX plotfiles (after run)

## Run Locally (if `run_local.sh` exists)
```bash
cd {run_dir.name}
./run_local.sh
```

## Submit Job (if `submit.sh` exists)
```bash
cd {run_dir.name}
sbatch submit.sh
```

## Monitor
```bash
squeue -u $USER  # batch only
tail -f run.out  # batch only
```
"""
    readme.write_text(readme_content)
    files["readme"] = str(readme)

    return files


def _copy_auxiliary_files(
    source_dir: Path,
    run_dir: Path,
    inputs_path: Path,
    max_size_mb: int = 50,
) -> dict[str, str]:
    """Copy auxiliary data files from baseline directory."""
    import re
    import shutil

    copied: dict[str, str] = {}
    max_bytes = max_size_mb * 1024 * 1024

    skip_patterns = [
        r"\.git",
        r"\.o$",
        r"\.mod$",
        r"\.exe$",
        r"\.ex$",
        r"^\.",
        r"GNUmakefile",
        r"CMake",
        r"build",
        r"inputs\.",
    ]

    data_extensions = [
        ".dat",
        ".yaml",
        ".yml",
        ".xml",
        ".json",
        ".txt",
        ".csv",
        ".h5",
        ".hdf5",
        ".inp",
        ".msh",
        ".geo",
    ]

    print(f"[INFO] Scanning {source_dir.name} for auxiliary files...")

    for item in source_dir.iterdir():
        if item.is_dir():
            continue
        if any(re.search(pattern, item.name) for pattern in skip_patterns):
            continue
        if item.stat().st_size > max_bytes:
            continue

        should_copy = False
        if any(item.suffix.lower() == ext for ext in data_extensions):
            should_copy = True

        if inputs_path.exists():
            with open(inputs_path) as handle:
                inputs_text = handle.read()
                if item.name in inputs_text:
                    should_copy = True

        if should_copy and item.name != "inputs":
            try:
                dest = run_dir / item.name
                shutil.copy2(item, dest)
                copied[item.name] = str(dest)
                print(
                    f"      [OK] Copied: {item.name} "
                    f"({item.stat().st_size / 1024:.1f} KB)"
                )
            except Exception as exc:
                print(f"      [WARN] Failed to copy {item.name}: {exc}")

    if copied:
        print(f"[OK] Copied {len(copied)} auxiliary file(s)")

    return copied


def find_latest_plotfile(directory: str = ".", pattern: str = "plt*") -> str | None:
    """
    Find the latest plotfile in a directory.

    Parameters
    ----------
    directory : str, optional
        Directory to search.
    pattern : str, optional
        Glob pattern for plotfiles.

    Returns
    -------
    Optional[str]
        Path to the latest plotfile or ``None``.
    """
    directory = Path(directory)
    plotfiles = list(directory.glob(pattern))
    if not plotfiles:
        return None
    return str(max(plotfiles, key=lambda p: p.stat().st_mtime))


def extract_plotfile_metadata(plotfile_path: str) -> dict[str, Any]:
    """
    Extract metadata from an AMReX plotfile path.

    Parameters
    ----------
    plotfile_path : str
        Plotfile directory path.

    Returns
    -------
    Dict[str, Any]
        Extracted plotfile metadata.
    """
    plotfile = Path(plotfile_path)
    metadata = {
        "plotfile": str(plotfile),
        "name": plotfile.name,
        "path": str(plotfile.parent),
    }

    header = plotfile / "Header"
    if header.exists():
        try:
            lines = header.read_text().splitlines()
            if lines:
                metadata["description"] = lines[0]
        except Exception:
            pass

    return metadata


def create_yt_sliceplot(
    plotfile_path: str,
    field: str = "density",
    axis: str = "z",
    output: str | None = None,
) -> str:
    """
    Create a sliceplot using yt and save as an image.

    Parameters
    ----------
    plotfile_path : str
        Plotfile directory path.
    field : str, optional
        Field to plot.
    axis : str, optional
        Axis to slice.
    output : Optional[str], optional
        Output image path.

    Returns
    -------
    str
        Path to the saved image.
    """
    import yt

    ds = yt.load(plotfile_path)
    slc = yt.SlicePlot(ds, axis, field)
    if output is None:
        output = f"slice_{field}_{axis}.png"
    slc.save(output)
    return output


def create_1d_profile(
    plotfile_path: str,
    field: str = "density",
    axis: str = "x",
    output: str | None = None,
) -> str:
    """
    Create a 1D profile plot using yt.

    Parameters
    ----------
    plotfile_path : str
        Plotfile directory path.
    field : str, optional
        Field to plot.
    axis : str, optional
        Axis to profile.
    output : Optional[str], optional
        Output image path.

    Returns
    -------
    str
        Path to the saved image.
    """
    import yt

    ds = yt.load(plotfile_path)
    prof = yt.create_profile(ds.all_data(), axis, fields=[field])

    if output is None:
        output = f"profile_{field}_{axis}.png"

    import matplotlib.pyplot as plt

    plt.figure()
    plt.plot(prof.x.value, prof[field])
    plt.xlabel(axis)
    plt.ylabel(field)
    plt.savefig(output)
    plt.close()

    return output


# Backward-compatible aliases
parse_pele_inputs = parse_amrex_inputs


def dict_to_pele_inputs(
    inputs_dict: dict[str, Any],
    output_path: str,
    base_inputs_dict: dict[str, Any] | None = None,
) -> Path:
    """
    Convert a Pele inputs dict to a file on disk.

    Parameters
    ----------
    inputs_dict : Dict[str, Any]
        Updated inputs mapping.
    output_path : str
        Path to write the inputs file.
    base_inputs_dict : Optional[Dict[str, Any]], optional
        Base inputs to overlay.

    Returns
    -------
    Path
        Path to the written inputs file.
    """
    return dict_to_amrex_inputs(inputs_dict, output_path, base_inputs_dict)


__all__ = [
    "parse_amrex_inputs",
    "dict_to_amrex_inputs",
    "parse_pele_inputs",
    "dict_to_pele_inputs",
    "setup_run_directory",
    "find_inputs_file",
    "copy_to_rundir",
    "find_latest_plotfile",
    "extract_plotfile_metadata",
    "create_yt_sliceplot",
    "create_1d_profile",
]