"""Base class for all dataset format loaders."""

from __future__ import annotations

from abc import ABC, abstractmethod

import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo


class BaseLoader(ABC):
    """Abstract base class for dataset format loaders.

    All loaders produce the same output contract:
    - DatasetInfo metadata dataclass
    - pd.DataFrame with columns: [episode_index, frame_index, action, state, timestamp]
    """

    @abstractmethod
    def can_load(self, source: str) -> bool:
        """Return True if this loader can handle the given source."""
        ...

    @abstractmethod
    def load_info(self, source: str) -> DatasetInfo:
        """Load dataset metadata without loading full episode data."""
        ...

    @abstractmethod
    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data as DataFrame.

        Returns DataFrame with columns:
            episode_index (int), frame_index (int), action (list/array),
            state (list/array or None), timestamp (float)
        """
        ...


def sample_episode_indices(total: int, max_samples: int = 100) -> list[int]:
    """Select evenly-spaced episode indices when total exceeds max_samples."""
    if total <= max_samples:
        return list(range(total))
    step = total / max_samples
    return [int(i * step) for i in range(max_samples)]


def _find_column(df: pd.DataFrame, target: str, candidates: list[str]) -> str | None:
    """Find the best matching column name from a list of candidates.

    Tries exact matches first, then substring matches on the column names.
    """
    # Already has the target column
    if target in df.columns:
        return target

    # Try exact candidate matches
    for candidate in candidates:
        if candidate in df.columns:
            return candidate

    # Try substring matching (e.g., 'action' in 'joint_action_command')
    hints = [target]  # always try the target name as a substring
    for col in df.columns:
        col_lower = col.lower()
        for candidate in candidates:
            if candidate.lower() in col_lower:
                return col

    return None


# Common names for action columns across robotics datasets.
# Order matters — more specific names first, then generic.
_ACTION_CANDIDATES = [
    "action",
    "actions",
    "observation.action",
    "end_pose",
    "delta_ee",
    "delta_action",
    "joint_command",
    "joint_cmd",
    "target_q",
    "desired_joint_positions",
    "commanded_joint_positions",
    "ee_command",
    "action_dict",
    "joint_action",
    "gripper_action",
]

# Common names for state/observation columns.
_STATE_CANDIDATES = [
    "state",
    "observation.state",
    "obs",
    "observation",
    "joint_positions",
    "joint_states",
    "joint_pos",
    "qpos",
    "robot_state",
    "proprioception",
    "proprio",
    "ee_pos",
    "eef_pos",
    "observation.environment_state",
    "start_pos",
]


def normalize_shapes(shapes: dict[str, list[int]]) -> dict[str, list[int]]:
    """Add canonical 'action'/'state' keys to shapes dict if missing.

    Mirrors the column renaming that normalize_dataframe() does so that
    DatasetInfo.shapes reflects the actual action/state dimensions after
    key detection.
    """
    if "action" not in shapes:
        for candidate in _ACTION_CANDIDATES:
            if candidate in shapes:
                shapes["action"] = shapes[candidate]
                break
    if "state" not in shapes:
        for candidate in _STATE_CANDIDATES:
            if candidate in shapes:
                shapes["state"] = shapes[candidate]
                break
    return shapes


def normalize_dataframe(
    df: pd.DataFrame,
    action_key: str | None = None,
    state_key: str | None = None,
) -> pd.DataFrame:
    """Ensure DataFrame has the expected columns in the correct order.

    Detects non-standard action and state column names (e.g., 'end_pose',
    'joint_command', 'qpos') and maps them to the canonical 'action' and
    'state' columns.

    Args:
        df: Raw DataFrame from a loader.
        action_key: User-specified action column name (overrides auto-detection).
        state_key: User-specified state column name (overrides auto-detection).
    """
    import logging

    logger = logging.getLogger(__name__)

    if df.empty:
        return pd.DataFrame(
            columns=["episode_index", "frame_index", "action", "state", "timestamp"]
        )

    # Find and rename action column
    if "action" not in df.columns:
        if action_key and action_key in df.columns:
            df = df.rename(columns={action_key: "action"})
        else:
            action_col = _find_column(df, "action", _ACTION_CANDIDATES)
            if action_col is not None:
                df = df.rename(columns={action_col: "action"})

    # Find and rename state column
    if "state" not in df.columns:
        if state_key and state_key in df.columns:
            df = df.rename(columns={state_key: "state"})
        else:
            state_col = _find_column(df, "state", _STATE_CANDIDATES)
            if state_col is not None:
                df = df.rename(columns={state_col: "state"})

    # Warn when action/state columns could not be found
    available_cols = [c for c in df.columns if c not in ("episode_index", "frame_index", "timestamp")]
    if "action" not in df.columns:
        logger.warning(
            "Could not find action data. Looked for columns: %s. "
            "Found columns: %s. Use --action-key <name> to specify the action column.",
            ", ".join(["action"] + _ACTION_CANDIDATES[:6]),
            ", ".join(available_cols[:15]),
        )
    if "state" not in df.columns:
        logger.warning(
            "Could not find state data. Looked for columns: %s. "
            "Found columns: %s. Use --state-key <name> to specify the state column.",
            ", ".join(["state"] + _STATE_CANDIDATES[:6]),
            ", ".join(available_cols[:15]),
        )

    expected = ["episode_index", "frame_index", "action", "state", "timestamp"]
    for col in expected:
        if col not in df.columns:
            df[col] = None

    return df[expected + [c for c in df.columns if c not in expected]]
