import numpy as np
from pathlib import Path


def load_from_dir(plane_dir):
    """
    Load all result files generated by `masknmf.run_plane()` from a given plane directory.

    Returns
    -------
    dict
        Dictionary with the following keys:
        - 'moco': np.ndarray
              Motion-corrected data array.
        - 'pmd_demixer': object
              Deserialized PMD demixer object containing the full demixing state.
        - 'a': np.ndarray
              Spatial components array with shape (height, width, n_components).
              Each component is a 2D spatial footprint.
        - 'c': np.ndarray
              Temporal components array with shape (n_frames, n_components).
              Each column is a time series for one component.
    """
    plane_dir = Path(plane_dir).expanduser()

    results = {}
    moco_path = plane_dir / "data_reg.npy"
    pmd_demixer_path = plane_dir / "pmd_demixer.npy"
    a_path = plane_dir / "a.npy"
    c_path = plane_dir / "a.npy"

    if moco_path.exists():
        results["moco"] = np.load(moco_path, allow_pickle=True)
    if pmd_demixer_path.exists():
        results["pmd_demixer"] = np.load(pmd_demixer_path, allow_pickle=True).item()
    if a_path.exists():
        results["a"] = np.load(a_path)
    if c_path.exists():
        results["c"] = np.load(c_path)
    return results
