"""Fractal internal module for axes handling."""

from collections.abc import Sequence
from enum import Enum
from typing import Literal, TypeAlias, TypeVar

from pydantic import BaseModel, ConfigDict, Field, model_validator

from ngio.utils import NgioValidationError, NgioValueError

T = TypeVar("T")
SlicingType: TypeAlias = slice | tuple[int, ...] | int

################################################################################################
#
# Axis Types and Units
# We define a small set of axis types and units that can be used in the metadata.
# This axis types are more restrictive than the OME standard.
# We do that to simplify the data processing.
#
#################################################################################################


class AxisType(str, Enum):
    """Allowed axis types."""

    channel = "channel"
    time = "time"
    space = "space"


SpaceUnits = Literal[
    "micrometer",
    "nanometer",
    "angstrom",
    "picometer",
    "millimeter",
    "centimeter",
    "decimeter",
    "meter",
    "inch",
    "foot",
    "yard",
    "mile",
    "kilometer",
    "hectometer",
    "megameter",
    "gigameter",
    "terameter",
    "petameter",
    "exameter",
    "parsec",
    "femtometer",
    "attometer",
    "zeptometer",
    "yoctometer",
    "zettameter",
    "yottameter",
]
DefaultSpaceUnit = "micrometer"

TimeUnits = Literal[
    "attosecond",
    "centisecond",
    "day",
    "decisecond",
    "exasecond",
    "femtosecond",
    "gigasecond",
    "hectosecond",
    "hour",
    "kilosecond",
    "megasecond",
    "microsecond",
    "millisecond",
    "minute",
    "nanosecond",
    "petasecond",
    "picosecond",
    "second",
    "terasecond",
    "yoctosecond",
    "yottasecond",
    "zeptosecond",
    "zettasecond",
]
DefaultTimeUnit = "second"


class Axis(BaseModel):
    """Axis infos model."""

    name: str
    unit: str | None = None
    axis_type: AxisType | None = None

    model_config = ConfigDict(extra="forbid", frozen=True)

    def implicit_type_cast(self, cast_type: AxisType) -> "Axis":
        unit = self.unit
        if cast_type == AxisType.time and unit is None:
            unit = DefaultTimeUnit

        if cast_type == AxisType.space and unit is None:
            unit = DefaultSpaceUnit

        return Axis(name=self.name, axis_type=cast_type, unit=unit)

    def canonical_axis_cast(self, canonical_name: str) -> "Axis":
        """Cast the implicit axis to the correct type."""
        match canonical_name:
            case "t":
                if self.axis_type != AxisType.time or self.unit is None:
                    return self.implicit_type_cast(AxisType.time)
            case "c":
                if self.axis_type != AxisType.channel:
                    return self.implicit_type_cast(AxisType.channel)
            case "z" | "y" | "x":
                if self.axis_type != AxisType.space or self.unit is None:
                    return self.implicit_type_cast(AxisType.space)
        return self


################################################################################################
#
# Axes Handling
# We define a unique mapping to match the axes on disk to the canonical axes.
# The canonical axes are the ones that are used consistently in the NGIO internal API.
# The canonical axes ordered are: t, c, z, y, x.
#
#################################################################################################


def canonical_axes_order() -> tuple[str, str, str, str, str]:
    """Get the canonical axes order."""
    return "t", "c", "z", "y", "x"


def canonical_label_axes_order() -> tuple[str, str, str, str]:
    """Get the canonical axes order."""
    return "t", "z", "y", "x"


class AxesSetup(BaseModel):
    """Axes setup model.

    This model is used to map the on disk axes to the canonical OME-Zarr axes.
    """

    x: str = "x"
    y: str = "y"
    z: str = "z"
    c: str = "c"
    t: str = "t"
    others: list[str] = Field(default_factory=list)
    allow_non_canonical_axes: bool = False
    strict_canonical_order: bool = False

    model_config = ConfigDict(extra="forbid", frozen=True)

    @model_validator(mode="after")
    def _validate_axes_values(self) -> "AxesSetup":
        """Validate the axes values."""
        canonical = {"x", "y", "z", "c", "t"}
        axes = {"x": self.x, "y": self.y, "z": self.z, "c": self.c, "t": self.t}

        for axis_name, axis_value in axes.items():
            reserved = canonical - {axis_name}
            if axis_value in reserved:
                raise NgioValueError(
                    f"The {axis_name} axis cannot be called: '{axis_value}'. "
                    f"{axis_value} is reserved. If you want to set a non canonical "
                    "axis order, please set the 'strict_canonical_order'to False."
                )
        return self

    @classmethod
    def from_ordered_list(
        cls, axes_names: Sequence[str], canonical_order: Sequence[str]
    ):
        """Create an AxesSetup from an ordered list of axes names."""
        # Make sure to only keep as many default axes as provided in axes_names
        if len(axes_names) > len(canonical_order):
            raise NgioValueError(
                f"Cannot create AxesSetup from axes names {axes_names} "
                f"and canonical order {canonical_order}. "
                "The number of axes names cannot be greater than the "
                "number of canonical axes."
            )
        canonical_order = list(canonical_order)
        chanonical_axes = canonical_axes_order()
        axes_mapping = {}
        for ax in reversed(axes_names):
            c_ax = canonical_order.pop()
            if ax in chanonical_axes:
                axes_mapping[ax] = ax
            else:
                axes_mapping[c_ax] = ax
        return cls(**axes_mapping)

    def canonical_map(self) -> dict[str, str]:
        """Get the canonical map of axes."""
        return {
            "t": self.t,
            "c": self.c,
            "z": self.z,
            "y": self.y,
            "x": self.x,
        }

    def get_on_disk_name(self, canonical_name: str) -> str | None:
        """Get the on disk name of the axis by its canonical name."""
        canonical_map = self.canonical_map()
        return canonical_map.get(canonical_name, None)

    def inverse_canonical_map(self) -> dict[str, str]:
        """Get the on disk map of axes."""
        return {
            self.t: "t",
            self.c: "c",
            self.z: "z",
            self.y: "y",
            self.x: "x",
        }

    def get_canonical_name(self, on_disk_name: str) -> str | None:
        """Get the canonical name of the axis by its on disk name."""
        inv_map = self.inverse_canonical_map()
        return inv_map.get(on_disk_name, None)


def _check_unique_names(axes: Sequence[Axis]):
    """Check if all axes on disk have unique names."""
    names = [ax.name for ax in axes]
    if len(set(names)) != len(names):
        duplicates = {item for item in names if names.count(item) > 1}
        raise NgioValidationError(
            f"All axes must be unique. But found duplicates axes {duplicates}"
        )


def _check_non_canonical_axes(axes_setup: AxesSetup):
    """Check if all axes are known."""
    if not axes_setup.allow_non_canonical_axes and len(axes_setup.others) > 0:
        raise NgioValidationError(
            f"Unknown axes {axes_setup.others}. Please set "
            "`allow_non_canonical_axes=True` to ignore them"
        )


def _check_axes_validity(axes: Sequence[Axis], axes_setup: AxesSetup):
    """Check if all axes are valid."""
    _axes_setup = axes_setup.model_dump(exclude={"others"})
    _all_known_axes = [*_axes_setup.values(), *axes_setup.others]
    for ax in axes:
        if ax.name not in _all_known_axes:
            raise NgioValidationError(
                f"Invalid axis name '{ax.name}'. "
                f"Please correct map `{ax.name}` "
                f"using the AxesSetup model {axes_setup}"
            )


def _check_canonical_order(axes: Sequence[Axis], axes_setup: AxesSetup):
    """Check if the axes are in the canonical order."""
    if not axes_setup.strict_canonical_order:
        return
    _names = [ax.name for ax in axes]
    _canonical_order = []
    for name in canonical_axes_order():
        mapped_name = getattr(axes_setup, name)
        if mapped_name in _names:
            _canonical_order.append(mapped_name)

    if _names != _canonical_order:
        raise NgioValidationError(
            f"Invalid axes order. The axes must be in the canonical order. "
            f"Expected {_canonical_order}, but found {_names}"
        )


def validate_axes(
    axes: Sequence[Axis],
    axes_setup: AxesSetup,
) -> None:
    """Validate the axes."""
    if axes_setup.allow_non_canonical_axes and axes_setup.strict_canonical_order:
        raise NgioValueError(
            "`allow_non_canonical_axes` and"
            "`strict_canonical_order` cannot be true at the same time."
            "If non canonical axes are allowed, the order cannot be checked."
        )
    _check_unique_names(axes=axes)
    _check_non_canonical_axes(axes_setup=axes_setup)
    _check_axes_validity(axes=axes, axes_setup=axes_setup)
    _check_canonical_order(axes=axes, axes_setup=axes_setup)


class AxesHandler:
    """This class is used to handle and operate on OME-Zarr axes.

    The class also provides:
        - methods to reorder, squeeze and expand axes.
        - methods to validate the axes.
        - methods to get axis by name or index.
        - methods to operate on the axes.
    """

    def __init__(
        self,
        # spec dictated args
        axes: Sequence[Axis],
        # user defined args
        axes_setup: AxesSetup | None = None,
    ):
        """Create a new AxesMapper object.

        Args:
            axes (list[Axis]): The axes on disk.
            axes_setup (AxesSetup, optional): The axis setup. Defaults to None.
        """
        axes_setup = axes_setup if axes_setup is not None else AxesSetup()

        validate_axes(
            axes=axes,
            axes_setup=axes_setup,
        )

        self._allow_non_canonical_axes = axes_setup.allow_non_canonical_axes
        self._strict_canonical_order = axes_setup.strict_canonical_order

        self._canonical_order = canonical_axes_order()

        self._axes = axes
        self._axes_setup = axes_setup

        self._index_mapping = self._compute_index_mapping()

        # Validate the axes type and cast them if necessary
        # This needs to be done after the name mapping is computed
        self.validate_axes_type()

    def _compute_index_mapping(self):
        """Compute the index mapping.

        The index mapping is a dictionary with keys the canonical axes names
        and values the on disk axes index.

        Example:
            If the on disk axes are ['channel', 't', 'z', 'y', 'x'],
            the index mapping will be:
            {
                'c': 0,
                'channel': 0,
                't': 1,
                'z': 2,
                'y': 3,
                'x': 4,
            }
        """
        _index_mapping = {}
        for i, ax in enumerate(self.axes_names):
            _index_mapping[ax] = i
        # If the axis is not in the canonical order we also set it.
        canonical_map = self._axes_setup.canonical_map()
        for canonical_name, on_disk_name in canonical_map.items():
            if on_disk_name in _index_mapping.keys():
                _index_mapping[canonical_name] = _index_mapping[on_disk_name]
        return _index_mapping

    @property
    def axes_setup(self) -> AxesSetup:
        """Return the axes setup."""
        return self._axes_setup

    @property
    def axes(self) -> tuple[Axis, ...]:
        return tuple(self._axes)

    @property
    def axes_names(self) -> tuple[str, ...]:
        """On disk axes names."""
        return tuple(ax.name for ax in self._axes)

    @property
    def allow_non_canonical_axes(self) -> bool:
        """Return if non canonical axes are allowed."""
        return self._allow_non_canonical_axes

    @property
    def strict_canonical_order(self) -> bool:
        """Return if strict canonical order is enforced."""
        return self._strict_canonical_order

    @property
    def space_unit(self) -> str | None:
        """Return the space unit for a given axis."""
        x_axis = self.get_axis("x")
        y_axis = self.get_axis("y")

        if x_axis is None or y_axis is None:
            raise NgioValidationError(
                "The dataset must have x and y axes to determine the space unit."
            )

        if x_axis.unit == y_axis.unit:
            return x_axis.unit
        else:
            raise NgioValidationError(
                "Inconsistent space units. "
                f"x={x_axis.unit} and y={y_axis.unit} should have the same unit."
            )

    @property
    def time_unit(self) -> str | None:
        """Return the time unit for a given axis."""
        t_axis = self.get_axis("t")
        if t_axis is None:
            return None
        return t_axis.unit

    def to_units(
        self,
        *,
        space_unit: SpaceUnits = DefaultSpaceUnit,
        time_unit: TimeUnits = DefaultTimeUnit,
    ) -> "AxesHandler":
        """Convert the pixel size to the given units.

        Args:
            space_unit(str): The space unit to convert to.
            time_unit(str): The time unit to convert to.
        """
        new_axes = []
        for ax in self.axes:
            if ax.axis_type == AxisType.space:
                new_ax = Axis(
                    name=ax.name,
                    axis_type=ax.axis_type,
                    unit=space_unit,
                )
                new_axes.append(new_ax)
            elif ax.axis_type == AxisType.time:
                new_ax = Axis(name=ax.name, axis_type=ax.axis_type, unit=time_unit)
                new_axes.append(new_ax)
            else:
                new_axes.append(ax)

        return AxesHandler(
            axes=new_axes,
            axes_setup=self.axes_setup,
        )

    def rename_axes(self, axes_names: Sequence[str]) -> "AxesHandler":
        """Rename the axes.

        Args:
            axes_names (Sequence[str]): The new axes names.
        """
        if len(axes_names) != len(self.axes):
            raise NgioValueError(
                f"Cannot rename axes. "
                f"Expected {len(self.axes)} axes, but got {len(axes_names)}."
            )
        new_axes = []
        axes_setup = self.axes_setup
        for ax, new_name in zip(self.axes, axes_names, strict=True):
            if ax.name == new_name:
                new_axes.append(ax)
                continue
            new_ax = Axis(name=new_name, axis_type=ax.axis_type, unit=ax.unit)
            match ax.name:
                case axes_setup.x:
                    axes_setup = axes_setup.model_copy(update={"x": new_name})
                case axes_setup.y:
                    axes_setup = axes_setup.model_copy(update={"y": new_name})
                case axes_setup.z:
                    axes_setup = axes_setup.model_copy(update={"z": new_name})
                case axes_setup.c:
                    axes_setup = axes_setup.model_copy(update={"c": new_name})
                case axes_setup.t:
                    axes_setup = axes_setup.model_copy(update={"t": new_name})
                case _:
                    if ax.name in axes_setup.others:
                        others = axes_setup.others.copy()
                        others.remove(ax.name)
                        others.append(new_name)
                        axes_setup = axes_setup.model_copy(update={"others": others})
            new_axes.append(new_ax)
            # Update the axes setup
        return AxesHandler(
            axes=new_axes,
            axes_setup=axes_setup,
        )

    def get_index(self, name: str) -> int | None:
        """Get the index of the axis by name."""
        return self._index_mapping.get(name, None)

    def has_axis(self, axis_name: str) -> bool:
        """Return whether the axis exists."""
        index = self.get_index(axis_name)
        if index is None:
            return False
        return True

    def get_canonical_name(self, name: str) -> str | None:
        """Get the canonical name of the axis by name."""
        return self._axes_setup.get_canonical_name(name)

    def get_axis(self, name: str) -> Axis | None:
        """Get the axis object by name."""
        index = self.get_index(name)
        if index is None:
            return None
        return self.axes[index]

    def validate_axes_type(self):
        """Validate the axes type.

        If the axes type is not correct, a warning is issued.
        and the axis is implicitly cast to the correct type.
        """
        new_axes = []
        for axes in self.axes:
            for name in self._canonical_order:
                if axes == self.get_axis(name):
                    new_axes.append(axes.canonical_axis_cast(name))
                    break
            else:
                new_axes.append(axes)
        self._axes = new_axes


def _build_axes_list_from_names(
    axes_names: Sequence[str],
    axes_setup: AxesSetup,
    space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
    time_unit: TimeUnits | str | None = DefaultTimeUnit,
) -> list[Axis]:
    """Build a list of Axis objects from a list of axis names."""
    axes = []
    for name in axes_names:
        c_name = axes_setup.get_canonical_name(name)
        match c_name:
            case "t":
                axes.append(Axis(name=name, axis_type=AxisType.time, unit=time_unit))
            case "c":
                axes.append(Axis(name=name, axis_type=AxisType.channel))
            case "z" | "y" | "x":
                axes.append(Axis(name=name, axis_type=AxisType.space, unit=space_unit))
            case _:
                axes.append(Axis(name=name, axis_type=AxisType.space))
    return axes


def build_canonical_axes_handler(
    axes_names: Sequence[str],
    canonical_channel_order: Sequence[str] | None = None,
    space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
    time_unit: TimeUnits | str | None = DefaultTimeUnit,
) -> AxesHandler:
    """Create a new canonical axes mapper.

    Args:
        axes_names (Sequence[str] | int): The axes names on disk.
            - The axes should be in ['t', 'c', 'z', 'y', 'x']
            - The axes should be in strict canonical order.
            - If an integer is provided, the axes are created from the last axis
              to the first
                e.g. 3 -> ["z", "y", "x"]
        canonical_channel_order (Sequence[str], optional): The canonical channel
            order. Defaults to None, which uses the default order.
        space_unit (SpaceUnits, optional): The space units. Defaults to None.
        time_unit (TimeUnits, optional): The time units. Defaults to None.
    """
    if canonical_channel_order is None:
        canonical_channel_order = canonical_axes_order()
    axes_setup = AxesSetup.from_ordered_list(
        axes_names=axes_names, canonical_order=canonical_channel_order
    )
    axes = _build_axes_list_from_names(
        axes_names=axes_names,
        axes_setup=axes_setup,
        space_unit=space_unit,
        time_unit=time_unit,
    )
    return AxesHandler(
        axes=axes,
        axes_setup=axes_setup,
    )


def build_axes_handler(
    axes_names: Sequence[str],
    axes_setup: AxesSetup | None = None,
    space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
    time_unit: TimeUnits | str | None = DefaultTimeUnit,
) -> AxesHandler:
    """Create a new axes mapper.

    Args:
        axes_names (Sequence[str]): The axes names on disk.
        axes_setup (AxesSetup, optional): The axis setup. Defaults to None.
        space_unit (SpaceUnits, optional): The space units. Defaults to None.
        time_unit (TimeUnits, optional): The time units. Defaults to None.
    """
    axes_setup = axes_setup if axes_setup is not None else AxesSetup()
    axes = _build_axes_list_from_names(
        axes_names=axes_names,
        axes_setup=axes_setup,
        space_unit=space_unit,
        time_unit=time_unit,
    )
    return AxesHandler(
        axes=axes,
        axes_setup=axes_setup,
    )
