"""
Specification for various utility types used throughout tools and CryoSPARC.
"""

from typing import Dict, List, Literal, Tuple, TypedDict, Union

Shape = Tuple[int, ...]
"""A numpy shape tuple from ndarray.shape"""

DType = Union[str, Tuple[str, Shape]]
"""

    Can just be a single string such as "f4", "3u4" or "O".
    A datatype description of a ndarray entry.

    Can also be the a tuple with a string datatype name and its shape. For
    example, the following dtypes are equivalent.

    - "3u4"
    - "<u4", (3,))
"""

Field = Union[Tuple[str, str], Tuple[str, str, Shape]]
"""
    Description of a column in a numpy array with named fields

    Examples:
    - ("uid", "u8")
    - ("coords", "3f4")
    - ("coords", "<f4", (3,))
"""

Datatype = Literal[
    "exposure",
    "particle",
    "template",
    "volume",
    "volume_multi",
    "mask",
    "live",
    "ml_model",
    "symmetry_candidate",
    "flex_mesh",
    "flex_model",
    "hyperparameter",
    "denoise_model",
    "annotation_model",
]
"""Supported data types for job inputs and outputs."""


LoadableSlots = Union[Literal["default", "passthrough", "all"], List[str]]
"""
Types of slots that may be included when loading an input or an output.
"default" means only slots required as inputs or newly-generated by outputs.
"passthrough" means only unused passthroughs for inputs/outputs. "all" means the
previous two combined. list means a specific list of slots.
"""


# Valid plot file types
TextFormat = Literal["txt", "csv", "html", "json", "xml", "bild", "bld", "log"]
"""
Supported job stream log asset file text formats.
"""

ImageFormat = Literal["pdf", "gif", "jpg", "jpeg", "png", "svg"]
"""
Supported job stream log asset file image formats.
"""

MicrographFormat = Literal["mrc", "mrcs", "stk", "tif", "tiff", "eer", "ecc", "bz2", "cmrcbz2"]
"""Supported micrograph extensions."""

DatasetFormat = Literal["cs", "npy"]
"""Supported extensions for dataset files."""

AssetFormat = Union[TextFormat, ImageFormat]
"""
Supported job stream log asset file formats.
"""

AnyFormat = Union[TextFormat, ImageFormat, MicrographFormat, DatasetFormat]
"""
Supported file extensions.
"""

TextContentType = Literal[
    "text/plain",
    "text/csv",
    "text/html",
    "application/json",
    "application/xml",
    "application/x-troff",
]
"""
Supported job stream log text asset MIME types.
"""

ImageContentType = Literal[
    "application/pdf",
    "image/gif",
    "image/jpeg",
    "image/png",
    "image/svg+xml",
]
"""
Supported job image asset MIME types renderable in the web
"""

MicrographContentType = Literal[
    "image/x-mrc",
    "image/tiff",
    "application/x-eer",
    "application/x-bzip2",
]
"""Supported micrograph content types."""

DatasetContentType = Literal["application/x-cryosparc-dataset", "application/x-numpy"]
"""Supported content types for dataset files."""

AssetContentType = Union[TextContentType, ImageContentType, Literal["application/octet-stream"]]
"""
Supported job asset MIME types.
"""

AnyContentType = Union[
    TextContentType,
    ImageContentType,
    MicrographContentType,
    DatasetContentType,
    Literal["application/octet-stream"],
]
"""
All known and supported content types used by CryoSPARC.
"""

TEXT_CONTENT_TYPES: Dict[TextFormat, TextContentType] = {
    "log": "text/plain",
    "txt": "text/plain",
    "csv": "text/csv",
    "html": "text/html",
    "json": "application/json",
    "xml": "application/xml",
    "bild": "application/x-troff",
    "bld": "application/x-troff",
}

IMAGE_CONTENT_TYPES: Dict[ImageFormat, ImageContentType] = {
    "pdf": "application/pdf",
    "gif": "image/gif",
    "jpeg": "image/jpeg",
    "jpg": "image/jpeg",
    "png": "image/png",
    "svg": "image/svg+xml",
}

MICROGRAPH_CONTENT_TYPES: Dict[MicrographFormat, MicrographContentType] = {
    "mrc": "image/x-mrc",
    "mrcs": "image/x-mrc",
    "stk": "image/x-mrc",
    "tif": "image/tiff",
    "tiff": "image/tiff",
    "eer": "application/x-eer",
    "ecc": "application/x-eer",
    "bz2": "application/x-bzip2",
    "cmrcbz2": "application/x-bzip2",
}

DATASET_CONTENT_TYPES: Dict[DatasetFormat, DatasetContentType] = {
    "cs": "application/x-cryosparc-dataset",  # but usually has numpy format internally
    "npy": "application/x-numpy",
}

ASSET_CONTENT_TYPES: Dict[AssetFormat, AssetContentType] = {**TEXT_CONTENT_TYPES, **IMAGE_CONTENT_TYPES}
ASSET_EXTENSIONS: Dict[AssetContentType, AssetFormat] = {v: k for k, v in ASSET_CONTENT_TYPES.items()}

KNOWN_CONTENT_TYPES: Dict[AnyFormat, AnyContentType] = {
    **TEXT_CONTENT_TYPES,
    **IMAGE_CONTENT_TYPES,
    **MICROGRAPH_CONTENT_TYPES,
    **DATASET_CONTENT_TYPES,
}


# Slot is defined in two classes like this because it's the only way to
# make the ``required`` key optional.
class _Slot(TypedDict):
    """
    :meta private:
    """

    name: str
    dtype: str


class Slot(_Slot, total=False):
    """
    Full slot dictionary specification type for items in the slots=... argument
    when creating inputs or outputs. e.g., ``{"name": "ctf", "dtype": "ctf"}``
    or ``{"name": "background_blob", "dtype": "stat_blob", "required": False}``

    See :py:type:`SlotSpec` for details.

    Attributes:
        name (str): where to find field in a corresponding .cs file e.g.,
            ``"background_blob"``, ``"ctf"``, ``"alignments_class_0"``
        dtype (str): name of known data type. e.g., ``"stat_blob"``, ``"ctf"``,
            ``"alignments3D"``.
        required (bool, optional): Whether this slot is required. Applies to
            input specs only. Defaults to True.

    """

    required: bool


class Datafield(TypedDict):
    """
    Deprecated. Use :py:class:`Slot` instead.
    """

    dtype: str
    prefix: str
    required: bool


SlotSpec = Union[str, Slot, Datafield]
"""
A result slot specification for items in the slots=... argument when creating
inputs or outputs.

In CryoSPARC, all jobs have one or more inputs and outputs. An input or output
has some broad :py:type:`Datatype`, such as ``"exposure"`` or ``"particle"``.
Each input or output also has a list of associated "low-level" results created
at various stages of processing, such as ``"location"`` for picked particles and
``blob`` for extracted particles. A slot represents one of these low-level
results.

In the CryoSPARC interface, open a job's "Inputs" or "Outputs" tab to see the
kinds of slots available. You may also download an output and load it with
:py:type:`~cryosparc.dataset.Dataset` to inspect the infomation encoded in its
results.

Provide each slot as either a string representing a name and result type, or a
full dictionary specification.

A string in the format ``"<slot>"`` is a shortcut for ``{"name": "<slot>",
"dtype": "<known slot type>", "required": True}``.

A string in the format ``"?<slot>"`` is a shortcut for ``{"name": "<slot>",
"dtype": "<known slot type>", "required": False}`` (input slots only).

Example strings::

    "ctf"
    "micrograph_blob"
    "?background_blob"

Example equivalent full specifications::

    {"name": "ctf", "dtype": "ctf"}
    {"name": "micrograph_blob", "dtype": "micrograph_blob", "required": True}
    {"name": "background_blob", "dtype": "stat_blob", "required": False}

Use the full specification when the ``dtype`` cannot be inferred from the
``name`` string because it is dynamic. For example, 3D Variability job
``particles`` outputs have slots named ``"components_mode_X"`` with dtype
``"components"`` where ``X`` is a mode number::

    [
        "blob",
        "?locations",
        {"name": "components_mode_0", "dtype": "components"},
        {"name": "components_mode_1", "dtype": "components", "required": False},
        {"name": "components_mode_2", "dtype": "components", "required": False},
    ]

Note that the ``required`` key only applies to input slots.
"""


class JobSection(TypedDict):
    """
    Deprecated. Use :py:class:`~cryosparc.models.job_register.JobRegister`
    instead.
    """

    name: str
    title: str
    description: str
    contains: List[str]
