import collections.abc
import itertools
from abc import ABC, abstractmethod
from collections import defaultdict
from datetime import datetime
from typing import (
    Collection,
    Dict,
    List,
    Mapping,
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
    Type,
    Union,
    cast,
)

import dagster._check as check
from dagster._annotations import PublicAttr, experimental, public
from dagster._core.definitions.multi_dimensional_partitions import (
    MultiPartitionKey,
    MultiPartitionsDefinition,
)
from dagster._core.definitions.partition import (
    PartitionsDefinition,
    PartitionsSubset,
    StaticPartitionsDefinition,
)
from dagster._core.definitions.time_window_partitions import TimeWindowPartitionsDefinition
from dagster._core.instance import DynamicPartitionsStore
from dagster._serdes import whitelist_for_serdes
from dagster._utils.cached_method import cached_method
from dagster._utils.warnings import disable_dagster_warnings


class UpstreamPartitionsResult(NamedTuple):
    """Represents the result of mapping a PartitionsSubset to the corresponding
    partitions in another PartitionsDefinition.

    partitions_subset (PartitionsSubset): The resulting partitions subset that was
        mapped to. Only contains partitions for existent partitions, filtering out nonexistent partitions.
    required_but_nonexistent_partition_keys (Sequence[str]): A list containing invalid partition keys in to_partitions_def
        that partitions in from_partitions_subset were mapped to.
    """

    partitions_subset: PartitionsSubset
    required_but_nonexistent_partition_keys: Sequence[str]


class PartitionMapping(ABC):
    """Defines a correspondence between the partitions in an asset and the partitions in an asset
    that it depends on.

    Overriding PartitionMapping outside of Dagster is not supported. The abstract methods of this
    class may change at any time.
    """

    @public
    @abstractmethod
    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        """Returns the subset of partition keys in the downstream asset that use the data in the given
        partition key subset of the upstream asset.

        Args:
            upstream_partitions_subset (Union[PartitionKeyRange, PartitionsSubset]): The
                subset of partition keys in the upstream asset.
            downstream_partitions_def (PartitionsDefinition): The partitions definition for the
                downstream asset.
        """

    @public
    @abstractmethod
    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        """Returns a UpstreamPartitionsResult object containing the partition keys the downstream
        partitions subset was mapped to in the upstream partitions definition.

        Valid upstream partitions will be included in UpstreamPartitionsResult.partitions_subset.
        Invalid upstream partitions will be included in UpstreamPartitionsResult.required_but_nonexistent_partition_keys.

        For example, if an upstream asset is time-partitioned and starts in June 2023, and the
        downstream asset is time-partitioned and starts in May 2023, this function would return a
        UpstreamPartitionsResult(PartitionsSubset("2023-06-01"), required_but_nonexistent_partition_keys=["2023-05-01"])
        when downstream_partitions_subset contains 2023-05-01 and 2023-06-01.
        """


@whitelist_for_serdes
class IdentityPartitionMapping(PartitionMapping, NamedTuple("_IdentityPartitionMapping", [])):
    """Expects that the upstream and downstream assets are partitioned in the same way, and maps
    partitions in the downstream asset to the same partition in the upstream asset.
    """

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        if downstream_partitions_subset is None:
            check.failed("downstream asset is not partitioned")

        if downstream_partitions_subset.partitions_def == upstream_partitions_def:
            return UpstreamPartitionsResult(downstream_partitions_subset, [])

        upstream_partition_keys = set(
            upstream_partitions_def.get_partition_keys(
                dynamic_partitions_store=dynamic_partitions_store
            )
        )
        downstream_partition_keys = set(downstream_partitions_subset.get_partition_keys())

        return UpstreamPartitionsResult(
            upstream_partitions_def.subset_with_partition_keys(
                list(upstream_partition_keys & downstream_partition_keys)
            ),
            list(downstream_partition_keys - upstream_partition_keys),
        )

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        if upstream_partitions_subset is None:
            check.failed("upstream asset is not partitioned")

        if upstream_partitions_subset.partitions_def == downstream_partitions_def:
            return upstream_partitions_subset

        upstream_partition_keys = set(upstream_partitions_subset.get_partition_keys())
        downstream_partition_keys = set(
            downstream_partitions_def.get_partition_keys(
                dynamic_partitions_store=dynamic_partitions_store
            )
        )

        return downstream_partitions_def.empty_subset().with_partition_keys(
            list(downstream_partition_keys & upstream_partition_keys)
        )


@whitelist_for_serdes
class AllPartitionMapping(PartitionMapping, NamedTuple("_AllPartitionMapping", [])):
    """Maps every partition in the downstream asset to every partition in the upstream asset.

    Commonly used in the case when the downstream asset is not partitioned, in which the entire
    downstream asset depends on all partitions of the usptream asset.
    """

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        upstream_subset = upstream_partitions_def.subset_with_all_partitions(
            current_time=current_time, dynamic_partitions_store=dynamic_partitions_store
        )
        return UpstreamPartitionsResult(upstream_subset, [])

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        raise NotImplementedError()


@whitelist_for_serdes
class LastPartitionMapping(PartitionMapping, NamedTuple("_LastPartitionMapping", [])):
    """Maps all dependencies to the last partition in the upstream asset.

    Commonly used in the case when the downstream asset is not partitioned, in which the entire
    downstream asset depends on the last partition of the upstream asset.
    """

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        last = upstream_partitions_def.get_last_partition_key(
            current_time=None, dynamic_partitions_store=dynamic_partitions_store
        )

        upstream_subset = upstream_partitions_def.empty_subset()
        if last is not None:
            upstream_subset = upstream_subset.with_partition_keys([last])

        return UpstreamPartitionsResult(upstream_subset, [])

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        raise NotImplementedError()


@whitelist_for_serdes
class SpecificPartitionsPartitionMapping(
    PartitionMapping,
    NamedTuple(
        "_SpecificPartitionsPartitionMapping", [("partition_keys", PublicAttr[Sequence[str]])]
    ),
):
    """Maps to a specific subset of partitions in the upstream asset.

    Example:
        .. code-block:: python

            from dagster import SpecificPartitionsPartitionMapping, StaticPartitionsDefinition, asset

            @asset(partitions_def=StaticPartitionsDefinition(["a", "b", "c"]))
            def upstream():
                ...

            @asset(
                ins={
                    "upstream": AssetIn(partition_mapping=SpecificPartitionsPartitionMapping(["a"]))
                }
            )
            def a_downstream(upstream):
                ...
    """

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        return UpstreamPartitionsResult(
            upstream_partitions_def.subset_with_partition_keys(self.partition_keys), []
        )

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        # if any of the partition keys in this partition mapping are contained within the upstream
        # partitions subset, then all partitions of the downstream asset are dependencies
        if any(key in upstream_partitions_subset for key in self.partition_keys):
            return downstream_partitions_def.subset_with_all_partitions(
                dynamic_partitions_store=dynamic_partitions_store
            )
        return downstream_partitions_def.empty_subset()


class DimensionDependency(NamedTuple):
    partition_mapping: PartitionMapping
    upstream_dimension_name: Optional[str] = None
    downstream_dimension_name: Optional[str] = None


class BaseMultiPartitionMapping(ABC):
    @abstractmethod
    def get_dimension_dependencies(
        self,
        upstream_partitions_def: PartitionsDefinition,
        downstream_partitions_def: PartitionsDefinition,
    ) -> Sequence[DimensionDependency]: ...

    def get_partitions_def(
        self, partitions_def: PartitionsDefinition, dimension_name: Optional[str]
    ) -> PartitionsDefinition:
        if isinstance(partitions_def, MultiPartitionsDefinition):
            if not isinstance(dimension_name, str):
                check.failed("Expected dimension_name to be a string")
            return partitions_def.get_partitions_def_for_dimension(dimension_name)
        return partitions_def

    def _get_dependency_partitions_subset(
        self,
        a_partitions_def: PartitionsDefinition,
        a_partitions_subset: PartitionsSubset,
        b_partitions_def: PartitionsDefinition,
        a_upstream_of_b: bool,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
        current_time: Optional[datetime] = None,
    ) -> Union[UpstreamPartitionsResult, PartitionsSubset]:
        """Given two partitions definitions a_partitions_def and b_partitions_def that have a dependency
        relationship (a_upstream_of_b is True if a_partitions_def is upstream of b_partitions_def),
        and a_partition_keys, a list of partition keys in a_partitions_def, returns a list of
        partition keys in the partitions definition b_partitions_def that are
        dependencies of the partition keys in a_partition_keys.
        """
        a_partition_keys_by_dimension = defaultdict(set)
        if isinstance(a_partitions_def, MultiPartitionsDefinition):
            for partition_key in a_partitions_subset.get_partition_keys():
                for dimension_name, key in cast(
                    MultiPartitionKey, partition_key
                ).keys_by_dimension.items():
                    a_partition_keys_by_dimension[dimension_name].add(key)
        else:
            for partition_key in a_partitions_subset.get_partition_keys():
                a_partition_keys_by_dimension[None].add(partition_key)

        # Maps the dimension name and key of a partition in a_partitions_def to the list of
        # partition keys in b_partitions_def that are dependencies of that partition
        dep_b_keys_by_a_dim_and_key: Dict[Optional[str], Dict[Optional[str], List[str]]] = (
            defaultdict(lambda: defaultdict(list))
        )
        required_but_nonexistent_upstream_partitions = set()

        b_dimension_partitions_def_by_name: Dict[Optional[str], PartitionsDefinition] = (
            {
                dimension.name: dimension.partitions_def
                for dimension in b_partitions_def.partitions_defs
            }
            if isinstance(b_partitions_def, MultiPartitionsDefinition)
            else {None: b_partitions_def}
        )

        if a_upstream_of_b:
            # a_partitions_def is upstream of b_partitions_def, so we need to map the
            # dimension names of a_partitions_def to the corresponding dependent dimensions of
            # b_partitions_def
            a_dim_to_dependency_b_dim = {
                dimension_mapping.upstream_dimension_name: (
                    dimension_mapping.downstream_dimension_name,
                    dimension_mapping.partition_mapping,
                )
                for dimension_mapping in self.get_dimension_dependencies(
                    a_partitions_def, b_partitions_def
                )
            }

            for a_dim_name, keys in a_partition_keys_by_dimension.items():
                if a_dim_name in a_dim_to_dependency_b_dim:
                    (
                        b_dim_name,
                        dimension_mapping,
                    ) = a_dim_to_dependency_b_dim[a_dim_name]
                    a_dimension_partitions_def = self.get_partitions_def(
                        a_partitions_def, a_dim_name
                    )
                    b_dimension_partitions_def = self.get_partitions_def(
                        b_partitions_def, b_dim_name
                    )
                    for key in keys:
                        # if downstream dimension mapping exists, for a given key, get the list of
                        # downstream partition keys that are dependencies of that key
                        dep_b_keys_by_a_dim_and_key[a_dim_name][key] = list(
                            dimension_mapping.get_downstream_partitions_for_partitions(
                                a_dimension_partitions_def.empty_subset().with_partition_keys(
                                    [key]
                                ),
                                b_dimension_partitions_def,
                                current_time=current_time,
                                dynamic_partitions_store=dynamic_partitions_store,
                            ).get_partition_keys()
                        )

        else:
            # a_partitions_def is downstream of b_partitions_def, so we need to map the
            # dimension names of a_partitions_def to the corresponding dependency dimensions of
            # b_partitions_def
            a_dim_to_dependency_b_dim = {
                dimension_mapping.downstream_dimension_name: (
                    dimension_mapping.upstream_dimension_name,
                    dimension_mapping.partition_mapping,
                )
                for dimension_mapping in self.get_dimension_dependencies(
                    b_partitions_def, a_partitions_def
                )
            }

            for a_dim_name, keys in a_partition_keys_by_dimension.items():
                if a_dim_name in a_dim_to_dependency_b_dim:
                    (
                        b_dim_name,
                        partition_mapping,
                    ) = a_dim_to_dependency_b_dim[a_dim_name]
                    a_dimension_partitions_def = self.get_partitions_def(
                        a_partitions_def, a_dim_name
                    )
                    b_dimension_partitions_def = self.get_partitions_def(
                        b_partitions_def, b_dim_name
                    )
                    for key in keys:
                        mapped_partitions_result = (
                            partition_mapping.get_upstream_mapped_partitions_result_for_partitions(
                                a_dimension_partitions_def.empty_subset().with_partition_keys(
                                    [key]
                                ),
                                b_dimension_partitions_def,
                                current_time=current_time,
                                dynamic_partitions_store=dynamic_partitions_store,
                            )
                        )
                        dep_b_keys_by_a_dim_and_key[a_dim_name][key] = list(
                            mapped_partitions_result.partitions_subset.get_partition_keys()
                        )
                        required_but_nonexistent_upstream_partitions.update(
                            set(mapped_partitions_result.required_but_nonexistent_partition_keys)
                        )

        b_partition_keys = set()

        mapped_a_dim_names = a_dim_to_dependency_b_dim.keys()
        mapped_b_dim_names = [mapping[0] for mapping in a_dim_to_dependency_b_dim.values()]
        unmapped_b_dim_names = list(
            set(b_dimension_partitions_def_by_name.keys()) - set(mapped_b_dim_names)
        )

        for key in a_partitions_subset.get_partition_keys():
            for b_key_values in itertools.product(
                *(
                    [
                        dep_b_keys_by_a_dim_and_key[dim_name][
                            (
                                cast(MultiPartitionKey, key).keys_by_dimension[dim_name]
                                if dim_name
                                else key
                            )
                        ]
                        for dim_name in mapped_a_dim_names
                    ]
                ),
                *[
                    b_dimension_partitions_def_by_name[dim_name].get_partition_keys()
                    for dim_name in unmapped_b_dim_names
                ],
            ):
                b_partition_keys.add(
                    MultiPartitionKey(
                        {
                            cast(str, (mapped_b_dim_names + unmapped_b_dim_names)[i]): key
                            for i, key in enumerate(b_key_values)
                        }
                    )
                    if len(b_key_values) > 1
                    else b_key_values[0]
                )

        mapped_subset = b_partitions_def.empty_subset().with_partition_keys(b_partition_keys)
        if a_upstream_of_b:
            return mapped_subset
        else:
            return UpstreamPartitionsResult(
                mapped_subset,
                required_but_nonexistent_partition_keys=list(
                    required_but_nonexistent_upstream_partitions
                ),
            )

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        if downstream_partitions_subset is None:
            check.failed("downstream asset is not partitioned")

        result = self._get_dependency_partitions_subset(
            cast(MultiPartitionsDefinition, downstream_partitions_subset.partitions_def),
            downstream_partitions_subset,
            cast(MultiPartitionsDefinition, upstream_partitions_def),
            a_upstream_of_b=False,
            dynamic_partitions_store=dynamic_partitions_store,
            current_time=current_time,
        )

        if not isinstance(result, UpstreamPartitionsResult):
            check.failed("Expected UpstreamPartitionsResult")

        return result

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        if upstream_partitions_subset is None:
            check.failed("upstream asset is not partitioned")

        result = self._get_dependency_partitions_subset(
            cast(MultiPartitionsDefinition, upstream_partitions_subset.partitions_def),
            upstream_partitions_subset,
            cast(MultiPartitionsDefinition, downstream_partitions_def),
            a_upstream_of_b=True,
            dynamic_partitions_store=dynamic_partitions_store,
        )

        if isinstance(result, UpstreamPartitionsResult):
            check.failed("Expected PartitionsSubset")

        return result


@experimental
@whitelist_for_serdes
class MultiToSingleDimensionPartitionMapping(
    BaseMultiPartitionMapping,
    PartitionMapping,
    NamedTuple(
        "_MultiToSingleDimensionPartitionMapping", [("partition_dimension_name", Optional[str])]
    ),
):
    """Defines a correspondence between an single-dimensional partitions definition
    and a MultiPartitionsDefinition. The single-dimensional partitions definition must be
    a dimension of the MultiPartitionsDefinition.

    This class handles the case where the upstream asset is multipartitioned and the
    downstream asset is single dimensional, and vice versa.

    For a partition key X, this partition mapping assumes that any multi-partition key with
    X in the selected dimension is a dependency.

    Args:
        partition_dimension_name (Optional[str]): The name of the partition dimension in the
            MultiPartitionsDefinition that matches the single-dimension partitions definition.
    """

    def __new__(cls, partition_dimension_name: Optional[str] = None):
        return super(MultiToSingleDimensionPartitionMapping, cls).__new__(
            cls,
            partition_dimension_name=check.opt_str_param(
                partition_dimension_name, "partition_dimension_name"
            ),
        )

    def get_dimension_dependencies(
        self,
        upstream_partitions_def: PartitionsDefinition,
        downstream_partitions_def: PartitionsDefinition,
    ) -> Sequence[DimensionDependency]:
        infer_mapping_result = _get_infer_single_to_multi_dimension_deps_result(
            upstream_partitions_def, downstream_partitions_def
        )

        if not infer_mapping_result.can_infer:
            check.invariant(isinstance(infer_mapping_result.inference_failure_reason, str))
            check.failed(cast(str, infer_mapping_result.inference_failure_reason))

        return [cast(DimensionDependency, infer_mapping_result.dimension_dependency)]


@whitelist_for_serdes
class DimensionPartitionMapping(
    NamedTuple(
        "_DimensionPartitionMapping",
        [
            ("dimension_name", str),
            ("partition_mapping", PartitionMapping),
        ],
    )
):
    """A helper class for MultiPartitionMapping that defines a partition mapping used to calculate
    the dependent partition keys in the selected downstream MultiPartitions definition dimension.

    Args:
        dimension_name (str): The name of the dimension in the downstream MultiPartitionsDefinition.
        partition_mapping (PartitionMapping): The partition mapping object used to calculate
            the downstream dimension partitions from the upstream dimension partitions and vice versa.
    """

    def __new__(
        cls,
        dimension_name: str,
        partition_mapping: PartitionMapping,
    ):
        return super(DimensionPartitionMapping, cls).__new__(
            cls,
            dimension_name=check.str_param(dimension_name, "dimension_name"),
            partition_mapping=check.inst_param(
                partition_mapping, "partition_mapping", PartitionMapping
            ),
        )


@experimental
@whitelist_for_serdes
class MultiPartitionMapping(
    BaseMultiPartitionMapping,
    PartitionMapping,
    NamedTuple(
        "_MultiPartitionMapping",
        [("downstream_mappings_by_upstream_dimension", Mapping[str, DimensionPartitionMapping])],
    ),
):
    """Defines a correspondence between two MultiPartitionsDefinitions.

    Accepts a mapping of upstream dimension name to downstream DimensionPartitionMapping, representing
    the explicit correspondence between the upstream and downstream MultiPartitions dimensions
    and the partition mapping used to calculate the downstream partitions.

    Examples:
        .. code-block:: python

            weekly_abc = MultiPartitionsDefinition(
                {
                    "abc": StaticPartitionsDefinition(["a", "b", "c"]),
                    "weekly": WeeklyPartitionsDefinition("2023-01-01"),
                }
            )
            daily_123 = MultiPartitionsDefinition(
                {
                    "123": StaticPartitionsDefinition(["1", "2", "3"]),
                    "daily": DailyPartitionsDefinition("2023-01-01"),
                }
            )

            MultiPartitionsMapping(
                {
                    "abc": DimensionPartitionMapping(
                        dimension_name="123",
                        partition_mapping=StaticPartitionMapping({"a": "1", "b": "2", "c": "3"}),
                    ),
                    "weekly": DimensionPartitionMapping(
                        dimension_name="daily",
                        partition_mapping=TimeWindowPartitionMapping(),
                    )
                }
            )

    For upstream or downstream dimensions not explicitly defined in the mapping, Dagster will
    assume an `AllPartitionsMapping`, meaning that all upstream partitions in those dimensions
    will be mapped to all downstream partitions in those dimensions.

    Examples:
        .. code-block:: python

            weekly_abc = MultiPartitionsDefinition(
                {
                    "abc": StaticPartitionsDefinition(["a", "b", "c"]),
                    "daily": DailyPartitionsDefinition("2023-01-01"),
                }
            )
            daily_123 = MultiPartitionsDefinition(
                {
                    "123": StaticPartitionsDefinition(["1", "2", "3"]),
                    "daily": DailyPartitionsDefinition("2023-01-01"),
                }
            )

            MultiPartitionsMapping(
                {
                    "daily": DimensionPartitionMapping(
                        dimension_name="daily",
                        partition_mapping=IdentityPartitionMapping(),
                    )
                }
            )

            # Will map `daily_123` partition key {"123": "1", "daily": "2023-01-01"} to the upstream:
            # {"abc": "a", "daily": "2023-01-01"}
            # {"abc": "b", "daily": "2023-01-01"}
            # {"abc": "c", "daily": "2023-01-01"}

    Args:
        downstream_mappings_by_upstream_dimension (Mapping[str, DimensionPartitionMapping]): A
            mapping that defines an explicit correspondence between one dimension of the upstream
            MultiPartitionsDefinition and one dimension of the downstream MultiPartitionsDefinition.
            Maps a string representing upstream dimension name to downstream DimensionPartitionMapping,
            containing the downstream dimension name and partition mapping.
    """

    def __new__(
        cls, downstream_mappings_by_upstream_dimension: Mapping[str, DimensionPartitionMapping]
    ):
        return super(MultiPartitionMapping, cls).__new__(
            cls,
            downstream_mappings_by_upstream_dimension=check.mapping_param(
                downstream_mappings_by_upstream_dimension,
                "downstream_mappings_by_upstream_dimension",
                key_type=str,
                value_type=DimensionPartitionMapping,
            ),
        )

    def get_dimension_dependencies(
        self,
        upstream_partitions_def: PartitionsDefinition,
        downstream_partitions_def: PartitionsDefinition,
    ) -> Sequence[DimensionDependency]:
        self._check_all_dimensions_accounted_for(
            upstream_partitions_def,
            downstream_partitions_def,
        )

        return [
            DimensionDependency(
                mapping.partition_mapping,
                upstream_dimension_name=upstream_dimension,
                downstream_dimension_name=mapping.dimension_name,
            )
            for upstream_dimension, mapping in self.downstream_mappings_by_upstream_dimension.items()
        ]

    def _check_all_dimensions_accounted_for(
        self,
        upstream_partitions_def: PartitionsDefinition,
        downstream_partitions_def: PartitionsDefinition,
    ) -> None:
        if any(
            not isinstance(partitions_def, MultiPartitionsDefinition)
            for partitions_def in (upstream_partitions_def, downstream_partitions_def)
        ):
            check.failed(
                "Both partitions defs provided to a MultiPartitionMapping must be multi-partitioned"
            )

        upstream_dimension_names = {
            dim.name
            for dim in cast(MultiPartitionsDefinition, upstream_partitions_def).partitions_defs
        }
        dimension_names = {
            dim.name
            for dim in cast(MultiPartitionsDefinition, downstream_partitions_def).partitions_defs
        }

        for (
            upstream_dimension_name,
            dimension_mapping,
        ) in self.downstream_mappings_by_upstream_dimension.items():
            if upstream_dimension_name not in upstream_dimension_names:
                check.failed(
                    "Dimension mapping has an upstream dimension name that is not in the upstream "
                    "partitions def"
                )
            if dimension_mapping.dimension_name not in dimension_names:
                check.failed(
                    "Dimension mapping has a downstream dimension name that is not in the"
                    " downstream partitions def"
                )

            upstream_dimension_names.remove(upstream_dimension_name)
            dimension_names.remove(dimension_mapping.dimension_name)


@whitelist_for_serdes
class StaticPartitionMapping(
    PartitionMapping,
    NamedTuple(
        "_StaticPartitionMapping",
        [
            (
                "downstream_partition_keys_by_upstream_partition_key",
                PublicAttr[Mapping[str, Union[str, Collection[str]]]],
            )
        ],
    ),
):
    """Define an explicit correspondence between two StaticPartitionsDefinitions.

    Args:
        downstream_partition_keys_by_upstream_partition_key (Dict[str, str | Collection[str]]):
          The single or multi-valued correspondence from upstream keys to downstream keys.
    """

    def __init__(
        self,
        downstream_partition_keys_by_upstream_partition_key: Mapping[
            str, Union[str, Collection[str]]
        ],
    ):
        check.mapping_param(
            downstream_partition_keys_by_upstream_partition_key,
            "downstream_partition_keys_by_upstream_partition_key",
            key_type=str,
            value_type=(str, collections.abc.Collection),
        )

        # cache forward and reverse mappings
        self._mapping = defaultdict(set)
        for (
            upstream_key,
            downstream_keys,
        ) in downstream_partition_keys_by_upstream_partition_key.items():
            self._mapping[upstream_key] = (
                {downstream_keys} if isinstance(downstream_keys, str) else set(downstream_keys)
            )

        self._inverse_mapping = defaultdict(set)
        for upstream_key, downstream_keys in self._mapping.items():
            for downstream_key in downstream_keys:
                self._inverse_mapping[downstream_key].add(upstream_key)

    @cached_method
    def _check_upstream(self, *, upstream_partitions_def: PartitionsDefinition):
        """Validate that the mapping from upstream to downstream is only defined on upstream keys."""
        check.inst(
            upstream_partitions_def,
            StaticPartitionsDefinition,
            "StaticPartitionMapping can only be defined between two StaticPartitionsDefinitions",
        )
        upstream_keys = upstream_partitions_def.get_partition_keys()
        extra_keys = set(self._mapping.keys()).difference(upstream_keys)
        if extra_keys:
            raise ValueError(
                f"mapping source partitions not in the upstream partitions definition: {extra_keys}"
            )

    @cached_method
    def _check_downstream(self, *, downstream_partitions_def: PartitionsDefinition):
        """Validate that the mapping from upstream to downstream only maps to downstream keys."""
        check.inst(
            downstream_partitions_def,
            StaticPartitionsDefinition,
            "StaticPartitionMapping can only be defined between two StaticPartitionsDefinitions",
        )
        downstream_keys = downstream_partitions_def.get_partition_keys()
        extra_keys = set(self._inverse_mapping.keys()).difference(downstream_keys)
        if extra_keys:
            raise ValueError(
                "mapping target partitions not in the downstream partitions definition:"
                f" {extra_keys}"
            )

    def get_downstream_partitions_for_partitions(
        self,
        upstream_partitions_subset: PartitionsSubset,
        downstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> PartitionsSubset:
        self._check_downstream(downstream_partitions_def=downstream_partitions_def)

        downstream_subset = downstream_partitions_def.empty_subset()
        downstream_keys = set()
        for key in upstream_partitions_subset.get_partition_keys():
            downstream_keys.update(self._mapping[key])
        return downstream_subset.with_partition_keys(downstream_keys)

    def get_upstream_mapped_partitions_result_for_partitions(
        self,
        downstream_partitions_subset: Optional[PartitionsSubset],
        upstream_partitions_def: PartitionsDefinition,
        current_time: Optional[datetime] = None,
        dynamic_partitions_store: Optional[DynamicPartitionsStore] = None,
    ) -> UpstreamPartitionsResult:
        self._check_upstream(upstream_partitions_def=upstream_partitions_def)

        upstream_subset = upstream_partitions_def.empty_subset()
        if downstream_partitions_subset is None:
            return UpstreamPartitionsResult(upstream_subset, [])

        upstream_keys = set()
        for key in downstream_partitions_subset.get_partition_keys():
            upstream_keys.update(self._inverse_mapping[key])

        return UpstreamPartitionsResult(upstream_subset.with_partition_keys(upstream_keys), [])


class InferSingleToMultiDimensionDepsResult(
    NamedTuple(
        "_InferSingleToMultiDimensionDepsResult",
        [
            ("can_infer", bool),
            ("inference_failure_reason", Optional[str]),
            ("dimension_dependency", Optional[DimensionDependency]),
        ],
    )
):
    def __new__(
        cls,
        can_infer: bool,
        inference_failure_reason: Optional[str] = None,
        dimension_dependency: Optional[DimensionDependency] = None,
    ):
        if can_infer and dimension_dependency is None:
            check.failed("dimension_dependency must be provided if can_infer is True")
        if not can_infer and inference_failure_reason is None:
            check.failed("inference_failure_reason must be provided if can_infer is False")

        return super(InferSingleToMultiDimensionDepsResult, cls).__new__(
            cls,
            can_infer,
            inference_failure_reason,
            dimension_dependency,
        )


def _get_infer_single_to_multi_dimension_deps_result(
    upstream_partitions_def: PartitionsDefinition,
    downstream_partitions_def: PartitionsDefinition,
    partition_dimension_name: Optional[str] = None,
) -> InferSingleToMultiDimensionDepsResult:
    from dagster._core.definitions.time_window_partition_mapping import TimeWindowPartitionMapping

    upstream_is_multipartitioned = isinstance(upstream_partitions_def, MultiPartitionsDefinition)

    multipartitions_defs = [
        partitions_def
        for partitions_def in [upstream_partitions_def, downstream_partitions_def]
        if isinstance(partitions_def, MultiPartitionsDefinition)
    ]
    if len(multipartitions_defs) != 1:
        return InferSingleToMultiDimensionDepsResult(
            False,
            "Can only use MultiToSingleDimensionPartitionMapping when upstream asset is"
            " multipartitioned and the downstream asset is single dimensional, or vice versa."
            f" Instead received {len(multipartitions_defs)} multi-partitioned assets.",
        )

    multipartitions_def = cast(MultiPartitionsDefinition, next(iter(multipartitions_defs)))

    single_dimension_partitions_def = next(
        iter(
            {
                upstream_partitions_def,
                downstream_partitions_def,
            }
            - set(multipartitions_defs)
        )
    )

    filtered_multipartition_dims = (
        multipartitions_def.partitions_defs
        if partition_dimension_name is None
        else [
            dim
            for dim in multipartitions_def.partitions_defs
            if dim.name == partition_dimension_name
        ]
    )

    if partition_dimension_name:
        if len(filtered_multipartition_dims) != 1:
            return InferSingleToMultiDimensionDepsResult(
                False,
                f"Provided partition dimension name {partition_dimension_name} not found in"
                f" multipartitions definition {multipartitions_def}.",
            )

    matching_dimension_defs = [
        dimension_def
        for dimension_def in filtered_multipartition_dims
        if dimension_def.partitions_def == single_dimension_partitions_def
    ]

    if len(matching_dimension_defs) == 1:
        return InferSingleToMultiDimensionDepsResult(
            True,
            dimension_dependency=DimensionDependency(
                IdentityPartitionMapping(),
                upstream_dimension_name=(
                    matching_dimension_defs[0].name if upstream_is_multipartitioned else None
                ),
                downstream_dimension_name=(
                    matching_dimension_defs[0].name if not upstream_is_multipartitioned else None
                ),
            ),
        )
    elif len(matching_dimension_defs) > 1:
        return InferSingleToMultiDimensionDepsResult(
            False,
            "partition dimension name must be specified when multiple dimensions of the"
            " MultiPartitionsDefinition match the single dimension partitions def",
        )

    time_dimensions = [
        dimension_def
        for dimension_def in filtered_multipartition_dims
        if isinstance(dimension_def.partitions_def, TimeWindowPartitionsDefinition)
    ]

    if len(time_dimensions) == 1 and isinstance(
        single_dimension_partitions_def, TimeWindowPartitionsDefinition
    ):
        return InferSingleToMultiDimensionDepsResult(
            True,
            dimension_dependency=DimensionDependency(
                TimeWindowPartitionMapping(),
                upstream_dimension_name=(
                    time_dimensions[0].name if upstream_is_multipartitioned else None
                ),
                downstream_dimension_name=(
                    time_dimensions[0].name if not upstream_is_multipartitioned else None
                ),
            ),
        )

    return InferSingleToMultiDimensionDepsResult(
        False,
        "MultiToSingleDimensionPartitionMapping can only be used when: \n(a) The single dimensional"
        " partitions definition is a dimension of the MultiPartitionsDefinition.\n(b) The single"
        " dimensional partitions definition is a TimeWindowPartitionsDefinition and the"
        " MultiPartitionsDefinition has a single time dimension.",
    )


def infer_partition_mapping(
    partition_mapping: Optional[PartitionMapping],
    downstream_partitions_def: Optional[PartitionsDefinition],
    upstream_partitions_def: Optional[PartitionsDefinition],
) -> PartitionMapping:
    from .time_window_partition_mapping import TimeWindowPartitionMapping

    if partition_mapping is not None:
        return partition_mapping
    elif upstream_partitions_def and downstream_partitions_def:
        if _get_infer_single_to_multi_dimension_deps_result(
            upstream_partitions_def, downstream_partitions_def
        ).can_infer:
            with disable_dagster_warnings():
                return MultiToSingleDimensionPartitionMapping()
        elif isinstance(upstream_partitions_def, TimeWindowPartitionsDefinition) and isinstance(
            downstream_partitions_def, TimeWindowPartitionsDefinition
        ):
            return TimeWindowPartitionMapping()
        else:
            return IdentityPartitionMapping()
    else:
        return AllPartitionMapping()


def get_builtin_partition_mapping_types() -> Tuple[Type[PartitionMapping], ...]:
    from dagster._core.definitions.time_window_partition_mapping import TimeWindowPartitionMapping

    return (
        AllPartitionMapping,
        IdentityPartitionMapping,
        LastPartitionMapping,
        SpecificPartitionsPartitionMapping,
        StaticPartitionMapping,
        TimeWindowPartitionMapping,
        MultiToSingleDimensionPartitionMapping,
        MultiPartitionMapping,
    )
