"""Universal dataset loader — handles ANY robotics data format.

Instead of relying on column names to find actions and states, this loader
uses the ArrayProfile intelligence layer to scan ALL numeric arrays in a
dataset, profile each one statistically, and pick the ones that look most
like robot actions and observations.

Works on top of the existing loader registry — once format is detected and
raw data is loaded, this module adds the "understanding" layer.

Usage:
    from orbit.loaders.universal import universal_load

    info, episodes_df, array_profiles = universal_load("path/to/dataset")
"""

from __future__ import annotations

import logging

import numpy as np
import pandas as pd

from orbit.analyzer.array_intelligence import (
    ArrayProfile,
    identify_action_vs_state,
    profile_array,
    profile_from_dataframe,
)
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.loaders import detect_and_load_episodes, detect_and_load_info

logger = logging.getLogger(__name__)


def universal_load(
    source: str,
    max_episodes: int | None = None,
    format: str | None = None,
    action_key: str | None = None,
    state_key: str | None = None,
    **loader_kwargs,
) -> tuple[DatasetInfo, pd.DataFrame, dict[str, ArrayProfile]]:
    """Load a dataset from any format with intelligent array profiling.

    Strategy:
    1. Use existing loader registry to detect format and load raw data.
    2. Profile every numeric column with ArrayProfile.
    3. If action/state columns are ambiguous, use statistical inference
       to identify which arrays are actions vs states.
    4. Return standardized output with array profiles attached.

    Args:
        source: Path, file, or HuggingFace repo ID.
        max_episodes: Maximum episodes to load.
        format: Explicit format override.
        action_key: Explicit action column name (skips auto-detection).
        state_key: Explicit state column name (skips auto-detection).
        **loader_kwargs: Extra kwargs for the loader.

    Returns:
        Tuple of (DatasetInfo, DataFrame, dict of array profiles).
        The dict maps column names ("action", "state", etc.) to ArrayProfile.
    """
    # Step 1: Load metadata
    info = detect_and_load_info(source, format=format, **loader_kwargs)

    # Step 2: Load episode data
    kw = dict(loader_kwargs)
    if action_key:
        kw["action_key"] = action_key
    if state_key:
        kw["state_key"] = state_key

    df = detect_and_load_episodes(
        source,
        max_episodes=max_episodes,
        format=format,
        **kw,
    )

    # Step 3: Profile all available numeric array columns
    profiles: dict[str, ArrayProfile] = {}

    # Profile the action column
    if "action" in df.columns and not df["action"].isna().all():
        profiles["action"] = profile_from_dataframe(df, "action")
        logger.info(
            "Action profile: type=%s (conf=%.2f), dims=%d, gripper=%s, dead=%s",
            profiles["action"].action_type,
            profiles["action"].action_type_confidence,
            profiles["action"].num_dimensions,
            profiles["action"].gripper_dimensions,
            profiles["action"].dead_dimensions,
        )

    # Profile the state column
    if "state" in df.columns and not df["state"].isna().all():
        profiles["state"] = profile_from_dataframe(df, "state")
        logger.info(
            "State profile: type=%s (conf=%.2f), dims=%d",
            profiles["state"].action_type,
            profiles["state"].action_type_confidence,
            profiles["state"].num_dimensions,
        )

    # Step 4: If no action column was found, try to identify it
    if "action" not in profiles:
        logger.warning(
            "No action column found. Scanning all numeric columns..."
        )
        profiles = _scan_all_columns(df, profiles)

        # Use intelligent identification
        if len(profiles) >= 2:
            assignments = identify_action_vs_state(profiles)
            for col_name, role in assignments.items():
                if role == "action" and "action" not in df.columns:
                    df = df.rename(columns={col_name: "action"})
                    logger.info("Identified '%s' as action column", col_name)
                elif role == "state" and "state" not in df.columns:
                    df = df.rename(columns={col_name: "state"})
                    logger.info("Identified '%s' as state column", col_name)

    return info, df, profiles


def _scan_all_columns(
    df: pd.DataFrame,
    existing_profiles: dict[str, ArrayProfile],
) -> dict[str, ArrayProfile]:
    """Scan all numeric columns in a DataFrame and profile each one."""
    skip_cols = {"episode_index", "frame_index", "timestamp", "action", "state"}
    profiles = dict(existing_profiles)

    for col in df.columns:
        if col in skip_cols or col in profiles:
            continue

        # Check if column contains array-like data
        sample = df[col].dropna().head(5)
        if len(sample) == 0:
            continue

        first = sample.iloc[0]
        if isinstance(first, (list, np.ndarray)):
            try:
                prof = profile_from_dataframe(df, col)
                if prof.num_dimensions > 0:
                    profiles[col] = prof
            except Exception as e:
                logger.debug("Could not profile column '%s': %s", col, e)
        elif isinstance(first, (int, float, np.integer, np.floating)):
            # Scalar column — skip for now (not array data)
            pass

    return profiles


def profile_loaded_episodes(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> dict[str, ArrayProfile]:
    """Profile pre-loaded episode arrays (alternative to universal_load).

    For cases where you already have loaded episodes as numpy arrays
    (e.g., from the CLI analysis pipeline).

    Args:
        action_episodes: List of per-episode action arrays.
        state_episodes: Optional list of per-episode state arrays.

    Returns:
        Dict mapping "action" (and optionally "state") to ArrayProfile.
    """
    profiles: dict[str, ArrayProfile] = {}

    if action_episodes:
        profiles["action"] = profile_array(action_episodes)

    if state_episodes:
        profiles["state"] = profile_array(state_episodes)

    return profiles
