"""Raw directory loader for numpy/CSV/parquet files."""

from __future__ import annotations

import json
import logging
from pathlib import Path

import numpy as np
import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.loaders.base import BaseLoader, normalize_dataframe, sample_episode_indices

logger = logging.getLogger(__name__)


class DirectoryLoader(BaseLoader):
    """Load datasets from directories containing numpy, CSV, or parquet files.

    Supports three directory structures:

    Structure A: One file per episode
        dataset/
        ├── episode_000.npz   # keys: "action" (T, A), "state" (T, D)
        ├── episode_001.npz
        └── ...

    Structure B: Flat arrays + episode boundaries
        dataset/
        ├── actions.npy       # (total_frames, A)
        ├── states.npy        # (total_frames, D)
        └── episodes.json     # [{"start": 0, "end": 50}, ...]

    Structure C: CSV/parquet with episode column
        dataset/
        └── data.parquet      # columns: episode_index, action_0..action_N, state_0..state_N
    """

    def can_load(self, source: str) -> bool:
        """True if source is a directory with numpy/CSV/parquet files (not LeRobot)."""
        p = Path(source).expanduser()
        if not p.is_dir():
            return False
        # Exclude LeRobot datasets
        if (p / "meta" / "info.json").exists():
            return False
        # Check for supported files
        has_npz = bool(list(p.glob("*.npz"))[:1])
        has_npy = bool(list(p.glob("*.npy"))[:1])
        has_csv = bool(list(p.glob("*.csv"))[:1])
        has_parquet = bool(list(p.glob("*.parquet"))[:1])
        return has_npz or has_npy or has_csv or has_parquet

    def load_info(self, source: str) -> DatasetInfo:
        """Infer dataset metadata from directory contents."""
        root = Path(source).expanduser().resolve()
        structure = self._detect_structure(root)

        if structure == "per_episode_npz":
            return self._info_from_npz(root)
        elif structure == "flat_arrays":
            return self._info_from_flat(root)
        elif structure == "tabular":
            return self._info_from_tabular(root)
        else:
            raise ValueError(
                f"Cannot determine dataset structure in {source}. "
                "Expected: episode_*.npz files, actions.npy + episodes.json, "
                "or data.parquet/data.csv"
            )

    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data from directory."""
        if max_episodes is None:
            max_episodes = 100

        root = Path(source).expanduser().resolve()
        structure = self._detect_structure(root)

        if structure == "per_episode_npz":
            return self._load_npz(root, episode_indices, max_episodes)
        elif structure == "flat_arrays":
            return self._load_flat(root, episode_indices, max_episodes)
        elif structure == "tabular":
            return self._load_tabular(root, episode_indices, max_episodes)
        else:
            raise ValueError(f"Cannot determine dataset structure in {source}")

    def _detect_structure(self, root: Path) -> str | None:
        """Detect which directory structure is present."""
        # Structure A: Per-episode npz files
        npz_files = sorted(root.glob("episode_*.npz")) + sorted(root.glob("ep_*.npz"))
        if npz_files:
            return "per_episode_npz"

        # Structure B: Flat arrays
        if (root / "actions.npy").exists():
            return "flat_arrays"

        # Structure C: Tabular (parquet or CSV)
        parquet_files = list(root.glob("*.parquet"))
        csv_files = list(root.glob("*.csv"))
        if parquet_files or csv_files:
            return "tabular"

        # Also check for any npz files (not episode_* named)
        any_npz = list(root.glob("*.npz"))
        if any_npz:
            return "per_episode_npz"

        return None

    # --- Structure A: Per-episode NPZ ---

    def _info_from_npz(self, root: Path) -> DatasetInfo:
        npz_files = sorted(root.glob("*.npz"))

        num_episodes = len(npz_files)
        num_frames = 0
        shapes: dict[str, list[int]] = {}

        if npz_files:
            sample = np.load(npz_files[0])
            if "action" in sample:
                shapes["action"] = list(sample["action"].shape[1:])
                num_frames = sum(
                    np.load(f)["action"].shape[0] for f in npz_files[:10]
                )
                # Estimate total from sample
                if len(npz_files) > 10:
                    avg_len = num_frames / 10
                    num_frames = int(avg_len * num_episodes)
            if "state" in sample:
                shapes["state"] = list(sample["state"].shape[1:])

        return DatasetInfo(
            repo_id=str(root),
            num_episodes=num_episodes,
            num_frames=num_frames,
            fps=30.0,  # Default; user can override
            robot_type="unknown",
            camera_names=[],
            video_keys=[],
            shapes=shapes,
        )

    def _load_npz(
        self, root: Path, episode_indices: list[int] | None, max_episodes: int
    ) -> pd.DataFrame:
        npz_files = sorted(root.glob("*.npz"))

        if episode_indices is not None:
            npz_files = [npz_files[i] for i in episode_indices if i < len(npz_files)]
        elif len(npz_files) > max_episodes:
            sampled = sample_episode_indices(len(npz_files), max_episodes)
            npz_files = [npz_files[i] for i in sampled]

        rows: list[dict] = []
        for ep_idx, f in enumerate(npz_files):
            data = np.load(f)
            actions = data.get("action", data.get("actions"))
            states = data.get("state", data.get("states"))

            n_frames = actions.shape[0] if actions is not None else (
                states.shape[0] if states is not None else 0
            )

            for frame_idx in range(n_frames):
                rows.append(
                    {
                        "episode_index": ep_idx,
                        "frame_index": frame_idx,
                        "action": (
                            actions[frame_idx].tolist() if actions is not None else None
                        ),
                        "state": (
                            states[frame_idx].tolist() if states is not None else None
                        ),
                        "timestamp": frame_idx / 30.0,
                    }
                )

        return normalize_dataframe(pd.DataFrame(rows) if rows else pd.DataFrame())

    # --- Structure B: Flat arrays ---

    def _info_from_flat(self, root: Path) -> DatasetInfo:
        actions = np.load(root / "actions.npy")
        states_path = root / "states.npy"
        states = np.load(states_path) if states_path.exists() else None

        episodes_file = root / "episodes.json"
        if episodes_file.exists():
            with open(episodes_file) as f:
                ep_boundaries = json.load(f)
            num_episodes = len(ep_boundaries)
        else:
            num_episodes = 1

        shapes: dict[str, list[int]] = {"action": list(actions.shape[1:])}
        if states is not None:
            shapes["state"] = list(states.shape[1:])

        return DatasetInfo(
            repo_id=str(root),
            num_episodes=num_episodes,
            num_frames=actions.shape[0],
            fps=30.0,
            robot_type="unknown",
            camera_names=[],
            video_keys=[],
            shapes=shapes,
        )

    def _load_flat(
        self, root: Path, episode_indices: list[int] | None, max_episodes: int
    ) -> pd.DataFrame:
        actions = np.load(root / "actions.npy")
        states_path = root / "states.npy"
        states = np.load(states_path) if states_path.exists() else None

        episodes_file = root / "episodes.json"
        if episodes_file.exists():
            with open(episodes_file) as f:
                ep_boundaries = json.load(f)
        else:
            ep_boundaries = [{"start": 0, "end": actions.shape[0]}]

        if episode_indices is not None:
            target_eps = [
                (i, ep_boundaries[i])
                for i in episode_indices
                if i < len(ep_boundaries)
            ]
        elif len(ep_boundaries) > max_episodes:
            sampled = sample_episode_indices(len(ep_boundaries), max_episodes)
            target_eps = [(i, ep_boundaries[i]) for i in sampled]
        else:
            target_eps = list(enumerate(ep_boundaries))

        rows: list[dict] = []
        for ep_idx, bounds in target_eps:
            start = bounds["start"]
            end = bounds["end"]
            for frame_idx, global_idx in enumerate(range(start, end)):
                rows.append(
                    {
                        "episode_index": ep_idx,
                        "frame_index": frame_idx,
                        "action": actions[global_idx].tolist(),
                        "state": (
                            states[global_idx].tolist() if states is not None else None
                        ),
                        "timestamp": frame_idx / 30.0,
                    }
                )

        return normalize_dataframe(pd.DataFrame(rows) if rows else pd.DataFrame())

    # --- Structure C: Tabular ---

    def _info_from_tabular(self, root: Path) -> DatasetInfo:
        df = self._read_tabular(root)

        num_episodes = (
            df["episode_index"].nunique() if "episode_index" in df.columns else 1
        )
        num_frames = len(df)

        # Detect action/state columns
        action_cols = sorted([c for c in df.columns if c.startswith("action")])
        state_cols = sorted([c for c in df.columns if c.startswith("state")])
        shapes: dict[str, list[int]] = {}
        if action_cols:
            shapes["action"] = [len(action_cols)]
        if state_cols:
            shapes["state"] = [len(state_cols)]

        return DatasetInfo(
            repo_id=str(root),
            num_episodes=num_episodes,
            num_frames=num_frames,
            fps=30.0,
            robot_type="unknown",
            camera_names=[],
            video_keys=[],
            shapes=shapes,
        )

    def _load_tabular(
        self, root: Path, episode_indices: list[int] | None, max_episodes: int
    ) -> pd.DataFrame:
        df = self._read_tabular(root)

        # Detect action/state columns
        action_cols = sorted([c for c in df.columns if c.startswith("action_")])
        state_cols = sorted([c for c in df.columns if c.startswith("state_")])

        # If the df already has 'action' and 'state' as single columns, use directly
        has_action_col = "action" in df.columns and not action_cols
        has_state_col = "state" in df.columns and not state_cols

        # Ensure episode_index exists
        if "episode_index" not in df.columns:
            df["episode_index"] = 0

        # Filter episodes
        all_episodes = sorted(df["episode_index"].unique())
        if episode_indices is not None:
            df = df[df["episode_index"].isin(episode_indices)]
        elif len(all_episodes) > max_episodes:
            sampled = sample_episode_indices(len(all_episodes), max_episodes)
            target = [all_episodes[i] for i in sampled if i < len(all_episodes)]
            df = df[df["episode_index"].isin(target)]

        # Combine action/state columns into lists
        rows: list[dict] = []
        for _, row in df.iterrows():
            action = None
            if has_action_col:
                action = row["action"]
            elif action_cols:
                action = [float(row[c]) for c in action_cols]

            state = None
            if has_state_col:
                state = row["state"]
            elif state_cols:
                state = [float(row[c]) for c in state_cols]

            rows.append(
                {
                    "episode_index": int(row["episode_index"]),
                    "frame_index": int(row.get("frame_index", 0)),
                    "action": action,
                    "state": state,
                    "timestamp": float(row.get("timestamp", 0.0)),
                }
            )

        return normalize_dataframe(pd.DataFrame(rows) if rows else pd.DataFrame())

    def _read_tabular(self, root: Path) -> pd.DataFrame:
        """Read the first parquet or CSV file found."""
        parquet_files = sorted(root.glob("*.parquet"))
        if parquet_files:
            return pd.read_parquet(parquet_files[0])

        csv_files = sorted(root.glob("*.csv"))
        if csv_files:
            return pd.read_csv(csv_files[0])

        raise ValueError(f"No parquet or CSV files found in {root}")
