from __future__ import annotations
from .taxonomy import IDENTITY_LABELS

_MHS_COL = {
    "race": "target_race",
    "religion": "target_religion",
    "gender": "target_gender",
    "sexuality": "target_sexuality",
    "disability": "target_disability",
    "origin": "target_origin",
    "age": "target_age",
}


def _truthy(v) -> float:
    if v is None:
        return -1.0
    try:
        return 1.0 if float(v) >= 0.5 else 0.0
    except (TypeError, ValueError):
        return 1.0 if bool(v) else 0.0


def load_identity_examples(max_rows=None, streaming=True):
    from datasets import load_dataset

    try:
        ds = load_dataset(
            "ucberkeley-dlab/measuring-hate-speech", split="train", streaming=streaming
        )
    except Exception as e:
        print(f"[identity] measuring-hate-speech unavailable ({e!r}); skipping")
        return
    n = 0
    for ex in ds:
        text = ex.get("text")
        if not text:
            continue
        labels, mask = ([], [])
        for lab in IDENTITY_LABELS:
            v = _truthy(ex.get(_MHS_COL[lab]))
            if v < 0:
                labels.append(0.0)
                mask.append(0.0)
            else:
                labels.append(v)
                mask.append(1.0)
        yield {"text": text, "labels": labels, "mask": mask}
        n += 1
        if max_rows and n >= max_rows:
            return


def licenses():
    return [("ucberkeley-dlab/measuring-hate-speech", "CC-BY-4.0", True)]