data_manager.py

import pandas as pd
import re
from pathlib import Path
from typing import List, Tuple, Optional
import json
import math


class DataManager:
    def __init__(self, data_dir: str):
        self.data_dir = Path(data_dir)
        self.master_df = self._load_all_data()

    def _load_old_format_folder(self, period_dir: Path) -> pd.DataFrame:
        all_xlsx_data = []
        period = period_dir.name

        for file_path in period_dir.iterdir():
            if file_path.suffix != ".xlsx":
                continue

            model_group = file_path.stem

            xls = pd.read_excel(file_path, sheet_name=None)

            for sheet_name, df in xls.items():
                if df.empty:
                    continue

                df = self._clean_dataframe(df)

                df["Period"] = period
                df["Metric"] = sheet_name
                df["Model Group"] = model_group

                all_xlsx_data.append(df)

        return all_xlsx_data

    def _load_new_format_folder(self, period_dir: Path) -> pd.DataFrame:
        """
        读取新格式 JSON 数据并转换为宽表。
        """
        raw_records = []
        period = period_dir.name

        for file_path in period_dir.rglob("*.json"):
            with open(file_path, "r", encoding="utf-8") as f:
                data = json.load(f)

            full_path = data["model_name_or_path"]
            model_name = full_path.split("/")[-1].replace(".pth", "")

            params = data["parameters count"]

            data_path = data["data_path"]
            source_col = data_path.split("-")[-1]

            neg_log_prob = data["neg_log_prob_sum"]
            avg_char = data["avg character count"]
            avg_bytes = data["avg bytes"]

            metrics = {}

            metrics["cr"] = data["compression_rate"]
            metrics["bpc"] = (neg_log_prob / avg_char) * (1 / math.log(2))
            metrics["bpb"] = (neg_log_prob / avg_bytes) * (1 / math.log(2))

            for metric_type, value in metrics.items():
                if value is not None:
                    raw_records.append(
                        {
                            "Name": model_name,
                            "Params (B)": params,
                            "Period": period,
                            "Metric": metric_type,
                            "Model Group": "other",
                            "Source": source_col,
                            "Value": value,
                        }
                    )

        if not raw_records:
            return []

        df_long = pd.DataFrame(raw_records)
        df_wide = df_long.pivot_table(
            index=["Name", "Params (B)", "Period", "Metric", "Model Group"],
            columns="Source",
            values="Value",
        ).reset_index()

        df_wide.columns.name = None

        def assign_group(p):
            if p >= 13:
                return "14b"
            if p >= 8:
                return "9b"
            if p >= 6:
                return "7b"
            if p >= 2.5:
                return "3b"
            if p >= 1:
                return "1b5"
            return "other"

        df_wide["Model Group"] = df_wide["Params (B)"].apply(assign_group)

        metadata_cols = ["Name", "Params (B)", "Period", "Metric", "Model Group"]
        new_columns = {}
        for col in df_wide.columns:
            if col not in metadata_cols:
                new_columns[col] = col.replace("_", " ")

        df_wide = df_wide.rename(columns=new_columns)

        return [df_wide]

    def _load_all_data(self) -> pd.DataFrame:
        all_records = []

        if not self.data_dir.exists():
            print(f"Warning: Directory {self.data_dir} does not exist.")
            return pd.DataFrame()

        period_dirs = [d for d in self.data_dir.iterdir() if d.is_dir() and re.match(r"^\d{4}-\d{2}$", d.name)]

        for period_dir in period_dirs:
            if period_dir.name <= "2025-11":
                all_records.extend(self._load_old_format_folder(period_dir))
            else:
                all_records.extend(self._load_new_format_folder(period_dir))

        if not all_records:
            return pd.DataFrame()

        final_df = pd.concat(all_records, ignore_index=True)

        exclude_cols = ["Name", "Period", "Metric", "Model Group"]
        numeric_cols = [c for c in final_df.columns if c not in exclude_cols]
        for col in numeric_cols:
            final_df[col] = pd.to_numeric(final_df[col], errors="coerce")

        return final_df

    def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.dropna(axis=1, how="all")
        new_columns = []
        for col in df.columns:
            col_str = str(col)
            if "Parameters" in col_str:
                new_columns.append("Params (B)")
            elif col_str == "Average (The lower the better)":
                new_columns.append("Average (lower=better)")
            else:
                new_columns.append(col_str.rsplit("_", maxsplit=1)[0].replace("\u200b", ""))  # 去除_202xxxxx后缀
        df.columns = new_columns
        column_mapping = {col: col.replace("_", " ") for col in df.columns}
        df = df.rename(columns=column_mapping)
        return df

    def get_available_periods(self) -> List[str]:
        """返回所有可用的时间周期，已排序，从旧到新"""
        if self.master_df.empty:
            return []
        return sorted(self.master_df["Period"].unique().tolist())

    def get_available_columns(self, period: str) -> List[str]:
        """获取特定时间段内的数据列（排除元数据列和全部为NaN的列）"""
        if self.master_df.empty:
            return []
        subset = self.master_df[self.master_df["Period"] == period]
        if subset.empty:
            return []

        metadata_cols = ["Name", "Params (B)", "Period", "Metric", "Model Group", "Average (lower=better)"]
        return [c for c in subset.columns if c not in metadata_cols and not subset[c].isna().all()]

    def query(
        self,
        period: str,
        metric_code: str,
        param_range: Tuple[float, float],
        model_groups: Optional[List[str]] = None,
        visible_columns: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """
        统一查询接口。

        Args:
            period: 时间周期 (e.g. "2025-12")
            metric_display_name: bpc, bpb, cr
            param_range: 参数量范围元组 (min, max)
            model_groups: (可选) 文件名列表，如 ['14b', '7b']
            visible_columns: (可选) 需要参与计算平均值的列名列表
        """

        mask = (
            (self.master_df["Period"] == period)
            & (self.master_df["Metric"] == metric_code)
            & (self.master_df["Params (B)"].between(param_range[0], param_range[1]))
        )

        if model_groups:
            mask = mask & (self.master_df["Model Group"].isin(model_groups))

        filtered_df = self.master_df.loc[mask].copy()

        if filtered_df.empty:
            return filtered_df

        # 始终排除的列
        exclude_cols = ["Period", "Metric", "Model Group"]
        # 始终保留的列
        metadata_cols = ["Name", "Params (B)", "Average (lower=better)"]

        if visible_columns is not None:
            valid_visible_cols = [c for c in visible_columns if c in filtered_df.columns]
            columns_to_keep = metadata_cols + valid_visible_cols
            cols_for_average = valid_visible_cols
        else:
            all_cols = [c for c in filtered_df.columns if c not in exclude_cols]
            columns_to_keep = all_cols
            cols_for_average = [c for c in all_cols if c not in metadata_cols]

        if not cols_for_average:
            return pd.DataFrame()

        filtered_df["Average (lower=better)"] = filtered_df[cols_for_average].mean(axis=1).round(3)
        columns_to_keep = [c for c in columns_to_keep if c not in exclude_cols]
        columns_to_keep = list(dict.fromkeys(columns_to_keep))
        filtered_df = filtered_df[columns_to_keep]

        if "Name" in filtered_df.columns:
            filtered_df["Name"] = filtered_df["Name"].apply(lambda x: x.replace(".pth", ""))

        filtered_df = filtered_df.sort_values(by="Average (lower=better)", ascending=True, kind="mergesort", na_position="last").reset_index(
            drop=True
        )

        # 定义列优先级顺序
        fixed_cols = ["Name", "Params (B)", "Average (lower=better)"]
        column_priority = [
            "ao3 english",
            "bbc news",
            "wikipedia english",
            "arxiv physics",
            "arxiv cs",
            "arxiv math",
            "github cpp",
            "github python",
            "github javascript",
            "github markdown",
        ]
        existing_cols = filtered_df.columns.tolist()
        ordered_cols = []
        for col in fixed_cols:
            if col in existing_cols:
                ordered_cols.append(col)
        for col in column_priority:
            if col in existing_cols and col not in ordered_cols:
                ordered_cols.append(col)
        for col in existing_cols:
            if col not in ordered_cols:
                ordered_cols.append(col)
        filtered_df = filtered_df[ordered_cols]

        return filtered_df


if __name__ == "__main__":

    dm = DataManager("data")
    periods = dm.get_available_periods()

    print(f"Total records loaded: {len(dm.master_df)}")
    print(f"Available periods: {periods}")
    print(f"Available columns: {dm.get_available_columns('2025-11')}")

    result = dm.query(
        period="2025-11",
        metric_code="cr",
        param_range=(0, 20),
        model_groups=["7b"],
        visible_columns=["wikipedia_english"],
    )

    print(result.head(20))