"""Abstract interfaces for every model in the stack."""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Optional

import numpy as np


class BaseWorldModel(ABC):
    """
    Predicts future world state from a context window of video frames.

    Input convention
    ----------------
    frames : (T, H, W, C) uint8 RGB
    actions : (T, action_dim) float32, optional — action conditioning

    Output
    ------
    world_embedding : (latent_dim,) float32
    """

    @abstractmethod
    def encode(self, frames: np.ndarray) -> np.ndarray:
        """Encode context frames into a latent world state."""

    @abstractmethod
    def predict_future(
        self,
        frames: np.ndarray,
        actions: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """
        Auto-regressively generate future video frames.
        Returns (T_future, H, W, C) uint8.
        """


class BasePolicy(ABC):
    """
    Generates an action chunk from the world embedding + task specification.

    Output
    ------
    actions : (chunk_len, action_dim) float32
    """

    @abstractmethod
    def act(
        self,
        world_embedding: np.ndarray,
        task_embedding: np.ndarray,
    ) -> np.ndarray:
        """Return (chunk_len, action_dim) float32."""


class BaseDiffusionPlanner(ABC):
    """
    Plans a trajectory from current state toward the goal using diffusion
    denoising. The planner runs at brain frequency (5 Hz) and outputs a
    longer-horizon path that conditions the low-level VLA policy.

    Input
    -----
    world_embedding : (latent_dim,) current world state from the world model
    task_embedding  : (embed_dim,) language-grounded goal from the VLA
    proprio         : (proprio_dim,) current robot joint / pose state

    Output
    ------
    trajectory : (horizon, action_dim) float32 — planned path of waypoints
    """

    @abstractmethod
    def plan(
        self,
        world_embedding: np.ndarray,
        task_embedding: np.ndarray,
        proprio: np.ndarray,
    ) -> np.ndarray:
        """Return (horizon, action_dim) float32 planned trajectory."""


class BaseLanguageModel(ABC):
    """
    Grounds natural-language instructions in visual observations.

    Returns a task embedding that conditions the policy.
    """

    @abstractmethod
    def encode(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        """
        Parameters
        ----------
        instruction : natural-language task description
        frame       : (H, W, C) uint8 RGB current frame

        Returns
        -------
        (embedding_dim,) float32
        """

    @abstractmethod
    def predict_action(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        """Direct action prediction (language → action_dim)."""