from astropy.wcs import WCS
from collections import OrderedDict
import numpy as np
import xarray as xr
from typing import List, Union
from .common import _c, _compute_world_sph_dims, _l_m_attr_notes
from xradio._utils.coord_math import _deg_to_rad
from xradio._utils.dict_helpers import (
    make_direction_location_dict,
    make_spectral_coord_reference_dict,
    make_quantity,
    make_skycoord_dict,
    make_time_coord_attrs,
)
from xradio._utils.logging import xradio_logger


def _is_galactic_frame(direction_reference: str) -> bool:
    """
    Check whether a direction reference frame is Galactic.

    Parameters
    ----------
    direction_reference : str
        Direction frame name.

    Returns
    -------
    bool
        True when the frame is Galactic, otherwise False.
    """
    return direction_reference.lower() == "galactic"


def _spherical_ctype_for_frame(direction_reference: str) -> list[str]:
    """
    Get WCS CTYPE axis tokens for a direction reference frame.

    Parameters
    ----------
    direction_reference : str
        Direction frame name.

    Returns
    -------
    list[str]
        Two-element CTYPE list for longitude/latitude axes.
    """
    if _is_galactic_frame(direction_reference):
        return ["GLON", "GLAT"]
    return ["RA", "Dec"]


def _spherical_coord_names_for_frame(direction_reference: str) -> tuple[str, str]:
    """
    Get coordinate variable names for spherical sky coordinates.

    Parameters
    ----------
    direction_reference : str
        Direction frame name.

    Returns
    -------
    tuple[str, str]
        Coordinate names for longitude and latitude.
    """
    if _is_galactic_frame(direction_reference):
        return ("galactic_longitude", "galactic_latitude")
    return ("right_ascension", "declination")


def _input_checks(
    phase_center: Union[list, np.ndarray],
    image_size: Union[list, np.ndarray],
    cell_size: Union[list, np.ndarray],
) -> None:
    """
    Validate input parameters for image creation functions.

    Parameters
    ----------
    phase_center : list or np.ndarray
        Image phase center coordinates. Must have exactly 2 elements.
    image_size : list or np.ndarray
        Number of pixels along each axis. Must have exactly 2 elements.
    cell_size : list or np.ndarray
        Size of pixels along each axis. Must have exactly 2 elements.

    Raises
    ------
    ValueError
        If any parameter does not have exactly 2 elements.

    Returns
    -------
    None
        This function validates inputs and raises on invalid shapes.
    """
    if len(image_size) != 2:
        raise ValueError("image_size must have exactly two elements")
    if len(phase_center) != 2:
        raise ValueError("phase_center must have exactly two elements")
    if len(cell_size) != 2:
        raise ValueError("cell_size must have exactly two elements")


def _make_coords(
    frequency_coords: Union[list, np.ndarray],
    time_coords: Union[list, np.ndarray],
) -> dict:
    """
    Build common time/frequency/velocity coordinate arrays.

    Parameters
    ----------
    frequency_coords : list or np.ndarray
        Frequency coordinate values in Hz.
    time_coords : list or np.ndarray
        Time coordinate values in MJD days.

    Returns
    -------
    dict
        Dictionary containing normalized coordinate arrays and a rest frequency.
    """
    if not isinstance(frequency_coords, list) and not isinstance(
        frequency_coords, np.ndarray
    ):
        frequency_coords = [frequency_coords]
    frequency_coords = np.array(frequency_coords, dtype=np.float64)
    restfreq = frequency_coords[len(frequency_coords) // 2]
    vel_coords = (1 - frequency_coords / restfreq) * _c.to("m/s").value
    if not isinstance(time_coords, list) and not isinstance(time_coords, np.ndarray):
        time_coords = [time_coords]
    time_coords = np.array(time_coords, dtype=np.float64)
    return dict(
        chan=frequency_coords, vel=vel_coords, time=time_coords, restfreq=restfreq
    )


def _add_common_attrs(
    xds: xr.Dataset,
    restfreq: float,
    spectral_reference: str,
    direction_reference: str,
    phase_center: Union[List[float], np.ndarray],
    cell_size: Union[List[float], np.ndarray],
    projection: str,
) -> xr.Dataset:
    """
    Attach common image-level coordinate attributes and metadata.

    Parameters
    ----------
    xds : xr.Dataset
        Dataset to enrich with metadata.
    restfreq : float
        Rest frequency in Hz.
    spectral_reference : str
        Spectral frame identifier.
    direction_reference : str
        Direction frame identifier.
    phase_center : list[float] or np.ndarray
        Two-element phase center in radians.
    cell_size : list[float] or np.ndarray
        Pixel cell size in radians.
    projection : str
        Projection code.

    Returns
    -------
    xr.Dataset
        Input dataset with updated coordinate attrs and dataset attrs.
    """
    xds.time.attrs = make_time_coord_attrs(units="d", scale="utc", time_format="mjd")
    freq_vals = np.array(xds.frequency)
    xds.frequency.attrs = {
        "observer": spectral_reference.lower(),
        "reference_frequency": make_spectral_coord_reference_dict(
            value=freq_vals[len(freq_vals) // 2].item(),
            units="Hz",
            observer=spectral_reference.lower(),
        ),
        "rest_frequencies": make_quantity(restfreq, "Hz"),
        "rest_frequency": make_quantity(restfreq, "Hz"),
        "type": "spectral_coord",
        "units": "Hz",
        "wave_units": "mm",
    }
    xds.velocity.attrs = {"doppler_type": "radio", "type": "doppler", "units": "m/s"}
    reference = make_skycoord_dict(
        data=phase_center,
        units="rad",
        frame=direction_reference,
    )
    if not _is_galactic_frame(direction_reference):
        reference["attrs"].update({"equinox": "j2000.0"})
    xds.attrs = {
        "data_groups": {"base": {}},
        "coordinate_system_info": {
            "reference_direction": reference,
            "native_pole_direction": make_direction_location_dict(
                [np.pi, 0.0], "rad", "native_projection"
            ),
            "pixel_coordinate_transformation_matrix": [[1.0, 0.0], [0.0, 1.0]],
            "projection": projection,
            "projection_parameters": [0.0, 0.0],
        },
        "type": "image",
    }
    return xds


def _make_common_coords(
    pol_coords: Union[list, np.ndarray],
    frequency_coords: Union[list, np.ndarray],
    time_coords: Union[list, np.ndarray],
) -> dict:
    """
    Build shared non-direction coordinates used by image constructors.

    Parameters
    ----------
    pol_coords : list or np.ndarray
        Polarization labels.
    frequency_coords : list or np.ndarray
        Frequency coordinate values in Hz.
    time_coords : list or np.ndarray
        Time coordinate values in MJD days.

    Returns
    -------
    dict
        Dictionary with assembled coordinate mapping and rest frequency.
    """
    some_coords = _make_coords(frequency_coords, time_coords)
    return {
        "coords": {
            "time": some_coords["time"],
            "frequency": some_coords["chan"],
            "velocity": ("frequency", some_coords["vel"]),
            "polarization": pol_coords,
        },
        "restfreq": some_coords["restfreq"],
    }


def _make_lm_values(
    image_size: Union[list, np.ndarray],
    cell_size: Union[list, np.ndarray],
) -> dict:
    """
    Build linear ``l`` and ``m`` coordinate arrays from image geometry.

    Parameters
    ----------
    image_size : list or np.ndarray
        Two-element image size in pixels.
    cell_size : list or np.ndarray
        Two-element cell size in radians.

    Returns
    -------
    dict
        Dictionary containing ``l`` and ``m`` coordinate arrays.
    """
    # l follows RA as far as increasing/decreasing, see AIPS Meme 27, change in alpha
    # definition three lines below Figure 2 and the first of the pair of equations 10.
    l = [
        (i - image_size[0] // 2) * (-1) * abs(cell_size[0])
        for i in range(image_size[0])
    ]
    m = [(i - image_size[1] // 2) * abs(cell_size[1]) for i in range(image_size[1])]
    return {"l": l, "m": m}


def _make_sky_coords(
    projection: str,
    image_size: Union[list, np.ndarray],
    cell_size: Union[list, np.ndarray],
    phase_center: Union[list, np.ndarray],
    direction_reference: str,
) -> dict:
    """
    Compute spherical sky-coordinate grids for the requested direction frame.

    Parameters
    ----------
    projection : str
        Spherical projection code.
    image_size : list or np.ndarray
        Two-element image size in pixels.
    cell_size : list or np.ndarray
        Two-element cell size in radians.
    phase_center : list or np.ndarray
        Two-element phase center in radians.
    direction_reference : str
        Direction frame identifier.

    Returns
    -------
    dict
        Mapping from spherical coordinate names to ``(dims, values)`` tuples.
    """
    long, lat = _compute_world_sph_dims(
        projection=projection,
        shape=image_size,
        ctype=_spherical_ctype_for_frame(direction_reference),
        crpix=[image_size[0] // 2, image_size[1] // 2],
        crval=phase_center,
        cdelt=[-abs(cell_size[0]), abs(cell_size[1])],
        cunit=["rad", "rad"],
    )["value"]
    lon_name, lat_name = _spherical_coord_names_for_frame(direction_reference)
    return {lon_name: (("l", "m"), long), lat_name: (("l", "m"), lat)}


def _add_lm_coord_attrs(xds: xr.Dataset) -> None:
    """
    Attach explanatory notes to ``l`` and ``m`` coordinates. The input Dataset is modified in-place.

    Parameters
    ----------
    xds : xr.Dataset
        Dataset containing ``l`` and ``m`` coordinates.

    Returns
    -------
    None
    """
    attr_note = _l_m_attr_notes()
    xds.l.attrs = {
        "note": attr_note["l"],
    }
    xds.m.attrs = {
        "note": attr_note["m"],
    }


def _make_empty_sky_image(
    phase_center: Union[list, np.ndarray],
    image_size: Union[list, np.ndarray],
    cell_size: Union[list, np.ndarray],
    frequency_coords: Union[list, np.ndarray],
    pol_coords: Union[list, np.ndarray],
    time_coords: Union[list, np.ndarray],
    direction_reference: str,
    projection: str,
    spectral_reference: str,
    do_sky_coords: bool,
) -> xr.Dataset:
    """
    Create an empty sky image dataset containing only coordinates.

    Parameters
    ----------
    phase_center : list or np.ndarray
        Two-element phase center in radians.
    image_size : list or np.ndarray
        Two-element image size in pixels.
    cell_size : list or np.ndarray
        Two-element cell size in radians.
    frequency_coords : list or np.ndarray
        Frequency coordinates in Hz.
    pol_coords : list or np.ndarray
        Polarization labels.
    time_coords : list or np.ndarray
        Time coordinates in MJD days.
    direction_reference : str
        Direction frame identifier.
    projection : str
        Projection code.
    spectral_reference : str
        Spectral frame identifier.
    do_sky_coords : bool
        Whether to add spherical sky-coordinate grids.

    Returns
    -------
    xr.Dataset
        Empty image dataset with coordinates and metadata.
    """
    _input_checks(phase_center, image_size, cell_size)
    cc = _make_common_coords(pol_coords, frequency_coords, time_coords)
    coords = cc["coords"]
    lm_values = _make_lm_values(image_size, cell_size)
    coords.update(lm_values)
    if do_sky_coords:
        coords.update(
            _make_sky_coords(
                projection, image_size, cell_size, phase_center, direction_reference
            )
        )
    xds = xr.Dataset(coords=coords)
    xds = _move_beam_param_dim_coord(xds)
    _add_lm_coord_attrs(xds)
    _add_common_attrs(
        xds,
        cc["restfreq"],
        spectral_reference,
        direction_reference,
        phase_center,
        cell_size,
        projection,
    )
    return xds


def _make_uv_coords(
    xds: xr.Dataset,
    image_size: Union[list, np.ndarray],
    sky_image_cell_size: Union[list, np.ndarray],
) -> xr.Dataset:
    """
    Attach ``u`` and ``v`` coordinates to a dataset.

    Parameters
    ----------
    xds : xr.Dataset
        Dataset to augment with uv coordinates.
    image_size : list or np.ndarray
        Two-element image size in pixels.
    sky_image_cell_size : list or np.ndarray
        Two-element sky image cell size in radians.

    Returns
    -------
    xr.Dataset
        Dataset with ``u`` and ``v`` coordinates and attrs.
    """
    uv_values = _make_uv_values(image_size, sky_image_cell_size)
    xds = xds.assign_coords(uv_values)
    attr = make_quantity(0.0, "lambda")
    xds.u.attrs = attr.copy()
    xds.v.attrs = attr.copy()
    return xds


def _make_uv_values(
    image_size: Union[list, np.ndarray],
    sky_image_cell_size: Union[list, np.ndarray],
) -> dict:
    """
    Compute linear ``u`` and ``v`` coordinate arrays.

    Parameters
    ----------
    image_size : list or np.ndarray
        Two-element image size in pixels.
    sky_image_cell_size : list or np.ndarray
        Two-element sky image cell size in radians.

    Returns
    -------
    dict
        Dictionary containing ``u`` and ``v`` coordinate arrays.
    """
    im_size_wave = 1 / np.array(sky_image_cell_size)
    uv_cell_size = im_size_wave / np.array(image_size)
    u_vals = [(i - image_size[0] // 2) * uv_cell_size[0] for i in range(image_size[0])]
    v_vals = [(i - image_size[1] // 2) * uv_cell_size[1] for i in range(image_size[1])]
    return {"u": u_vals, "v": v_vals}


def _make_empty_aperture_image(
    phase_center: Union[list, np.ndarray],
    image_size: Union[list, np.ndarray],
    sky_image_cell_size: Union[list, np.ndarray],
    frequency_coords: Union[list, np.ndarray],
    pol_coords: Union[list, np.ndarray],
    time_coords: Union[list, np.ndarray],
    direction_reference: str,
    projection: str,
    spectral_reference: str,
) -> xr.Dataset:
    """
    Create an empty aperture image dataset containing only coordinates.

    Parameters
    ----------
    phase_center : list or np.ndarray
        Two-element phase center in radians.
    image_size : list or np.ndarray
        Two-element image size in pixels.
    sky_image_cell_size : list or np.ndarray
        Two-element sky image cell size in radians.
    frequency_coords : list or np.ndarray
        Frequency coordinates in Hz.
    pol_coords : list or np.ndarray
        Polarization labels.
    time_coords : list or np.ndarray
        Time coordinates in MJD days.
    direction_reference : str
        Direction frame identifier.
    projection : str
        Projection code.
    spectral_reference : str
        Spectral frame identifier.

    Returns
    -------
    xr.Dataset
        Empty aperture-image dataset with coordinates and metadata.
    """
    _input_checks(phase_center, image_size, sky_image_cell_size)
    cc = _make_common_coords(pol_coords, frequency_coords, time_coords)
    coords = cc["coords"]
    xds = xr.Dataset(coords=coords)
    xds = _make_uv_coords(xds, image_size, sky_image_cell_size)
    _add_common_attrs(
        xds,
        cc["restfreq"],
        spectral_reference,
        direction_reference,
        phase_center,
        sky_image_cell_size,
        projection,
    )
    xds = _move_beam_param_dim_coord(xds)
    return xds


def _move_beam_param_dim_coord(xds: xr.Dataset) -> xr.Dataset:
    """
    Add beam_params_label coordinate to an xarray Dataset.

    Parameters
    ----------
    xds : xr.Dataset
        Input Dataset to which beam parameters will be added.

    Returns
    -------
    xr.Dataset
        Dataset with beam_params_label coordinate containing ['major', 'minor', 'pa'].
    """
    return xds.assign_coords(
        beam_params_label=("beam_params_label", ["major", "minor", "pa"])
    )


def _make_empty_lmuv_image(
    phase_center: Union[list, np.ndarray],
    image_size: Union[list, np.ndarray],
    sky_image_cell_size: Union[list, np.ndarray],
    frequency_coords: Union[list, np.ndarray],
    pol_coords: Union[list, np.ndarray],
    time_coords: Union[list, np.ndarray],
    direction_reference: str,
    projection: str,
    spectral_reference: str,
    do_sky_coords: bool,
) -> xr.Dataset:
    """
    Create an empty image dataset with both lm and uv coordinates.

    Parameters
    ----------
    phase_center : list or np.ndarray
        Two-element phase center in radians.
    image_size : list or np.ndarray
        Two-element image size in pixels.
    sky_image_cell_size : list or np.ndarray
        Two-element sky image cell size in radians.
    frequency_coords : list or np.ndarray
        Frequency coordinates in Hz.
    pol_coords : list or np.ndarray
        Polarization labels.
    time_coords : list or np.ndarray
        Time coordinates in MJD days.
    direction_reference : str
        Direction frame identifier.
    projection : str
        Projection code.
    spectral_reference : str
        Spectral frame identifier.
    do_sky_coords : bool
        Whether to include spherical sky-coordinate grids.

    Returns
    -------
    xr.Dataset
        Empty image dataset with lm and uv coordinate systems.
    """
    xds = _make_empty_sky_image(
        phase_center,
        image_size,
        sky_image_cell_size,
        frequency_coords,
        pol_coords,
        time_coords,
        direction_reference,
        projection,
        spectral_reference,
        do_sky_coords,
    )
    xds = _make_uv_coords(xds, image_size, sky_image_cell_size)
    xds = _move_beam_param_dim_coord(xds)
    return xds


def detect_store_type(store):
    """
    Detect the storage format type of an image store.

    Parameters
    ----------
    store : str or dict
        Path to the image store or a dictionary representation.

    Returns
    -------
    str
        The detected store type: 'fits', 'casa', or 'zarr'.

    Raises
    ------
    ValueError
        If the directory structure is unknown or the path does not exist.
    """
    import os

    if isinstance(store, str):
        if os.path.isfile(store):
            store_type = "fits"
        elif os.path.isdir(store):
            if "table.info" in os.listdir(store):
                store_type = "casa"
            elif ".zattrs" in os.listdir(store):
                store_type = "zarr"
            else:
                xradio_logger().error("Unknown directory structure.")
                raise ValueError("Unknown directory structure." + str(store))
        else:
            xradio_logger().error("Path does not exist.")
            raise ValueError(
                "Path does not exist. The current path: "
                + str(os.system("pwd"))
                + " .The current casa directory: "
                + str(os.system("ls 3c286_Band6_5chans_lsrk_robust_0.5_niter_99_casa"))
                + ". The current fits directory: "
                + str(os.system("ls 3c286_Band6_5chans_lsrk_robust_0.5_niter_99_fits"))
                + " The given store "
                + str(store)
            )
    else:
        store_type = "zarr"

    return store_type


def detect_image_type(store):
    """
    Detect the image type from the store name or path.

    Infers the image type based on common naming patterns in the store path.

    Parameters
    ----------
    store : str or other
        Path to the image store. If not a string, returns 'ALL'.

    Returns
    -------
    str
        The detected image type. Possible values include:
        - 'SKY': Sky image (contains 'image' or 'im')
        - 'POINT_SPREAD_FUNCTION': PSF image (contains 'psf')
        - 'MODEL': Model image (contains 'model')
        - 'RESIDUAL': Residual image (contains 'residual')
        - 'MASK_DECONVOLVE': Mask image (contains 'mask')
        - 'PRIMARY_BEAM': Primary beam image (contains 'pb')
        - 'APERTURE': Aperture image (contains 'aperture')
        - 'VISIBILITY': Visibility image (contains 'visibility')
        - 'VISIBILITY_NORMALIZATION': Visibility normalization (contains 'sumwt')
        - 'UNKNOWN': Could not detect type from name
        - 'ALL': Non-string store type
    """
    import os

    if isinstance(store, str):
        if "fits" in store.lower():
            image_type = "SKY"
        elif "image" in store.lower():
            image_type = "SKY"
        elif "sky" in store.lower():
            image_type = "SKY"
        elif "psf" in store.lower():
            image_type = "POINT_SPREAD_FUNCTION"
        elif "model" in store.lower():
            image_type = "MODEL"
        elif "residual" in store.lower():
            image_type = "RESIDUAL"
        elif "dirty" in store.lower():
            image_type = "DIRTY"
        elif "pb" in store.lower():
            image_type = "PRIMARY_BEAM"
        elif "aperture" in store.lower():
            image_type = "APERTURE"
        elif "uv" in store.lower():
            # Must precede the "im" check so uv images aren't misclassified as SKY.
            image_type = "APERTURE"
        elif "visibility" in store.lower():
            image_type = "VISIBILITY"
        elif "sumwt" in store.lower():
            image_type = "VISIBILITY_NORMALIZATION"
        elif "zarr" in store.lower():
            image_type = "ALL"
        elif "im" in store.lower():
            # Must precede the "mask" check so *mask*.im images aren't misclassified.
            image_type = "SKY"
        elif "mask" in store.lower():
            image_type = "MASK_DECONVOLVE"
        else:
            image_type = "UNKNOWN"
    else:
        image_type = "ALL"

    return image_type


def create_store_dict(store_to_label):
    """
    Create a standardized dictionary mapping image types to their store information.

    Converts various input formats (string, list, or dict) into a consistent
    dictionary format where keys are image types and values contain store metadata.

    Parameters
    ----------
    store_to_label : str, list, or dict
        Input store specification. Can be:
        - str: Single store path
        - list: List of store paths (image types will be auto-detected)
        - dict: Mapping of image types to store paths

    Returns
    -------
    dict
        Dictionary with image types as keys. Each value is a dict with:
        - 'store_type': str, the format ('casa', 'fits', or 'zarr')
        - 'store': str, the path to the store

    Raises
    ------
    ValueError
        If image type cannot be detected or duplicate image types are found.
    """
    store_list = None
    if isinstance(store_to_label, str):
        store_list = [store_to_label]  # So can iterate over it.
    elif isinstance(store_to_label, list):
        store_list = store_to_label

    if (store_list is not None) and isinstance(store_list, list):
        store_dict_to_label = {i: v for i, v in enumerate(store_list)}
    else:
        store_dict_to_label = store_to_label

    store_dict = {}
    for image_type, store in store_dict_to_label.items():

        if isinstance(image_type, int):
            image_type = detect_image_type(store)

        image_type = image_type.upper()

        store_type = detect_store_type(store)

        if image_type == "UNKNOWN":
            xradio_logger().error(f"Could not detect image type for store {store}. ")
            example = "store={'sky': 'path/to/image.fits'}"
            raise ValueError(
                f"Could not detect image type for store {store}. Please label the store with the image type explicitly. For example: {example}"
            )

        if image_type in store_dict:
            xradio_logger().error(
                f"Duplicate image type {image_type} detected in store list."
            )
            raise ValueError(
                f"Duplicate image type {image_type} detected in store list. Please ensure each image type is unique. The store dict"
                + str(store_dict)
            )

        if store_type == "zarr":
            image_type = "ALL"  # Zarr can have multiple data variables.

        store_dict[image_type] = {"store_type": store_type, "store": store}

    data_groups = {}

    for image_type in store_dict.keys():
        if "sky" in image_type.lower():
            if "sky" == image_type.lower():
                data_groups["base"] = {"sky": image_type}
            else:
                data_group_name = image_type.lower().replace("sky_", "")
                data_groups[data_group_name] = {"sky": image_type}
        if "aperture" == image_type.lower():
            data_groups["base"] = {"aperture": image_type}

    return store_dict, data_groups


def create_image_xds_from_store(
    store: Union[list, dict, str],
    access_store_casa: callable,
    casa_kwargs: dict,
    access_store_fits: callable,
    fits_kwargs: dict,
    access_store_zarr: callable,
    zarr_kwargs: dict,
) -> xr.Dataset:
    """
    Create an xarray Dataset from one or more image stores.

    This function reads image data from CASA, FITS, or zarr format stores and
    combines them into a single xarray Dataset with appropriate metadata and
    data variables.

    Parameters
    ----------
    store : str, list, or dict
        Image store specification:
        - str: Single store path
        - list: List of store paths
        - dict: Mapping of image types to store paths
    access_store_casa : callable
        Function to read CASA format images. Should accept a store path and
        keyword arguments, returning an xr.Dataset.
    casa_kwargs : dict
        Keyword arguments to pass to access_store_casa.
    access_store_fits : callable or None
        Function to read FITS format images. Should accept a store path and
        keyword arguments, returning an xr.Dataset. Can be None if FITS support
        is not needed.
    fits_kwargs : dict
        Keyword arguments to pass to access_store_fits.
    access_store_zarr : callable
        Function to read zarr format images. Should accept a store path and
        keyword arguments, returning an xr.Dataset.
    zarr_kwargs : dict
        Keyword arguments to pass to access_store_zarr.

    Returns
    -------
    xr.Dataset
        An xarray Dataset containing the image data and metadata. The Dataset
        includes:
        - Data variables for each image type (e.g., 'SKY', 'MODEL', 'RESIDUAL')
        - Coordinates shared across all images
        - Attributes including 'type' and 'data_groups'

    Raises
    ------
    ValueError
        If zarr store with multiple data variables is combined with other stores.
    RuntimeError
        If FITS format is requested but access_store_fits is None, or if an
        unrecognized image format is encountered.

    Notes
    -----
    - Zarr stores can contain multiple data variables and will be returned as-is.
    - For other formats, data from multiple stores is combined into one Dataset.
    - BEAM_FIT_PARAMS from SKY images take precedence over POINT_SPREAD_FUNCTION.
    - Masks are renamed to MASK_<IMAGE_TYPE> for internal masks.
    """
    store_dict, data_groups = create_store_dict(store)
    if "ALL" in store_dict and len(store_dict) > 1:
        xradio_logger().error(
            "When using a zarr store with multiple data variables, no other stores can be specified."
        )
        raise ValueError(
            "When using a zarr store with multiple data variables, no other stores can be specified."
        )

    if "ALL" in store_dict:
        zarr_store = store if isinstance(store, str) else list(store.values())[0]
        img_xds = access_store_zarr(zarr_store, **zarr_kwargs)
        return img_xds

    img_xds = xr.Dataset()
    # Loop over all the input CASA and Fits images.
    for image_type, store_description in store_dict.items():

        store_type = store_description["store_type"]
        store = store_description["store"]

        fits_kwargs["image_type"] = image_type
        casa_kwargs["image_type"] = image_type

        if store_type == "casa":
            xds = access_store_casa(store, **casa_kwargs)
        elif store_type == "fits":
            if access_store_fits is None:
                xradio_logger().error("FITS not currently supported.")
                raise RuntimeError("FITS not currently supported.")
            xds = access_store_fits(store, **fits_kwargs)
        else:
            xradio_logger().error(
                f"Unrecognized image format for path {store}. Supported types are CASA, FITS, and zarr.\n"
            )
            raise RuntimeError(
                f"Unrecognized image format for path {store}. Supported types are CASA, FITS, and zarr.\n"
            )

        img_xds.attrs = img_xds.attrs | xds.attrs
        img_xds[image_type] = xds[image_type]
        img_xds[image_type].attrs["type"] = image_type.lower()

        active_data_group_name = None
        # If sky image, handle internal masks and beam fit params.
        if "sky" in image_type.lower():
            for data_group_name, data_group in data_groups.items():
                if data_group["sky"] == image_type:
                    active_data_group_name = data_group_name

            if "BEAM_FIT_PARAMS_" + image_type.upper() in xds:
                img_xds["BEAM_FIT_PARAMS_" + image_type.upper()] = xds[
                    "BEAM_FIT_PARAMS_" + image_type.upper()
                ]
                data_groups[active_data_group_name]["beam_fit_params_sky"] = (
                    "BEAM_FIT_PARAMS_" + image_type.upper()
                )
            expected_flag_name = "FLAG_" + image_type

            # TODO remove this mask logic and everything that still makes it necessary
            def _add_flag_to_group(
                img_xds: xr.Dataset,
                flag_array: xr.DataArray,
                expected_flag_name: str,
                active_group: dict,
            ):
                img_xds[expected_flag_name] = flag_array
                img_xds[expected_flag_name].attrs["type"] = "flag"
                active_group["flag"] = expected_flag_name

            if expected_flag_name in xds:
                _add_flag_to_group(
                    img_xds,
                    xds[expected_flag_name],
                    expected_flag_name,
                    data_groups[active_data_group_name],
                )

            if "MASK_0" in xds:
                _add_flag_to_group(
                    img_xds,
                    xds["MASK_0"],
                    expected_flag_name,
                    data_groups[active_data_group_name],
                )
                """
                TODO delete old code when certain new function works
                img_xds[expected_flag_name] = xds["MASK_0"]
                data_groups[active_data_group_name]["flag"] = expected_flag_name
                img_xds[expected_flag_name].attrs["type"] = "flag"
                """
            if "MASK" in xds:
                _add_flag_to_group(
                    img_xds,
                    xds["MASK"],
                    expected_flag_name,
                    data_groups[active_data_group_name],
                )
                """
                TODO delete old code when certain new function works
                img_xds[expected_flag_name] = xds["MASK"]
                data_groups[active_data_group_name]["flag"] = expected_flag_name
                img_xds[expected_flag_name].attrs["type"] = "flag"
                """
            img_xds[image_type].attrs["type"] = "sky"

        # If point spread function, handle beam fit params.
        if "point_spread_function" in image_type.lower():
            if "BEAM_FIT_PARAMS_" + image_type.upper() in xds:
                img_xds["BEAM_FIT_PARAMS_" + image_type.upper()] = xds[
                    "BEAM_FIT_PARAMS_" + image_type.upper()
                ]

        # Figure out data groups.
        # Each sky image gets its own data group and shares all other images between them.
        if "sky" not in image_type.lower():
            for data_group_name, data_group in data_groups.items():
                data_group[image_type.lower()] = image_type

                if "point_spread_function" in image_type.lower():
                    if "BEAM_FIT_PARAMS_" + image_type.upper() in xds:
                        data_group["beam_fit_params_point_spread_function"] = (
                            "BEAM_FIT_PARAMS_" + image_type.upper()
                        )
    if (
        "visibility_normalization" not in image_type.lower()
        or len(img_xds.data_vars) > 1
    ):
        # if beam_param coord not in image type it is not auto assigned to img_xds
        # but it must be present even if unused
        if "beam_params_label" not in img_xds.dims:
            img_xds.expand_dims(beam_params_label=3)

        if "beam_params_label" not in img_xds.coords:
            img_xds = _move_beam_param_dim_coord(img_xds)
    img_xds.attrs["type"] = "image_dataset"
    img_xds.attrs["data_groups"] = data_groups
    return img_xds
