NexQuant/rdagent/scenarios/data_science/debug/data.py

import json
import os
import shutil
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
from tqdm import tqdm

try:
    import bson  # pip install pymongo
except:
    pass


class DataHandler:
    """Base DataHandler interface."""

    def load(self, path) -> pd.DataFrame:
        raise NotImplementedError

    def dump(self, df: pd.DataFrame, path):
        raise NotImplementedError


class GenericDataHandler(DataHandler):
    """
    A generic data handler that automatically detects file type based on suffix
    and uses the correct pandas method for load/dump.
    """

    def load(self, path) -> pd.DataFrame:
        path = Path(path)
        suffix = path.suffix.lower()

        if suffix == ".csv":
            return pd.read_csv(path, encoding="utf-8")
        elif suffix == ".pkl":
            return pd.read_pickle(path)
        elif suffix == ".parquet":
            return pd.read_parquet(path)
        elif suffix in [".h5", ".hdf", ".hdf5"]:
            # Note: for HDF, you need a 'key' in read_hdf. If you expect a single key,
            # you might do: pd.read_hdf(path, key='df') or something similar.
            # Adjust as needed based on your HDF structure.
            return pd.read_hdf(path, key="data")
        elif suffix == ".jsonl":
            # Read JSON Lines file
            return pd.read_json(path, lines=True)
        elif suffix == ".json":
            # Not each json file is able to be converted to a DataFrame
            try:
                return pd.read_json(path, lines=False)
            except:
                return None
        elif suffix == ".bson":
            data = bson.decode_file_iter(open(path, "rb"))
            df = pd.DataFrame(data)
            return df
        else:
            raise ValueError(f"Unsupported file type: {suffix}")

    def dump(self, df: pd.DataFrame | dict, path):
        path = Path(path)
        suffix = path.suffix.lower()

        if suffix == ".csv":
            df.to_csv(path, index=False, encoding="utf-8")
        elif suffix == ".pkl":
            df.to_pickle(path)
        elif suffix == ".parquet":
            df.to_parquet(path, index=True)
        elif suffix in [".h5", ".hdf", ".hdf5"]:
            # Similarly, you need a key for HDF.
            df.to_hdf(path, key="data", mode="w")
        elif suffix == ".jsonl":
            # Save DataFrame to JSON Lines file
            df.to_json(path, orient="records", lines=True)
        elif suffix == ".json":
            df.to_json(path, orient="records", lines=False)
        elif suffix == ".bson":
            data = df.to_dict(orient="records")
            with open(path, "wb") as file:
                # Write each record in the list to the BSON file
                for record in data:
                    file.write(bson.BSON.encode(record))
        else:
            raise ValueError(f"Unsupported file type: {suffix}")


class DataReducer:
    """Base DataReducer interface."""

    def __init__(self, min_frac=0.02, min_num=5):
        self.min_frac = min_frac
        self.min_num = min_num
        self.sampled_files = []

    def reduce(self, df: pd.DataFrame) -> pd.DataFrame:
        raise NotImplementedError


class RandDataReducer(DataReducer):
    """
    Example random sampler: ensures at least `min_num` rows
    or at least `min_frac` fraction of the data (whichever is larger).
    """

    def reduce(self, df: pd.DataFrame, frac: float = None) -> pd.DataFrame:
        frac = max(self.min_frac, self.min_num / len(df)) if frac is None else frac
        # print(f"Sampling {frac * 100:.2f}% of the data ({len(df)} rows)")
        if frac >= 1:
            return df
        return df.sample(frac=frac, random_state=1)


class FolderReducer(DataReducer):
    """
    Sample folder from a large number of folders.
    """

    def reduce(self, array: list, frac: float = None) -> list:
        frac = max(self.min_frac, self.min_num / len(array)) if frac is None else frac
        if frac >= 1:
            return array
        train_items = [x for x in array if "train" in str(x)]
        test_items = [x for x in array if "test" in str(x)]

        # 至少保留一个 train 和一个 test
        mandatory = []
        if train_items:
            mandatory.append(np.random.choice(train_items, size=1, replace=False)[0])
        if test_items:
            mandatory.append(np.random.choice(test_items, size=1, replace=False)[0])
        mandatory.extend(np.random.choice(array, size=int(len(array) * frac) - len(mandatory), replace=False))
        return mandatory


class FileReducer(DataReducer):
    """
    Sample file from a large number of files, keep min_num of files for each folder.
    """

    def reduce(self, files: list[Path]) -> list:
        folder_dict = defaultdict(list)
        for file in files:
            folder_dict[file.parent].append(file)

        sampled_files = []
        for folder, folder_files in folder_dict.items():
            n = min(max(int(len(folder_files) * self.min_frac), self.min_num), len(folder_files))
            sampled_files.extend(np.random.choice(folder_files, size=n, replace=False))
        return sampled_files


class FileKeepReducer(DataReducer):
    """
    Sample file from a large number of files, keep min_num of files for each folder.
    """

    def reduce(self, files: list[Path]) -> list:
        folder_dict = defaultdict(list)
        for file in files:
            folder_dict[file.parent].append(file)

        sampled_files = []
        max_num = max(len(folder_files) for folder_files in folder_dict.values())
        for folder, folder_files in folder_dict.items():
            print(f"[INFO] Folder {folder} contains {len(folder_files)} files.")
            if len(folder_files) < max_num * self.min_frac:
                print(f"[INFO] Folder {folder} less than {max_num * self.min_frac} files.")
                sampled_files.extend(folder_files)
                continue
            n = min(max(int(len(folder_files) * self.min_frac), self.min_num), len(folder_files))
            sampled_files.extend(np.random.choice(folder_files, size=n, replace=False))
        return sampled_files


class SingleFileReducer(DataReducer):
    """
    Sample file from a large number of files, keep at least 1 file.
    """

    def reduce(self, files: list[Path]) -> list:
        n = min(max(int(len(files) * self.min_frac), 1), len(files))
        return np.random.choice(files, size=n, replace=False)


class UniqueIDDataReducer(DataReducer):
    def reduce(self, df: pd.DataFrame) -> pd.DataFrame:
        if not len(df):
            return df

        random_reducer = RandDataReducer(self.min_frac, self.min_num)
        if not isinstance(df, pd.DataFrame):
            return random_reducer.reduce(df)

        def is_valid_label(column):
            if not isinstance(column.iloc[0], (int, float, str, tuple, frozenset, bytes, complex, type(None))):
                return False

            if not (0 < column.nunique() < df.shape[0] * 0.5):
                return False

            if pd.api.types.is_numeric_dtype(column) and all(isinstance(x, float) for x in column.dropna()):
                return False

            return True

        label_col = df.iloc[:, -1]
        if not is_valid_label(label_col) and df.shape[1] > 2:
            label_col = df.iloc[:, 1]

        if not is_valid_label(label_col):
            return random_reducer.reduce(df)

        unique_labels = label_col.unique()
        unique_count = len(unique_labels)
        print(f"Unique labels: {unique_count} / {df.shape[0]}")

        sampled_rows = df.groupby(label_col, group_keys=False).apply(lambda x: x.sample(n=1, random_state=1))
        frac = max(self.min_frac, self.min_num / len(df))

        if int(len(df) * frac) < unique_count:
            return sampled_rows.reset_index(drop=True)

        remain_df = df.drop(index=sampled_rows.index)
        remaining_frac = frac - unique_count / len(df)

        remaining_sampled = random_reducer.reduce(remain_df, remaining_frac)
        result_df = pd.concat([sampled_rows, remaining_sampled]).sort_index()
        return result_df


class JsonReducer(DataReducer):

    def extract_filename(self, item: Any) -> Optional[str]:
        if isinstance(item, str):
            return item

        if isinstance(item, dict):
            for key in ("file_name", "filename", "path", "file", "url"):
                if key in item and isinstance(item[key], str):
                    return item[key]

            for v in item.values():
                if isinstance(v, str):
                    if "/" in v or re.search(r"\.\w{2,4}$", v):
                        return v

        return None

    def reduce(self, data: dict) -> dict:
        """
        1. 找到最大列表
        2. 随机采样并替换
        """
        candidates: List[Tuple[Union[Dict, str, int, List], Union[str, int], List[Any]]] = []
        self._find_all_lists(data, None, None, candidates)

        for parent, key, lst in sorted(candidates, key=lambda x: len(x[2]), reverse=True):
            sampled = self._sample_list(lst)
            if isinstance(parent, dict):
                parent[key] = sampled  # type: ignore
            else:
                parent[key] = sampled  # type: ignore  # parent 是 list，key 是 index, list.__setitem__(key, sampled)
            self.sampled_files.extend([self.extract_filename(i) for i in sampled])
            break
        if len(self.sampled_files) <= 0:
            raise AssertionError("sampled_files must contain at least one file")
        return data

    def _find_all_lists(
        self,
        current: Any,
        parent: Union[Dict, List, None],
        key: Union[str, int, None],
        out: List[Tuple[Union[Dict, List], Union[str, int], List[Any]]],
    ) -> None:
        """
        out => (parent_container, key_or_index, the_list)。
        """
        if isinstance(current, dict):
            for k, v in current.items():
                if isinstance(v, list):
                    out.append((current, k, v))
                    self._find_all_lists(v, current, k, out)
                elif isinstance(v, (dict, list)):
                    self._find_all_lists(v, current, k, out)

        elif isinstance(current, list):
            if parent is not None and key is not None:
                out.append((parent, key, current))
            for idx, item in enumerate(current):
                if isinstance(item, (dict, list)):
                    self._find_all_lists(item, current, idx, out)

    def _sample_list(self, lst: List[Any]) -> List[Any]:
        target = max(self.min_num, int(len(lst) * self.min_frac))
        if target >= len(lst):
            return lst[:]
        return np.random.choice(lst, size=target, replace=False)


class DataSampler:
    """Base DataSampler interface."""

    def __init__(self, data_folder, sample_folder, reducer):
        self.data_folder = data_folder
        self.sample_folder = sample_folder
        self.data_reducer = reducer
        self.included_extensions = {".csv", ".pkl", ".parquet", ".h5", ".hdf", ".hdf5", ".jsonl", ".bson"}
        self.data_handler = GenericDataHandler()

    def sample(self) -> None:
        raise NotImplementedError


class DefaultSampler(DataSampler):
    def sample(self) -> None:
        # Traverse the folder and exclude specific file types, without json currently

        files_to_process = [file for file in self.data_folder.rglob("*") if file.is_file()]
        file_types_count = count_files_in_folder(files_to_process)
        sample_json = False
        if isinstance(self.data_reducer, JsonReducer):
            self.included_extensions.add(".json")
            sample_json = True

        skip_subfolder_data = any(
            f.is_file() and f.suffix in self.included_extensions
            for f in self.data_folder.iterdir()
            if f.name.startswith(("train", "test"))
        )
        processed_files = []
        sample_used_file_names = set()
        has_id_col = False

        for file_path in tqdm(files_to_process, desc="Processing data", unit="file"):
            sampled_file_path = self.sample_folder / file_path.relative_to(self.data_folder)
            if sampled_file_path.exists():
                continue

            if file_path.suffix.lower() not in self.included_extensions:
                continue

            if skip_subfolder_data and file_path.parent != self.data_folder:
                continue  # bypass files in subfolders

            sampled_file_path.parent.mkdir(parents=True, exist_ok=True)

            # Load the original data
            if sample_json:
                if file_path.suffix.lower() == ".json":
                    data = json.load(file_path.open())
                    data_sampled = self.data_reducer.reduce(data)
                    sample_used_file_names = [file_path.parent / i for i in self.data_reducer.sampled_files]
                    print("sample_used_file_names", len(sample_used_file_names))
            else:
                df = self.data_handler.load(file_path)
                if df is None:
                    continue

                # Create a sampled subset
                df_sampled = self.data_reducer.reduce(df)
                processed_files.append(file_path)
                # Dump the sampled data
                try:
                    self.data_handler.dump(df_sampled, sampled_file_path)
                    # Extract possible file references from the sampled data
                    if "submission" in file_path.stem:
                        continue  # Skip submission files
                    for col in df_sampled.columns:
                        if "id" in col:
                            has_id_col = True
                            sample_used_file_names.extend([df_sampled[col].astype(str).unique()])
                            continue
                    for col in df_sampled.columns:
                        sample_used_file_names.extend([df_sampled[col].astype(str).unique()])
                except Exception as e:
                    print(f"Error processing {file_path}: {e}")
                    continue

        # Process non-data files
        subfolder_dict = {}
        global_groups = defaultdict(list)
        for file_path in files_to_process:
            if file_path in processed_files:
                continue  # Already handled above
            rel_dir = file_path.relative_to(self.data_folder).parts[0]
            subfolder_dict.setdefault(rel_dir, []).append(file_path)
            global_groups[file_path.stem].append(Path(file_path))

        # For each subfolder, decide which files to copy
        selected_groups = []
        extra_tag = [".txt", ".json"]
        for rel_dir, file_list in tqdm(subfolder_dict.items(), desc="Processing files", unit="file"):
            used_files = []
            not_used_files = []
            extra_files = []

            # Check if each file is in the "used" list
            for fp in file_list:
                if (
                    str(fp.name) in sample_used_file_names
                    or str(fp.stem) in sample_used_file_names
                    or fp in sample_used_file_names
                ):
                    used_files.append(fp)
                else:
                    for tag in extra_tag:
                        if file_types_count.get(tag, 1000) < 100 and fp.suffix.lower() == tag:
                            extra_files.append(fp)
                    not_used_files.append(fp)

            # Directly copy used files
            for uf in used_files:
                copy_file(uf, self.sample_folder, self.data_folder)

            # If no files are used, randomly sample files to keep the folder from being empty
            if len(used_files) == 0:
                if len(file_list) <= self.data_reducer.min_num:
                    num_to_keep = len(file_list)
                else:
                    num_to_keep = max(int(len(file_list) * self.data_reducer.min_frac), self.data_reducer.min_num)

                # Use a greedy strategy to select groups so that the total number of files is as close as possible to num_to_keep
                total_files = 0
                np.random.shuffle(not_used_files)
                for nf in not_used_files:
                    if total_files > num_to_keep:
                        break
                    if nf.stem in selected_groups:
                        total_files += 1
                    else:
                        selected_groups.append(nf.stem)
                        total_files += 1

                print(f"Sampling {num_to_keep} files without label from {total_files} files in {rel_dir}")

                # Flatten the selected groups into a single list of files
                sampled_not_used = [
                    nf for group, value in global_groups.items() if group in selected_groups for nf in value
                ]

                # Copy the selected files to the target directory (all files with the same base name will be copied)
                for nf in sampled_not_used:
                    # Construct the target path based on the relative path of nf from data_folder
                    sampled_file_path = self.sample_folder / nf.relative_to(self.data_folder)
                    if sampled_file_path.exists():
                        continue
                    sampled_file_path.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy(nf, sampled_file_path)

            # Copy extra files
            print(f"Copying {len(extra_files)} extra files")
            for uf in extra_files:
                copy_file(uf, self.sample_folder, self.data_folder)

        final_files_count = sum(1 for _ in self.sample_folder.rglob("*") if _.is_file())
        print(
            f"[INFO] After sampling, the sample folder `{self.sample_folder}` contains {final_files_count} files in total."
        )


class FolderSampler(DataSampler):
    """
    Sample data from a large number of folders.
    """

    def sample(self) -> None:
        sample_used_file_names = []
        current_level = [d for d in self.data_folder.iterdir() if d.is_dir()]
        last_count = 0
        subdirs = []
        sample_dirs = []
        sample_files = []
        extra_files = [d for d in self.data_folder.iterdir() if d.is_file()]
        level = 1
        while current_level:
            subdirs = [d for current_dir in current_level for d in current_dir.iterdir() if d.is_dir()]
            subdirs_names = [d.name for d in subdirs]
            extra_files.extend([d for current_dir in current_level for d in current_dir.iterdir() if d.is_file()])
            if not subdirs:
                print("current_level", len(current_level))
                subfiles = [d for current_dir in current_level for d in current_dir.iterdir() if d.is_file()]
                sample_files = self.data_reducer.reduce(subfiles)
                extra_files = list(set(extra_files) - set(subfiles))
                print(f"sample {len(sample_files)} files from {len(subfiles)}")
                break

            print(
                f"subdirs count: {len(set(subdirs_names))}, last_count: {last_count}, subdirs[0]: {subdirs[0]}, sample_used_file_names count: {len(set(sample_used_file_names))}"
            )
            if sample_used_file_names and set(sample_used_file_names).issubset(set(subdirs_names)):
                sample_dirs = [d for d in subdirs if d.name in sample_used_file_names]
                print(f"sample {len(sample_dirs)} folders from {len(subdirs)}")
                break

            if len(subdirs_names) > 100 or (last_count and 1 < len(sample_dirs) < last_count):
                sample_dirs = self.data_reducer.reduce(subdirs)
                print(f"sample {len(sample_dirs)} folders from {len(subdirs)}")
                break
            last_count = len(set(subdirs_names))
            current_level = subdirs
            level += 1

        print(
            f"[INFO] After sampling, the sample folder `{self.sample_folder}` contains extra_files {len(extra_files)} folders in total."
        )
        for i in sample_dirs:
            copy_folder(i, self.sample_folder, self.data_folder)
        for i in sample_files:
            copy_file(i, self.sample_folder, self.data_folder)
        for i in set(extra_files):
            copy_file(i, self.sample_folder, self.data_folder)


class SingleFilePerFolderSampler(DataSampler):
    """
    For each leaf (final) subfolder under data_folder, keep exactly one file (randomly chosen).
    Files in non-leaf folders are copied unchanged.
    """

    def sample(self) -> None:
        data_folder = Path(self.data_folder)
        sample_folder = Path(self.sample_folder)

        # Find all leaf directories (no subdirectories)
        leaf_dirs = [Path(root) for root, dirs, _ in os.walk(data_folder) if not dirs]
        print(f"Found {len(leaf_dirs)} leaf directories")

        # Sample one file per leaf directory
        for leaf in tqdm(leaf_dirs, desc="Processing files", unit="file"):
            files = [f for f in leaf.iterdir() if f.is_file()]
            if not files:
                continue
            chosen = self.data_reducer.reduce(files)
            for f in chosen:
                copy_file(f, sample_folder, data_folder)

        # Copy all files in non-leaf directories
        # i.e. any file whose parent is not a leaf dir
        # Copy all files in non-leaf directories
        for root, _, files in os.walk(data_folder):
            current_dir = Path(root)
            if current_dir in leaf_dirs:
                continue
            for fname in files:
                file_path = current_dir / fname
                copy_file(file_path, sample_folder, data_folder)

        total = sum(1 for _ in sample_folder.rglob("*") if _.is_file())
        print(f"[INFO] SingleFilePerFolderSampler: copied {total} files to {sample_folder}")


def copy_file(src_fp, target_folder, data_folder):
    """
    Construct the target file path based on the file's relative location from data_folder,
    then copy the file if it doesn't already exist.
    """
    target_fp = target_folder / src_fp.relative_to(data_folder)
    if not target_fp.exists():
        target_fp.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy(src_fp, target_fp)


def copy_folder(src_fp, target_folder, data_folder):
    """
    Copy a folder recursively.
    """
    target_fp = target_folder / src_fp.relative_to(data_folder)
    if not target_fp.exists():
        target_fp.parent.mkdir(parents=True, exist_ok=True)
        shutil.copytree(src_fp, target_fp)


def count_files_in_folder(files_to_process):
    """
    Count the number of each file type in a folder, including files in subfolders.
    """
    total_files_count = len(files_to_process)
    print(f"[INFO] Original dataset folder has {total_files_count} files in total (including subfolders).")
    file_types_count = Counter(file.suffix.lower() for file in files_to_process)
    print("File type counts:")
    for file_type, count in file_types_count.items():
        print(f"{file_type}: {count}")
    return file_types_count


def map_competition(competition: str) -> tuple[DataReducer, DataSampler]:
    cls_map = {
        "google-research-identify-contrails-reduce-global-warming": (FolderReducer, FolderSampler),
        "smartphone-decimeter-2022": (FolderReducer, FolderSampler),
        "herbarium-2020-fgvc7": (SingleFileReducer, SingleFilePerFolderSampler),
        "herbarium-2021-fgvc8": (SingleFileReducer, SingleFilePerFolderSampler),
        "herbarium-2022-fgvc9": (SingleFileReducer, SingleFilePerFolderSampler),
        "vesuvius-challenge-ink-detection": (FileReducer, FolderSampler),
        "3d-object-detection-for-autonomous-vehicles": (FileKeepReducer, FolderSampler),
    }
    return cls_map.get(competition, (UniqueIDDataReducer, DefaultSampler))


def create_debug_data(
    competition: str,
    dataset_path: str | Path,
    min_frac=0.01,
    min_num=5,
    sample_path=None,
):
    """
    Reads the original data file, creates a reduced sample,
    and renames/moves files for easier debugging.
    Automatically detects file type (csv, pkl, parquet, hdf, etc.).
    """
    if sample_path is None:
        sample_path = Path(dataset_path) / "sample"

    # Prepare data handler and reducer
    reduce_method, sample_method = map_competition(competition)
    data_reducer = reduce_method(min_frac=min_frac, min_num=min_num)
    sampler = sample_method(Path(dataset_path) / competition, Path(sample_path) / competition, data_reducer)
    print(f"processing {competition}, sample_method: {sample_method}, reduce_method: {reduce_method}")
    sampler.sample()