data_loader.py

"""
Data loading utilities for cellmap challenge results viewer.
Loads evaluation results and metadata for interactive visualization.
"""

import json
import glob
import os
import pandas as pd
import numpy as np
from upath import UPath


CROP_MANIFEST_URL = (
    "https://raw.githubusercontent.com/janelia-cellmap/cellmap-segmentation-challenge/"
    "refs/heads/main/src/cellmap_segmentation_challenge/utils/test_crop_manifest.csv"
)
CONFIG = json.load(open(".config.json", "r"))


def load_crop_manifest(url=CROP_MANIFEST_URL):
    """
    Load the test crop manifest CSV and return a mapping of crop number to dataset name.

    Args:
        url: URL to the crop manifest CSV

    Returns:
        dict: Mapping of crop number string (e.g., '982') to dataset name (e.g., 'jrc_mus-nacc-1')
    """
    try:
        df = pd.read_csv(url)
        mapping = (
            df.drop_duplicates(subset=["crop_name"])[["crop_name", "dataset"]]
            .set_index("crop_name")["dataset"]
            .to_dict()
        )
        # Convert keys to strings for matching with crop names like "crop982"
        return {str(k): v for k, v in mapping.items()}
    except Exception as e:
        print(f"Warning: Could not load crop manifest from {url}: {e}")
        return {}


def load_eval_results(results_dir=CONFIG["RESULTS_DIR"], eval_ids=None):
    """
    Load eval_*.results files into a pandas DataFrame.

    Args:
        results_dir: Directory containing eval_*.results files
        eval_ids: Optional set/list of evaluation IDs to load. If None, loads all.

    Returns:
        tuple: (df_results, df_aggregate) where df_results has per-crop/organelle metrics
               and df_aggregate has overall scores
    """
    results = []
    aggregate_results = []

    # Convert to set for O(1) lookup if provided
    if eval_ids is not None:
        eval_ids = set(eval_ids)

    pattern = os.path.join(results_dir, "eval_*.results")
    files = [f for f in glob.glob(pattern) if not f.endswith("_submitted_only.results")]

    for filepath in files:
        with open(filepath, "r") as f:
            data = json.load(f)

        file_name = UPath(filepath).name
        # Extract evaluation_id from filename (eval_123.results -> 123)
        # Skip files that don't have numeric IDs (e.g., eval_test.results)
        id_str = file_name.replace("eval_", "").replace(".results", "")
        try:
            eval_id = int(id_str)
        except ValueError:
            print(f"Skipping file with non-numeric ID: {file_name}")
            continue

        # Skip if not in the allowed set of evaluation IDs
        if eval_ids is not None and eval_id not in eval_ids:
            continue

        aggregate = {"file_name": file_name, "evaluation_id": eval_id}

        for crop, crop_data in data.items():
            # Put top-level aggregates into a separate table
            if crop in {
                "overall_instance_score",
                "overall_semantic_score",
                "overall_score",
            }:
                aggregate[crop] = crop_data
            elif crop in {
                "label_scores",
                "total_evals",
                "num_evals_done",
                "status",
                "is_missing",
                "git_version",
            }:
                continue
            elif isinstance(crop_data, dict):
                for organelle, metrics in crop_data.items():
                    row = {
                        "file_name": file_name,
                        "evaluation_id": eval_id,
                        "crop": crop,
                        "organelle": organelle,
                    }

                    for k, v in metrics.items():
                        row[k] = v

                    results.append(row)
            else:
                # Unknown top-level scalar → treat as aggregate metric
                aggregate[crop] = crop_data
        aggregate_results.append(aggregate)

    return pd.DataFrame(results), pd.DataFrame(aggregate_results)


def load_evaluations_csv(csv_path):
    """
    Load evaluations.csv metadata file.

    Args:
        csv_path: Path to evaluations.csv

    Returns:
        pandas.DataFrame with evaluation metadata
    """
    df = pd.read_csv(csv_path)

    # Clean up column names (remove BOM if present)
    df.columns = [col.replace("\ufeff", "") for col in df.columns]

    # Parse created_at as datetime
    if "created_at" in df.columns:
        df["created_at"] = pd.to_datetime(df["created_at"], errors="coerce", utc=True)

    return df


def merge_data(df_results, df_aggregate, df_evals):
    """
    Merge results with evaluation metadata.

    Args:
        df_results: DataFrame from load_eval_results (per-crop/organelle)
        df_aggregate: DataFrame from load_eval_results (overall scores)
        df_evals: DataFrame from load_evaluations_csv

    Returns:
        tuple: (df_results_merged, df_aggregate_merged) with metadata joined
    """
    # Merge per-crop results with metadata
    df_results_merged = df_results.merge(
        df_evals[
            [
                "evaluation_id",
                "username",
                "submission_name",
                "data_path",
                "status",
                "created_at",
            ]
        ],
        on="evaluation_id",
        how="left",
    )

    # Merge aggregate results with metadata
    df_aggregate_merged = df_aggregate.merge(
        df_evals[
            [
                "evaluation_id",
                "username",
                "submission_name",
                "data_path",
                "status",
                "created_at",
            ]
        ],
        on="evaluation_id",
        how="left",
    )

    return df_results_merged, df_aggregate_merged


def get_filter_options(df_results, df_aggregate):
    """
    Extract unique values for filter dropdowns.

    Args:
        df_results: Merged results DataFrame
        df_aggregate: Merged aggregate DataFrame

    Returns:
        dict: {metrics: [...], crops: [...], organelles: [...]}
    """
    # Get metrics from column names
    exclude_cols = {
        "file_name",
        "evaluation_id",
        "crop",
        "organelle",
        "username",
        "submission_name",
        "data_path",
        "status",
        "status_x",
        "status_y",
        "created_at",
        "num_voxels",
        "voxel_size",
        "is_missing",
    }

    # Metrics from per-crop results
    metrics_results = [col for col in df_results.columns if col not in exclude_cols]

    # Overall metrics: derive dynamically from aggregate columns
    aggregate_meta_cols = {
        "file_name",
        "evaluation_id",
        "username",
        "submission_name",
        "data_path",
        "status",
        "status_x",
        "status_y",
        "created_at",
    }
    overall_metrics = [
        col for col in df_aggregate.columns if col not in aggregate_meta_cols
    ]

    all_metrics = sorted(set(metrics_results + overall_metrics))

    # Get unique crops and organelles
    crops = (
        sorted(df_results["crop"].unique().tolist())
        if "crop" in df_results.columns
        else []
    )
    organelles = (
        sorted(df_results["organelle"].unique().tolist())
        if "organelle" in df_results.columns
        else []
    )

    # Get created_at date range
    created_at_min = None
    created_at_max = None
    for df in [df_aggregate, df_results]:
        if "created_at" in df.columns:
            valid = df["created_at"].dropna()
            if not valid.empty:
                lo = valid.min()
                hi = valid.max()
                if created_at_min is None or lo < created_at_min:
                    created_at_min = lo
                if created_at_max is None or hi > created_at_max:
                    created_at_max = hi

    return {
        "metrics": all_metrics,
        "overall_metrics": sorted(overall_metrics),
        "crops": crops,
        "organelles": organelles,
        "created_at_min": created_at_min.date().isoformat() if created_at_min is not None else None,
        "created_at_max": created_at_max.date().isoformat() if created_at_max is not None else None,
    }


def compute_metric_series(df, metric_or_expr):
    """
    Compute metric values from DataFrame.
    Supports both simple column names and expressions.

    Args:
        df: pandas DataFrame
        metric_or_expr: Column name or expression (e.g., "accuracy", "accuracy * 0.5")

    Returns:
        pandas.Series with metric values
    """
    # Fast path: plain column name
    if metric_or_expr in df.columns:
        s = df[metric_or_expr]
        return s.astype(float)

    # Expression path
    try:
        s = pd.eval(
            metric_or_expr,
            engine="python",
            parser="pandas",
            local_dict=df.to_dict("series"),
        )
    except Exception as e:
        raise ValueError(
            f"Couldn't evaluate metric expression: {metric_or_expr!r}. "
            f"Available columns: {sorted(df.columns)}. Error: {e}"
        )

    if not isinstance(s, pd.Series):
        # pd.eval can return scalar; make it broadcast if needed
        s = pd.Series([float(s)] * len(df), index=df.index)

    return s.astype(float)


def filter_dataframe(df, crop=None, organelle=None, eval_id=None, username=None, date_from=None, date_to=None):
    """
    Filter DataFrame by crop, organelle, eval_id, username, and/or date range.

    Args:
        df: pandas DataFrame
        crop: Crop name to filter by (optional)
        organelle: Organelle name to filter by (optional)
        eval_id: Evaluation ID to filter by (optional)
        username: Username to filter by (optional)
        date_from: Start date string 'YYYY-MM-DD' inclusive (optional)
        date_to: End date string 'YYYY-MM-DD' inclusive (optional)

    Returns:
        Filtered pandas DataFrame
    """
    mask = pd.Series(True, index=df.index)

    # Apply eval_id filter first (most restrictive)
    if eval_id is not None:
        mask &= df["evaluation_id"] == int(eval_id)

    # Apply username filter
    if username and username != "":
        mask &= df["username"] == username

    # Apply date range filter on created_at
    if "created_at" in df.columns:
        if date_from:
            ts_from = pd.Timestamp(date_from, tz="UTC")
            mask &= df["created_at"] >= ts_from
        if date_to:
            ts_to = pd.Timestamp(date_to, tz="UTC") + pd.Timedelta(days=1)
            mask &= df["created_at"] < ts_to

    # Apply existing crop/organelle filters
    if crop and crop != "":
        mask &= df["crop"] == crop

    if organelle and organelle != "":
        mask &= df["organelle"] == organelle

    return df[mask]


def compute_histogram_data(
    df,
    metric,
    crop=None,
    organelle=None,
    bins=20,
    exclude_missing=True,
    eval_id=None,
    username=None,
    date_from=None,
    date_to=None,
):
    """
    Compute histogram data for a given metric with filters.

    Args:
        df: pandas DataFrame
        metric: Metric name or expression
        crop: Crop filter (optional)
        organelle: Organelle filter (optional)
        bins: Number of histogram bins
        exclude_missing: Whether to exclude missing values
        eval_id: Evaluation ID filter (optional)
        username: Username filter (optional)
        date_from: Start date 'YYYY-MM-DD' inclusive (optional)
        date_to: End date 'YYYY-MM-DD' inclusive (optional)

    Returns:
        dict: {bins: [...], counts: [...], bin_edges: [...]}
    """
    # Filter data
    filtered = filter_dataframe(df, crop, organelle, eval_id, username, date_from, date_to)

    # Exclude missing if requested
    if exclude_missing and "is_missing" in filtered.columns:
        filtered = filtered[~filtered["is_missing"]]

    # Get metric values
    values = compute_metric_series(filtered, metric)

    # Drop NaNs
    values = values[values.notna()].values

    if len(values) == 0:
        return {"bins": [], "counts": [], "bin_edges": [], "bin_width": 0}

    # Compute histogram
    counts, bin_edges = np.histogram(values, bins=bins)

    # Bin centers for plotting
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
    bin_width = bin_edges[1] - bin_edges[0] if len(bin_edges) > 1 else 0

    return {
        "bins": bin_centers.tolist(),
        "counts": counts.tolist(),
        "bin_edges": bin_edges.tolist(),
        "bin_width": bin_width,
    }


def get_submissions_in_bin(
    df,
    metric,
    bin_min=None,
    bin_max=None,
    crop=None,
    organelle=None,
    eval_id=None,
    username=None,
    date_from=None,
    date_to=None,
):
    """
    Get submissions with metric values in a specific bin range, or all submissions if no bin specified.

    Args:
        df: pandas DataFrame
        metric: Metric name
        bin_min: Minimum bin value (optional, if None returns all)
        bin_max: Maximum bin value (optional, if None returns all)
        crop: Crop filter (optional)
        organelle: Organelle filter (optional)
        eval_id: Evaluation ID filter (optional)
        username: Username filter (optional)
        date_from: Start date 'YYYY-MM-DD' inclusive (optional)
        date_to: End date 'YYYY-MM-DD' inclusive (optional)

    Returns:
        list: Submission records with eval_id, username, score, etc.
    """
    # Filter data
    filtered = filter_dataframe(df, crop, organelle, eval_id, username, date_from, date_to)

    # Get metric values
    values = compute_metric_series(filtered, metric)

    # Filter by bin range if specified
    if bin_min is not None and bin_max is not None:
        # To match numpy histogram behavior, check if this might be the last bin
        # by seeing if any values equal bin_max (last bin includes right edge)
        max_value = values.max()
        if max_value == bin_max:
            bin_mask = (values >= bin_min) & (values <= bin_max)
        else:
            bin_mask = (values >= bin_min) & (values < bin_max)

        in_bin = filtered[bin_mask].copy()
        in_bin["_metric_value"] = values[bin_mask]
    else:
        # No bin filter - return all non-NaN submissions
        valid_mask = values.notna()
        in_bin = filtered[valid_mask].copy()
        in_bin["_metric_value"] = values[valid_mask]

    # Build results using vectorized operations
    result_df = pd.DataFrame(
        {
            "eval_id": in_bin["evaluation_id"].astype(int),
            "username": in_bin["username"].fillna("Unknown"),
            "submission_name": in_bin["submission_name"].fillna(""),
            "score": in_bin["_metric_value"].astype(float).round(10),
            # "status": in_bin["status"].fillna("") if "status" in in_bin.columns else "",
            "created_at": (
                in_bin["created_at"].fillna("")
                if "created_at" in in_bin.columns
                else ""
            ),
            "crop": in_bin["crop"].fillna("") if "crop" in in_bin.columns else "",
            "organelle": (
                in_bin["organelle"].fillna("") if "organelle" in in_bin.columns else ""
            ),
        }
    )

    # Sort by score descending
    result_df = result_df.sort_values("score", ascending=False)

    return result_df.to_dict("records")