from __future__ import annotations

from dataclasses import dataclass, replace

from wisent.core.contrastive_pairs.core.atoms import AtomResponse
from wisent.core.activations.core.atoms import LayerActivations, RawActivationMap

__all__ = [
    "Response",
    "PositiveResponse",
    "NegativeResponse",
]


@dataclass(frozen=True, slots=True)
class Response(AtomResponse):
    """A model's response to a prompt, with optional activations and label.

    attributes:
        model_response: The text response generated by the model.
        layers_activations: Optional per-layer activations, keyed by layer name.
            See **What is LayerActivations?** below for details.
        label: Optional label for the response, e.g., "harmless", "toxic", etc.

    what is LayerActivations?
        'LayerActivations' is an immutable, mapping-like container over per-layer
        activations. It behaves like a 'Mapping[str, torch.Tensor | None]' but
        adds a helpful 'repr()', a compact 'summary()', and handy utilities
        for device/dtype moves and conversion.

        keys:
            Layer names as strings (e.g., "blocks.0.mlp", "attn.3").

        values:
            Either a 'torch.Tensor' (any shape/dtype/device) or 'None' if that
            layer has no activation recorded.

        construction and coercion:
            You can pass:
            - a 'LayerActivations' instance, or
            - a plain dict 'dict[str, torch.Tensor | np.ndarray | None]'.
            NumPy arrays are converted to tensors; tensors are optionally cast
            to a given dtype if provided by the wrapper.

        methods:
            - 'summary()' → small dict of shape/dtype/device per layer.
            - 'to(*args, **kwargs)' → like 'Tensor.to' for all non-'None' values.
            - 'cpu()', 'detach()' → convenience variants.
            - 'numpy()' → convert stored tensors to NumPy arrays (on cpu).
            - 'to_dict()' → plain 'dict[str, torch.Tensor | None]'.

    serialization notes:
        'Response.to_dict()' returns tensors as tensors. This is convenient for
        in-process use but not JSON-serializable. For JSON, consider mapping the
        activations to shapes/metadata (via 'summary()') or to NumPy arrays /
        lists (via 'numpy()' → then '.tolist()') before encoding.

    examples:
        >>> resp = Response(
        ...     model_response="OK",
        ...     layers_activations={"blocks.0.mlp": torch.randn(2, 4), "attn.1": None},
        ...     label="harmless",
        ... )
        >>> print(resp.layers_activations)
        LayerActivations(
          blocks.0.mlp: Tensor(shape=(2, 4), dtype=torch.float32, device=cpu)
          attn.1: None
        )
        >>> resp.layers_activations.summary()
        {'blocks.0.mlp': {'shape': (2, 4), 'dtype': 'torch.float32', 'device': 'cpu', 'requires_grad': False},
         'attn.1': {'shape': None, 'dtype': None, 'device': None, 'requires_grad': None}}

        # Update fields immutably:
        >>> resp2 = resp.with_label("toxic")
        >>> resp3 = resp.with_activations({"blocks.0.mlp": torch.zeros(2, 4)})
    """

    model_response: str
    layers_activations: LayerActivations | None = None
    label: str | None = None

    def __post_init__(self) -> None:
        if not isinstance(self.model_response, str) or not self.model_response.strip():
            raise ValueError("'model_response' must be a non-empty string.")
        la = self.layers_activations
        if la is None or isinstance(la, LayerActivations):
            coerced = la
        else:
            coerced = LayerActivations(la) 
        object.__setattr__(self, "layers_activations", coerced)

    def with_activations(self, layers_activations: LayerActivations | RawActivationMap | None) -> Response:
        new_val = layers_activations if isinstance(layers_activations, LayerActivations) or layers_activations is None \
                  else LayerActivations(layers_activations)
        return replace(self, layers_activations=new_val)

    def with_label(self, label: str | None) -> Response:
        return replace(self, label=label)

    def to_dict(self) -> dict[str, RawActivationMap | str | None]:
        """Return a plain dict representation of this Response.

        returns:
            A dictionary with keys 'model_response', 'layers_activations', and 'label'.

        example:
            {
                "model_response": "OK",
                "layers_activations": {"blocks.0.mlp": torch.randn(2, 4), "attn.1": None},
                "label": "harmless"
            }
        """
        return {
            "model_response": self.model_response,
            "layers_activations": (
                None if self.layers_activations is None else self.layers_activations.to_dict()
            ),
            "label": self.label,
        }

    @classmethod
    def from_dict(cls, data: dict[str, str | RawActivationMap | None]) -> Response:
        ''' Create a Response from a plain dict.

        arguments:
            data: A dictionary with keys 'model_response', 'layers_activations', and 'label'.   
                    'layers_activations' should be a dict or None.
        raises:
            ValueError: If 'model_response' is missing or not a non-empty string.

        example:
         >>> data = {
         ...     "model_response": "OK",
         ...     "layers_activations": {"blocks.0.mlp": torch.randn(2, 4), "attn.1": None},
         ...     "label": "harmless"
         ... }
         >>> resp = Response.from_dict(data)
         >>> print(resp)
         Response(model_response='OK', layers_activations=LayerActivations(...), label='harmless')          
        '''
        return cls(
            model_response=str(data["model_response"]),
            layers_activations=(
                None if data.get("layers_activations") is None
                else LayerActivations(data["layers_activations"])  
            ),
            label=data.get("label") if isinstance(data.get("label"), str) else None,
        )


class PositiveResponse(Response): ...
class NegativeResponse(Response): ...