src/matcha_sentiment/data.py

from __future__ import annotations

from pathlib import Path
from typing import Sequence

import pandas as pd

from .config import LABEL2ID
from .text import compact_for_key, normalize_label, normalize_text


BAD_TEXT_VALUES = {"", "x", "-", ".", "n/a", "na", "none", "null"}


def load_binary_dataset(path: str | Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    required = {"text", "label", "label_name"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"Dataset is missing required columns: {sorted(missing)}")
    df = df.copy()
    df["text"] = df["text"].map(normalize_text)
    df["label"] = df["label"].astype(int)
    df["label_name"] = df["label_name"].map(normalize_label)
    df = df[df["label_name"].isin(LABEL2ID)]
    df = df[df["text"].str.len() > 0]
    return df.reset_index(drop=True)


def prepare_binary_dataset(
    input_path: str | Path | Sequence[str | Path],
    output_path: str | Path,
    *,
    sheet_name: str = "Data",
    dedupe: bool = True,
    balance: bool = False,
    random_state: int = 42,
) -> tuple[pd.DataFrame, dict]:
    if isinstance(input_path, (str, Path)):
        input_paths = [Path(input_path)]
    else:
        input_paths = [Path(path) for path in input_path]
    output_path = Path(output_path)

    rows: list[dict] = []
    seen: set[str] = set()
    summary = {
        "input_files": [path.name for path in input_paths],
        "output_file": str(output_path).replace("\\", "/"),
        "original_rows": 0,
        "dropped_netral": 0,
        "dropped_other_label": 0,
        "dropped_bad_text": 0,
        "dropped_duplicates": 0,
        "balanced": bool(balance),
        "random_state": int(random_state),
    }

    for path in input_paths:
        raw = pd.read_excel(path, sheet_name=sheet_name)
        raw.columns = [normalize_text(c) for c in raw.columns]
        summary["original_rows"] += int(len(raw))

        for _, row in raw.iterrows():
            label_name = normalize_label(row.get("sentimen"))
            if label_name == "Netral":
                summary["dropped_netral"] += 1
                continue
            if label_name not in LABEL2ID:
                summary["dropped_other_label"] += 1
                continue

            source_column = "perbaikan"
            text = normalize_text(row.get("perbaikan"))
            if not text:
                source_column = "textTranslated"
                text = normalize_text(row.get("textTranslated"))
            if not text:
                source_column = "text"
                text = normalize_text(row.get("text"))

            if text.lower() in BAD_TEXT_VALUES:
                summary["dropped_bad_text"] += 1
                continue

            key = compact_for_key(text)
            if dedupe and key in seen:
                summary["dropped_duplicates"] += 1
                continue
            seen.add(key)

            rows.append(
                {
                    "text": text,
                    "label": LABEL2ID[label_name],
                    "label_name": label_name,
                    "kategori": normalize_text(row.get("kategori")),
                    "stars": normalize_text(row.get("stars")),
                    "source_column": source_column,
                    "source_file": path.name,
                }
            )

    columns = ["text", "label", "label_name", "kategori", "stars", "source_column", "source_file"]
    df = pd.DataFrame(rows, columns=columns)
    summary["kept_rows_before_balance"] = int(len(df))
    summary["labels_before_balance"] = df["label_name"].value_counts().to_dict()

    if balance and not df.empty:
        label_counts = df["label_name"].value_counts()
        if len(label_counts) < len(LABEL2ID):
            raise ValueError(f"Cannot balance dataset with labels: {label_counts.to_dict()}")
        target = int(label_counts.min())
        balanced_frames = [
            group.sample(n=target, random_state=random_state)
            for _, group in df.groupby("label_name", sort=False)
        ]
        df = pd.concat(balanced_frames, ignore_index=True).sample(frac=1, random_state=random_state).reset_index(drop=True)
        summary["balance_target_per_label"] = target
        summary["dropped_by_balance"] = int(summary["kept_rows_before_balance"] - len(df))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_path, index=False)
    summary["kept_rows"] = int(len(df))
    summary["labels"] = df["label_name"].value_counts().to_dict()
    return df, summary