# /// script
# dependencies = [
#     "datasets==4.5.0",
#     "gliner2==1.2.4",
#     "ipython==9.9.0",
#     "marimo>=0.19.10",
#     "requests==2.32.5",
#     "scikit-learn>=1.0.0",
#     "tqdm>=4.0.0",
#     "torch @ https://download.pytorch.org/whl/cu130/torch-2.10.0%2Bcu130-cp312-cp312-win_amd64.whl",
#     "torchvision @ https://download.pytorch.org/whl/cu130/torchvision-0.25.0%2Bcu130-cp312-cp312-win_amd64.whl",
#     "pyzmq>=27.1.0",
# ]
# requires-python = ">=3.12"
#
# ///

import marimo

__generated_with = "0.19.11"
app = marimo.App(width="medium")


@app.cell
def _():
    import time
    import torch
    from tqdm import tqdm
    from sklearn.metrics import precision_recall_fscore_support
    from gliner2 import GLiNER2
    from datasets import load_dataset

    # Load dataset and filter by language
    ds = load_dataset("ai4privacy/open-pii-masking-500k-ai4privacy", split="validation")
    ds_filtered = ds.filter(lambda x: x["language"] == "nl") # in ["en", "nl"]

    # Detect CUDA availability
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    if device == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    # Load model
    extractor = GLiNER2.from_pretrained("fastino/gliner2-multi-v1")
    extractor.to(device)

    # Define label schema with descriptions for GLiNER2
    label_schema = {
        "surname": "Family or last name of a person",
        "given_name": "First or given name of a person",
        "title": "Honorific title (e.g., Mr., Mrs., Dr., Prof.)",
        "gender": "Gender identity of a person",
        "age": "Age or date of birth of a person",
        "email": "Email address or electronic mail",
        "phone_number": "Telephone or mobile phone number",
        "date": "Calendar date (birth date, event date, appointment date)",
        "time": "Time of day (hour, minute, specific moment)",
        "street": "Street name or address line",
        "building_number": "Building or house number",
        "city": "City or municipality name",
        "zip_code": "Postal or ZIP code",
        "id_card_number": "National identity card or ID number",
        "passport_number": "Passport identification number",
        "driver_license_number": "Driver's license or driving permit number",
        "social_number": "Social security or national identification number",
        "tax_number": "Tax identification number (VAT, income tax)",
        "credit_card_number": "Credit card number or financial card number"
    }

    # Mapping from original dataset labels to new schema labels
    label_mapping = {
        "SURNAME": "surname",
        "GIVENNAME": "given_name",
        "TITLE": "title",
        "SEX": "gender",
        "GENDER": "gender",
        "AGE": "age",
        "EMAIL": "email",
        "TELEPHONENUM": "phone_number",
        "DATE": "date",
        "TIME": "time",
        "STREET": "street",
        "BUILDINGNUM": "building_number",
        "CITY": "city",
        "ZIPCODE": "zip_code",
        "IDCARDNUM": "id_card_number",
        "PASSPORTNUM": "passport_number",
        "DRIVERLICENSENUM": "driver_license_number",
        "SOCIALNUM": "social_number",
        "TAXNUM": "tax_number",
        "CREDITCARDNUMBER": "credit_card_number"
    }

    # Simple benchmark
    start = time.time()

    nl_true_labels = []
    nl_pred_labels = []
    en_true_labels = []
    en_pred_labels = []

    # Collect all samples for batch processing
    samples_list = list(ds_filtered)
    texts = []
    sample_metadata = []

    for sample in samples_list:
        text = sample["source_text"]
        language = sample["language"]
        true_entities = sample["privacy_mask"]
        labels = [item["label"] for item in true_entities]

        if not labels:
            continue

        texts.append(text)
        sample_metadata.append({
            "language": language,
            "true_entities": true_entities
        })

    print(f"Labels in schema: {list(label_schema.keys())}")

    # Batch extraction with tqdm progress
    batch_size = 32
    results = []

    for i in tqdm(range(0, len(texts), batch_size), desc="Extracting entities"):
        batch_texts = texts[i:i+batch_size]
        batch_results = extractor.batch_extract_entities(
            batch_texts,
            label_schema,
            include_spans=True,
            batch_size=16
        )
        results.extend(batch_results)

    # Process each result
    for result, metadata in zip(results, sample_metadata):
        language = metadata["language"]
        true_entities = metadata["true_entities"]

        # Flatten gliner2 result to list with spans
        pred_entities = []
        for label, entities in result["entities"].items():
            for entity in entities:
                pred_entities.append({
                    "label": label,
                    "start": entity["start"],
                    "end": entity["end"]
                })

        # Compare each true entity with predictions
        for true_entity in true_entities:
            true_label = label_mapping.get(true_entity["label"], true_entity["label"])
            true_span = (true_entity["start"], true_entity["end"])

            # Find if gliner2 predicted at this exact span
            pred_label = None
            for pred_entity in pred_entities:
                if (pred_entity["start"], pred_entity["end"]) == true_span:
                    pred_label = pred_entity["label"]
                    break

            # If no prediction found at this span, use "O" (no entity)
            if pred_label is None:
                pred_label = "O"

            if language == "nl":
                nl_true_labels.append(true_label)
                nl_pred_labels.append(pred_label)
            else:  # en
                en_true_labels.append(true_label)
                en_pred_labels.append(pred_label)

    elapsed = time.time() - start
    count = len(sample_metadata)

    # Compute metrics for NL
    if nl_true_labels:
        nl_precision, nl_recall, nl_fscore, _ = precision_recall_fscore_support(
            nl_true_labels, nl_pred_labels, average="weighted", zero_division=0
        )
    else:
        nl_precision = nl_recall = nl_fscore = 0

    # Compute metrics for EN
    if en_true_labels:
        en_precision, en_recall, en_fscore, _ = precision_recall_fscore_support(
            en_true_labels, en_pred_labels, average="weighted", zero_division=0
        )
    else:
        en_precision = en_recall = en_fscore = 0

    print(f"Processed {count} samples in {elapsed:.2f} seconds")
    print(f"Average time per sample: {elapsed/count*1000:.2f} ms")
    print(f"\n--- Dutch (NL) ---")
    print(f"Precision: {nl_precision:.4f}")
    print(f"Recall: {nl_recall:.4f}")
    print(f"F1-Score: {nl_fscore:.4f}")
    print(f"\n--- English (EN) ---")
    print(f"Precision: {en_precision:.4f}")
    print(f"Recall: {en_recall:.4f}")
    print(f"F1-Score: {en_fscore:.4f}")
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()
