"""Dimensionality reduction abstractions and checkpoint restoration."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import ClassVar

import numpy as np
from pydantic import BaseModel, ConfigDict, Field, JsonValue

JSONDict = dict[str, JsonValue]


class ReducerConfigModel(BaseModel):
    """Base class for reducer config schemas."""

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


class ReducerStateModel(BaseModel):
    """Base class for reducer state schemas."""

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


class ReductionPayload(BaseModel):
    """Serialized reducer payload stored in checkpoints."""

    kind: str = Field(min_length=1)
    config: JSONDict
    state: JSONDict

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


_REDUCER_REGISTRY: dict[str, type["Reducer"]] = {}


class Reducer(ABC):
    """Abstract base class for checkpointable reducers."""

    kind: ClassVar[str]

    @abstractmethod
    def fit(self, embeddings: np.ndarray) -> "Reducer":
        """Fit the reducer on embeddings."""

    @abstractmethod
    def transform(self, embeddings: np.ndarray) -> np.ndarray:
        """Transform embeddings into reduced space."""

    def fit_transform(self, embeddings: np.ndarray) -> np.ndarray:
        """Fit and transform embeddings in one step."""
        self.fit(embeddings)
        return self.transform(embeddings)

    @abstractmethod
    def checkpoint_payload(self) -> ReductionPayload:
        """Return serialized reducer payload for checkpoints."""

    @classmethod
    @abstractmethod
    def from_checkpoint_payload(cls, payload: ReductionPayload) -> "Reducer":
        """Restore a fitted reducer from a validated checkpoint payload."""


def register_reducer(reducer_cls: type[Reducer]) -> type[Reducer]:
    """Register a reducer class for checkpoint restoration."""
    reducer_kind = reducer_cls.kind
    if not reducer_kind:
        raise ValueError(
            f"Reducer class {reducer_cls.__name__} must define a non-empty kind"
        )
    existing = _REDUCER_REGISTRY.get(reducer_kind)
    if existing is not None and existing is not reducer_cls:
        raise ValueError(
            f"Reducer kind '{reducer_kind}' is already registered to "
            f"{existing.__name__}; cannot register {reducer_cls.__name__}"
        )
    _REDUCER_REGISTRY[reducer_kind] = reducer_cls
    return reducer_cls


def restore_reducer(payload: ReductionPayload | None) -> Reducer | None:
    """Restore a reducer from a checkpoint payload."""
    if payload is None:
        return None

    reducer_cls = _REDUCER_REGISTRY.get(payload.kind)
    if reducer_cls is None:
        raise ValueError(f"Unknown reducer kind '{payload.kind}' in checkpoint")

    return reducer_cls.from_checkpoint_payload(payload)
