class ModelStateMapper(abc.ABC):
    """
   The abstract base class for all model state transformation operations.

   This class serves as the interface between the definition of a transformation
   topology and the actual execution of tensor operations.

   It enforces a Declarative vs. Imperative separation of concerns:

   1.  Declarative (Topology): Through `state_dependency_groups()`, the mapper
       announces *what* it intends to do without handling any data. This allows the system to build execution graphs,
       validate chains, detect collisions, and shard tasks *before* allocating memory.
   2.  Imperative (Execution): Through `apply()`, the mapper performs the
       actual logic (PyTorch operations) on model states.
   """

    @abc.abstractmethod
    def state_dependency_groups(self) -> frozenset[StateGroup]:
        """
        Calculates and returns the set of independent dependency groups this mapper handles.

        Returns:
            A frozenset of `StateGroup` objects. Each group
            represents a disjoint operation. For example, a mapper that renames ten
            independent tensors would return ten distinct `StateGroup` objects,
            allowing them to be sharded or processed individually.
        """
        ...

    @abc.abstractmethod
    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """
        Executes the transformation logic on a specific dictionary of tensors.

        The orchestration system guarantees that the `group` dictionary passed here contains
        all keys listed in the `inputs` of the active `StateGroup`.

        Implementation of this method should guarantee that the result will contain all keys listed in the `outputs`.

        Args:
           group: A dictionary containing the source data.
               Keys match `StateGroup.inputs`.

        Returns:
           A dictionary containing the transformed data. Keys must strictly match `StateGroup.outputs`.
        """
        ...

@dataclasses.dataclass(frozen=True)
class StateGroup:
    """
    Represents an atomic unit of dependency in the model state transformation graph.

    A `StateGroup` defines a strict contract between a set of input keys (source)
    and a set of output keys (destination).

    Attributes:
        inputs: The complete set of keys required from the source state dictionary to satisfy this dependency.
        outputs: The complete set of keys that will be produced as a result of this transformation.
    """

    inputs: frozenset[str]
    outputs: frozenset[str]

class ModelStateMapperParallel(ModelStateMapper):
    """
    Executes a list of states mappers independently alongside each other.

    This class aggregates multiple mappers into a single logical unit.
    It enforces strict isolation between the mappers: no two mappers can
    consume the same input key (input collision) or produce the same output
    key (output collision).

    During execution (`apply`), it routes the specific subset of the input dictionary
    to the sub-mapper responsible for those keys.
    """

    def __init__(self, mappers: Sequence[ModelStateMapper]):
        mappers_lst = filter_empty_mappers(mappers)

        all_groups = set()
        inputs_to_mapper = {}

        seen_inputs: set[str] = set()
        seen_outputs: set[str] = set()
        for mapper in mappers_lst:
            sub_groups = mapper.state_dependency_groups()

            for sub_group in sub_groups:
                if not seen_inputs.isdisjoint(sub_group.inputs):
                    raise ValueError(f"Found a colliding input group: {sub_group.inputs}")
                seen_inputs.update(sub_group.inputs)

                if not seen_outputs.isdisjoint(sub_group.outputs):
                    raise ValueError(f"Found colliding output keys: {sub_group.outputs}")
                seen_outputs.update(sub_group.outputs)

                all_groups.add(sub_group)
                inputs_to_mapper[sub_group.inputs] = mapper

        self._all_groups = frozenset(all_groups)
        self._inputs_to_mapper = inputs_to_mapper

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return self._all_groups

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        group_keys = frozenset(group.keys())

        if group_keys not in self._inputs_to_mapper:
            raise ValueError("Tried to run a parallel mapper with undefined group. Perhaps you sent groups that are "
                             "not isolated?")

        return self._inputs_to_mapper[group_keys].apply(group)

class ModelStateMapperSequential(ModelStateMapper):
    """
    Executes a list of mappers in a specific sequence (pipeline).

    This class manages the data flow from one mapper to the next. It abstracts
    away intermediate states, exposing only the inputs required by the first
    relevant stage and the outputs produced by the final relevant stage.

    Key Features:

    1. **Gap Filling**: Automatically injects `Identity` mappers if a tensor needs
       to pass through a stage without modification to reach a later stage or
       the final output.

    2. **Group Merging**: Computes the net dependency graph. If Stage A requires 'x'
       and produces 'y', and Stage B requires 'y' and produces 'z', the
       Sequential mapper reports a single group `{x} -> {z}`.
    """

    def __init__(self, mappers: list[ModelStateMapper]):
        mappers = filter_empty_mappers(mappers)
        if not mappers:
            raise ValueError("Mappers list cannot be empty.")

        mappers = self._fill_gaps(mappers)

        self._groups = self._compute_pipeline_groups(mappers)
        self._mappers = mappers

    @staticmethod
    def _fill_gaps(mappers: list[ModelStateMapper]) -> list[ModelStateMapper]:
        mappers = mappers.copy()

        # propagate inputs from bottom to top
        for stage_i in range(1, len(mappers))[::-1]:
            groups_current = mappers[stage_i].state_dependency_groups()
            groups_prev = mappers[stage_i - 1].state_dependency_groups()
            current_stage_requires = frozenset.union(*(x.inputs for x in groups_current))
            prev_stage_produces = frozenset.union(*(x.outputs for x in groups_prev))

            needs_to_pass_through = current_stage_requires - prev_stage_produces

            mappers[stage_i - 1] = ModelStateMapperParallel(
                [mappers[stage_i - 1]] + [ModelStateMapperIdentity(x) for x in needs_to_pass_through]
            )

        # propagate outputs from top to bottom
        for stage_i in range(0, len(mappers) - 1):
            groups_current = mappers[stage_i].state_dependency_groups()
            groups_next = mappers[stage_i + 1].state_dependency_groups()
            current_stage_produces = frozenset.union(*(x.outputs for x in groups_current))
            next_stage_requires = frozenset.union(*(x.inputs for x in groups_next))

            needs_to_pass_through = current_stage_produces - next_stage_requires

            mappers[stage_i + 1] = ModelStateMapperParallel(
                [mappers[stage_i + 1]] + [ModelStateMapperIdentity(x) for x in needs_to_pass_through]
            )

        return mappers

    @staticmethod
    def _compute_pipeline_groups(mappers: list[ModelStateMapper]) -> frozenset[StateGroup]:
        outputs_depend_on_inputs = {}

        # given a fully connected graph, we can just go upwards
        for last_group_traced in mappers[-1].state_dependency_groups():
            required_inputs = last_group_traced.inputs

            for mapper_i in range(0, len(mappers) - 1)[::-1]:
                next_visit_groups = [x for x in mappers[mapper_i].state_dependency_groups()
                                     if not x.outputs.isdisjoint(required_inputs)]

                required_inputs = frozenset.union(*(x.inputs for x in next_visit_groups))

            outputs_depend_on_inputs[last_group_traced.outputs] = required_inputs

        return ModelStateMapperSequential._merge_groups(list(outputs_depend_on_inputs.items()))

    @staticmethod
    def _merge_groups(groups: Sequence[tuple[AbstractSet[str], AbstractSet[str]]]) -> frozenset[StateGroup]:
        saved_groups: list[tuple[set[str], set[str]]] = []

        saved_groups_modified = True
        while saved_groups_modified:
            saved_groups_modified = False
            for output_names, input_names in groups:
                was_new_group_created = False
                for group in saved_groups:
                    if group[0].intersection(input_names) or group[1].intersection(output_names):
                        group[0].update(input_names)
                        group[1].update(output_names)
                        was_new_group_created = True
                        saved_groups_modified = True

                if not was_new_group_created:
                    saved_groups.append((set(input_names), set(output_names)))

            groups = saved_groups
            saved_groups = []

        return frozenset(StateGroup(inputs=frozenset(x[0]), outputs=frozenset(x[1])) for x in groups)

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return self._groups

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        current_state = group
        next_state = {}
        for mapper in self._mappers:
            for deps in mapper.state_dependency_groups():
                if not deps.inputs <= current_state.keys():
                    continue

                next_state.update(mapper.apply({k: v for k, v in current_state.items() if k in deps.inputs}))

            current_state = next_state
            next_state = {}

        return current_state

class ModelStateMapperShard(ModelStateMapper):
    """
    Wraps another state mapper and restricts its execution to a specific subset (shard)
    of dependency groups.

    This is primarily used for parallelizing model loading across multiple processes
    or nodes. By assigning a different `current_shard` index to each process,
    the total set of tensors required by the `sub_mapper` is split evenly,
    preventing every process from loading the entire checkpoint.
    """

    def __init__(self, sub_mapper: ModelStateMapper, total_shards: int, current_shard: int):
        self._groups = self._shard_groups(
            sub_mapper.state_dependency_groups(),
            n_shards=total_shards, shard=current_shard
        )
        self._sub_mapper = sub_mapper
        self._total_shards = total_shards
        self._current_shard = current_shard

    @staticmethod
    def _shard_groups(groups: frozenset[StateGroup], n_shards: int, shard: int) -> frozenset[StateGroup]:
        groups_sorted = sorted(groups, key=lambda x: sorted(x.inputs))
        groups_shard = [x for i, x in enumerate(groups_sorted) if i % n_shards == shard]
        return frozenset(groups_shard)

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return self._groups

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return self._sub_mapper.apply(group)

class ModelStateMapperDistribute(ModelStateMapper):
    """
    Converts a single local Tensor object into a DTensor object with specified
    `device_mesh` and `placements`.
    """

    def __init__(self, name: str, device_mesh: DeviceMesh | None, placements: Sequence[Placement] | None):
        self._name = name

        self._device_mesh = device_mesh
        self._placements = placements

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([StateGroup(inputs=frozenset([self._name]), outputs=frozenset([self._name]))])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return {
            self._name: distribute_tensor(
                group[self._name],
                device_mesh=self._device_mesh,
                placements=self._placements,
                src_data_rank=None  # do not communicate here
            )
        }

class ModelStateMapperGatherFullTensor(ModelStateMapper):
    """
    Gathers a single DTensor object into a full Tensor object.
    """

    def __init__(self, name: str):
        self._name = name

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([StateGroup(inputs=frozenset([self._name]), outputs=frozenset([self._name]))])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        tensor = group[self._name]

        if not isinstance(tensor, DTensor):
            raise ValueError("Cannot gather anything but DTensor")

        return {
            self._name: tensor.full_tensor()
        }

class ModelStateMapperIdentity(ModelStateMapper):
    """
    Passes a single state tensor through unchanged.
    """

    def __init__(self, name: str):
        self._name = name

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([
            StateGroup(
                inputs=frozenset([self._name]),
                outputs=frozenset([self._name])
            )
        ])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return group

class ModelStateMapperRename(ModelStateMapper):
    """
    Renames a single state tensor from `name_from` to `name_to`.
    """

    def __init__(self, name_from: str, name_to: str):
        self._name_from = name_from
        self._name_to = name_to

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([
            StateGroup(
                inputs=frozenset([self._name_from]),
                outputs=frozenset([self._name_to])
            )
        ])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        return {
            self._name_to: group[self._name_from]
        }

class ModelStateMapperSelectChildModules(ModelStateMapper):
    """
    Selects a set of keys belonging to a specific parent module (prefix) and
    renames them by removing that prefix.

    This is effectively a batch rename operation that "hoists" parameters
    from a submodule scope to the current scope.
    """

    def __init__(self, base_names: list[str], parent_name: str):
        self._base_names = base_names
        self._parent_prefix = f"{parent_name}."

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([
            StateGroup(
                inputs=frozenset([self._parent_prefix + name]),
                outputs=frozenset([name])
            )
            for name in self._base_names
        ])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        name, value = next(iter(group.items()))
        if name.startswith(self._parent_prefix):
            return {
                name[len(self._parent_prefix):]: value
            }
        else:
            return {

            }

class ModelStateMapperStackTensors(ModelStateMapper):
    """
    Stacks multiple input tensors with names `source_names` into a single output tensor with name `target_name`
    producing new `stack_dim` dimension.
    """

    def __init__(self, source_names: list[str], target_name: str, stack_dim: int):
        self._source_names = source_names
        self._target_name = target_name
        self._stack_dim = stack_dim

    def state_dependency_groups(self) -> frozenset[StateGroup]:
        return frozenset([
            StateGroup(
                inputs=frozenset(self._source_names),
                outputs=frozenset([self._target_name])
            )
        ])

    def apply(self, group: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        source_tensors = [group[name] for name in self._source_names]
        return {
            self._target_name: torch.stack(source_tensors, dim=self._stack_dim)
        }