"""Layer statistics helpers."""

from typing import Dict

from datasets import load_dataset

from config import REPO_ID, qwen25_dataset_path

def load_layer_statistics(layer_id: int) -> Dict:
    """Load statistics for a layer from the HuggingFace dataset."""
    try:
        data_file = qwen25_dataset_path(
            f"statistics/layer_{layer_id}-00000-of-00001.parquet"
        )
        print(f"[info] Loading statistics via data_files: layer={layer_id} file={data_file}")
        ds = load_dataset(
            REPO_ID,
            data_files={"train": [data_file]},
            split="train",
            streaming=False,
        )
        print(f"[info] statistics loaded: file={data_file} rows={len(ds)}")
        if len(ds) == 0:
            return {"total_valid_neurons": 0, "modalities": {}}
        row = ds[0]
        return {
            "total_valid_neurons": row.get("total_valid_neurons", 0),
            "modalities": row.get("modalities", {}),
        }
    except Exception as e:  # pragma: no cover - network/cache dependent
        print(f"[error] Failed to load statistics for layer {layer_id}: {e}")
        return {"total_valid_neurons": 0, "modalities": {}}


def format_layer_statistics(stats: Dict) -> str:
    """Format layer statistics into Markdown for display."""
    total_neurons = stats.get("total_valid_neurons", 0)
    modalities = stats.get("modalities", {})

    markdown = f"**Total Valid Neurons:** {total_neurons}\n\n"
    markdown += "**Modality Distribution:**\n\n"

    if modalities:
        for modality, count in modalities.items():
            percentage = (count / total_neurons * 100) if total_neurons > 0 else 0
            markdown += f"- {modality}: {count} ({percentage:.1f}%)\n"
    else:
        markdown += "No data available."

    return markdown