src/utils/loader.py

"""Data loading utilities for fetching experiment results from HuggingFace."""
import pandas as pd
import streamlit as st
from huggingface_hub import hf_hub_download, list_repo_files

import json
from pathlib import Path
from typing import Any, Optional

REPO_ID = "withcomment/mad"
REPO_TYPE = "dataset"

TO_NORMALIZE = ('tilde_mu', 'tilde_pi')
def normalize(d: dict[str, float]) -> dict[str, float]:
    """Normalize a distribution"""
    total = sum(d.values())
    if total == 0:
        return {k: 0.0 for k in d}
    return {k: v / total for k, v in d.items()}

@st.cache_data(ttl=3600)
def get_available_runs() -> dict[str, Any]:
    """List all available run directories in the HuggingFace repo.

    Returns:
        A list of records: {"run_path": str, "config": Optional[dict]}
    """
    try:
        files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
        # Find directories containing manifest.parquet
        runs = set()
        for f in files:
            if f.endswith("manifest.parquet"):
                # Extract run path (everything before the dataset name)
                parts = f.split("/")
                if len(parts) >= 2:
                    run_path = "/".join(parts[:-2])  # e.g., "mcmc/2026-01-10/12-30-00"
                    runs.add(run_path)
        run_paths = sorted(runs)
        return {p: download_config(p) for p in run_paths}
    except Exception as e:
        st.warning(f"Failed to list runs: {e}")
        return {}


def _get_nested_value(obj: Any, path: list[str]) -> Any:
    """Get a nested value from dict/list structures using a path.

    Supports dict keys and list indices (when the path segment is an int).
    Returns None when the path cannot be fully resolved.
    """
    current: Any = obj
    for segment in path:
        if isinstance(current, dict):
            if segment not in current:
                return None
            current = current[segment]
            continue

        if isinstance(current, list):
            if not segment.isdigit():
                return None
            idx = int(segment)
            if idx < 0 or idx >= len(current):
                return None
            current = current[idx]
            continue

        return None
    return current


def _parse_filter_value(raw: str) -> Any:
    """Parse a user-provided filter value.

    Tries JSON first to support numbers/bools/null/lists/dicts; falls back to string.
    """
    raw = raw.strip()
    if raw == "":
        return ""
    try:
        import json

        return json.loads(raw)
    except Exception:
        return raw


def filter_runs_by_config_query(
    runs: list[dict[str, Any]],
    query: str,
) -> tuple[list[dict[str, Any]], Optional[str]]:
    """Filter run records by a config query of the form key1.key2=value.

    Args:
        runs: Output of list_available_runs().
        query: A string like "a.b.c=123".

    Returns:
        (filtered_runs, error_message). error_message is None when successful.
    """
    query = (query or "").strip()
    if not query:
        return runs, None

    if "=" not in query:
        return [], "Invalid filter: expected 'path.to.key=value'"

    left, right = query.split("=", 1)
    path = [p for p in left.strip().split(".") if p]
    if not path:
        return [], "Invalid filter: empty key path"

    target_value = _parse_filter_value(right)

    filtered: list[dict[str, Any]] = []
    for record in runs:
        config = record.get("config")
        if not isinstance(config, dict):
            continue
        actual_value = _get_nested_value(config, path)
        if actual_value == target_value:
            filtered.append(record)

    return filtered, None


@st.cache_data(ttl=3600)
def list_datasets_for_run(run_path: str) -> list[str]:
    """List available datasets for a specific run."""
    try:
        files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
        datasets = set()
        for f in files:
            if f.startswith(run_path) and f.endswith("manifest.parquet"):
                parts = f.split("/")
                # Dataset name is the second-to-last part
                if len(parts) >= 2:
                    datasets.add(parts[-2])
        return sorted(datasets)
    except Exception as e:
        st.warning(f"Failed to list datasets: {e}")
        return []


@st.cache_data(ttl=3600)
def hf_download(run_path: str, dataset: str, file_type: str, suffix: str = 'parquet') -> Optional[pd.DataFrame]:
    """
    Download a parquet file from HuggingFace.

    Args:
        run_path: Path to the run (e.g., "mcmc/2026-01-10/12-30-00")
        dataset: Dataset name (e.g., "csqa")
        file_type: Either "manifest" or "traces"

    Returns:
        DataFrame or None if download fails
    """
    try:
        file_path = f"{run_path}/{dataset}/{file_type}.{suffix}"
        local_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=file_path,
            repo_type=REPO_TYPE,
        )
        if suffix == 'json':
            with open(local_path, 'r') as f:
                data = json.load(f)
        elif suffix == 'parquet':
            data = pd.read_parquet(local_path)
        return data
    except Exception as e:
        st.warning(f"Failed to download {file_type}.{suffix}: {e}")
        return None


@st.cache_data(ttl=3600)
def download_config(run_path: str) -> Optional[dict]:
    """Download config.json for a run."""
    try:
        file_path = f"{run_path}/config.json"
        local_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=file_path,
            repo_type=REPO_TYPE,
        )
        import json
        with open(local_path, 'r') as f:
            return json.load(f)
    except Exception as e:
        st.warning(f"Failed to download config: {e}")
        return None


def load_manifest(run_path: str, dataset: str) -> Optional[pd.DataFrame]:
    """Load manifest dataframe for a run/dataset."""
    df = hf_download(run_path, dataset, "manifest")
    if df is None:
        return None
    df['correct'] = df['argmax'] == df['label']
    return df


def load_metadata(run_path: str, dataset: str) -> Optional[dict]:
    return hf_download(run_path, dataset, "metadata", suffix="json")


def load_traces(run_path: str, dataset: str) -> Optional[pd.DataFrame]:
    """Load traces dataframe for a run/dataset."""
    traces = hf_download(run_path, dataset, "traces")
    if traces is None:
        return None
    for name in TO_NORMALIZE:
        if name not in traces.columns:
            continue
        traces[name] = traces[name].apply(lambda d: normalize(d) if isinstance(d, dict) else d)
    return traces

def load_merged(run_path: str, dataset: str) -> Optional[pd.DataFrame]:
    """Load and merge manifest and traces dataframes for a run/dataset."""
    manifest_df = load_manifest(run_path, dataset)
    traces_df = load_traces(run_path, dataset)
    return traces_df.merge(manifest_df, on='chain_idx')


def get_traces_for_chain(traces_df: pd.DataFrame, chain_idx: int) -> pd.DataFrame:
    """Get all traces for a specific chain, sorted by step."""
    chain_traces = traces_df[traces_df['chain_idx'] == chain_idx].copy()
    return chain_traces.sort_values('step')


def extract_run_metadata(runs: list[str]) -> pd.DataFrame:
    """
    Extract metadata from run paths for filtering.
    Returns DataFrame with columns for filtering.
    """
    records = []
    for run in runs:
        parts = run.split("/")
        records.append({
            'run_path': run,
            'date': parts[1] if len(parts) > 1 else '',
            'time': parts[2] if len(parts) > 2 else '',
        })
    return pd.DataFrame(records)

def get_run_type(run_path: str) -> str:
    return run_path.split("/")[0]

def get_run_timestamp(run_path: str) -> str:
    parts = run_path.split("/")
    if len(parts) >= 3:
        return f"{parts[1]}/{parts[2]}"
    return ""

def get_unique_config_values(runs: dict[str, Any], config_path: list[str]) -> list:
    """Extract unique values for a config path across all runs."""
    values = set()
    for config in runs.values():
        if not config:
            continue
        val = _get_nested_value(config, config_path)
        if val is not None:
            # Convert lists/dicts to string for display
            if isinstance(val, (list, dict)):
                import json
                val = json.dumps(val)
            values.add(val)
    return sorted(values, key=lambda x: (isinstance(x, str), x))


def filter_runs_by_config_value(runs: dict[str, Any], key: list[str] | str, selected_values: list) -> dict[str, Any]:
    """Filter runs by selected config values."""
    if not selected_values:
        return runs
    filtered = {}
    for r, config in runs.items():
        if not config:
            continue

        val = _get_nested_value(config, key)
        if val is not None:
            # Convert lists/dicts to string for comparison
            if isinstance(val, (list, dict)):
                import json
                val = json.dumps(val)
            if val in selected_values:
                filtered[r] = config
    return filtered

@st.cache_data(ttl=3600)
def load_data(run_path: str, dataset: str) -> Optional[pd.DataFrame]:
    """Load merged data (manifest + traces) for a run and dataset."""
    return load_merged(run_path, dataset)