"""Explicit tokenizer pretraining and cleaning shared by both token layers."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
from .fitting import ProfileFitter

class TrainingMixin:
    def preprocess(
        self,
        values: np.ndarray,
        *,
        profile: str | None = None,
        gripper_column: int | None = None,
        **kwargs: Any,
    ) -> np.ndarray:
        """Clean a canonical ``(T,D_active)`` profile trajectory for explicit tokenizer pretraining."""
        if profile is None:
            profiles = {component.space for component in self.action_space.components}
            if len(profiles) != 1:
                raise ValueError("profile is required for layouts combining joint and EEF")
            profile = profiles.pop()
        if profile not in self.preprocessing:
            raise ValueError(f"no recorded preprocessing for profile {profile!r}")
        values = np.asarray(values)
        if "ease_limits" not in kwargs and values.ndim == 2:
            kwargs["ease_limits"] = self._joint_ease_limits(
                profile, values.shape[1], gripper_column, kwargs.get("active_dims")
            )
        return self.preprocessing[profile].apply(values, gripper_column=gripper_column, **kwargs)

    def _joint_ease_limits(self, profile, width, gripper_column, active_dims):
        processor = self.action_tokenizer.profiles[profile].processor
        component = processor.config.physical.components.get("generalized_joint_position")
        if profile != "joint" or component is None or component.primitive_order != 1:
            return None
        domain = component.constraints["order_1"]
        limits = np.full(processor.action_dim, np.inf)
        limits[list(component.dims)] = np.minimum(np.abs(domain.minimum), np.abs(domain.maximum))
        dims = tuple(range(width)) if active_dims is None else tuple(active_dims)
        limits = limits[list(dims)]
        if gripper_column is not None:
            limits[gripper_column] = np.inf
        return limits

    def fit(
        self,
        episodes,
        *,
        fps: float | None = None,
        current_state=None,
        previous_first_order=None,
        bootstrap_context: bool = False,
        apply_preprocessing: bool = False,
        preprocessing=None,
        backend: str = "auto",
        threads: int = 0,
        batch_size: int = 256,
        progress: bool = False,
        max_cell_bytes: int | None = None,
        work_dir: str | Path | None = None,
    ) -> TrainingMixin:
        """Fit vocabularies from episodes in declared source layouts."""
        fitter = ProfileFitter(self, bootstrap_context=bootstrap_context,
                      apply_preprocessing=apply_preprocessing, preprocessing=preprocessing)
        router, report = fitter.fit(
            episodes,
            work_dir=work_dir,
            metadata=dict(
                fps=fps, current_state=current_state, previous_first_order=previous_first_order
            ),
            backend=backend,
            threads=threads,
            batch_size=batch_size,
            progress=progress,
            max_cell_bytes=max_cell_bytes,
        )
        replacement = type(self)(router, action_space=self.action_space.name)
        replacement._fitted_profiles = self.fitted_profiles | report["profiles"].keys()
        replacement.training_report = report
        replacement.preprocessing = dict(fitter.preprocessing)
        replacement._warned_default_fps = self._warned_default_fps
        self.__dict__.update(replacement.__dict__)
        return self