"""End-to-end NER accuracy test against ground-truth crime data.

This integration test:

1. Reads complaint texts from ``tests/data/complaints.csv``.
2. Runs the full extraction pipeline against a live Ollama instance.
3. Compares each extraction against ground-truth labels in:
   - ``tests/data/crimes_perceived_detailed.csv``
   - ``tests/data/perceived_suspects.csv``
4. Reports per-field accuracy and asserts minimum thresholds.

Run with::

    pytest tests/test_ner_accuracy.py -v -m integration

Skip if Ollama is unavailable::

    pytest tests/ -m "not integration"
"""

from __future__ import annotations

import difflib
import json
import sys
from pathlib import Path

# Add project root to PYTHONPATH to allow running this script directly
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import numpy as np
import pandas as pd
import pytest
from scipy.optimize import linear_sum_assignment
from tqdm import tqdm

from tests.schema.crime_schema import CrimeExtraction, create_crime_extractor

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

DATA_DIR = Path(__file__).parent / "data"
COMPLAINTS_CSV = DATA_DIR / "complaints.csv"
CRIMES_CSV = DATA_DIR / "crimes_perceived_detailed.csv"
SUSPECTS_CSV = DATA_DIR / "perceived_suspects.csv"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _parse_range(s: str) -> tuple[float, float] | None:
    """Try to parse a string like '165-180' or '165' as a numeric range."""
    s = s.strip()
    if "-" in s:
        parts = s.split("-", 1)
        try:
            return float(parts[0]), float(parts[1])
        except ValueError:
            return None
    try:
        v = float(s)
        return v, v
    except ValueError:
        return None


def _safe_comp(reference: str, predicted: str) -> float:
    """Fuzzy string comparison returning a score in [0.0, 1.0].

    Returns 1.0 when:
    * strings are identical (case-insensitive, ``_`` -> space),
    * reference is the special wildcard ``"any"``,
    * both are numeric ranges and their bounds match exactly, or
    * both are non-null, non-numeric and SequenceMatcher similarity > 0.80.
    """
    rv = str(reference).lower().replace("nan", "null").replace("_", " ")
    lv = str(predicted).lower().replace("none", "null").replace("_", " ")

    if rv == lv or rv == "any":
        return 1.0
    if rv == "null" or lv == "null":
        return 0.0

    # Numeric / range comparison — exact bounds match
    ref_range = _parse_range(rv)
    pred_range = _parse_range(lv)
    if ref_range is not None and pred_range is not None:
        return 1.0 if ref_range == pred_range else 0.0

    # Fallback: fuzzy string match for non-numeric fields
    similarity = difflib.SequenceMatcher(None, rv, lv).ratio()
    if similarity > 0.80:
        return 1.0
    return 0.0


# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------

SUSPECT_ATTRS = ["genero", "nacionalidad", "edad", "altura", "color_pelo"]
SUSPECT_ALIASES = ["género", "nacionalidad", "edad", "altura", "color_pelo"]


@pytest.mark.integration
def test_ner_full_accuracy() -> None:
    """Measure extraction accuracy across all complaint records."""
    df_complaints = pd.read_csv(COMPLAINTS_CSV, sep=";")
    df_crimes = pd.read_csv(CRIMES_CSV, sep=";")
    df_suspects = pd.read_csv(SUSPECTS_CSV, sep=";")

    extractor = create_crime_extractor()

    total = len(df_complaints)
    hits_date = 0
    hits_address = 0
    hits_mo = 0
    hits_num_suspects = 0
    hits_suspect_attrs: float = 0.0
    total_expected_attrs = 0

    print(f"\nRunning accuracy test against {total} ground-truth records...")

    for _, row in tqdm(df_complaints.iterrows(), total=total, desc="Processing"):
        crime_id: str = row["crime_id"]
        complaint_text: str = row["complaint_text"]

        tqdm.write(f"\n{'=' * 60}\n[{crime_id}]\n{'-' * 60}")

        extraction: CrimeExtraction | None = extractor.extract_one(  # type: ignore[assignment]
            complaint_text, retry_on_null=True
        )
        if extraction is None:
            tqdm.write("  FATAL: extractor returned None - skipping.")
            continue

        real_crime = df_crimes[df_crimes["crime_id"] == crime_id].iloc[0]
        real_suspects = df_suspects[df_suspects["crime_id"] == crime_id]

        # ── Datetime ─────────────────────────────────────────────────────
        real_date = str(real_crime.get("datetime_perc", "nan"))
        extracted_date_field = (
            extraction.facts.datetime_perc if extraction.facts else None
        )
        extracted_date = (
            extracted_date_field.value
            if extracted_date_field and extracted_date_field.value is not None
            else None
        )
        evidence_date = (
            extracted_date_field.evidence
            if extracted_date_field and extracted_date_field.evidence is not None
            else "null"
        )
        if real_date == "nan":
            match_date = extracted_date is None
        else:
            match_date = extracted_date is not None and real_date[:13] in str(
                extracted_date
            )
        if match_date:
            hits_date += 1
        tqdm.write(
            f"  [{'OK' if match_date else 'XX'}] datetime  "
            f"expected='{real_date}'  got='{extracted_date}'  ev='{evidence_date}'"
        )

        # ── Address (Calle) ──────────────────────────────────────────────
        real_address = str(real_crime.get("address_perc", "nan"))
        extracted_address_field = (
            extraction.facts.address_perc if extraction.facts else None
        )
        extracted_address = (
            extracted_address_field.value
            if extracted_address_field and extracted_address_field.value is not None
            else "null"
        )
        evidence_address = (
            extracted_address_field.evidence
            if extracted_address_field and extracted_address_field.evidence is not None
            else "null"
        )
        # We can use _safe_comp for fuzzy text matching
        match_address = (
            _safe_comp(real_address, extracted_address) == 1.0
            if real_address != "nan"
            else extracted_address == "null" or extracted_address is None
        )
        if match_address:
            hits_address += 1
        tqdm.write(
            f"  [{'OK' if match_address else 'XX'}] address_perc  "
            f"expected='{real_address}'  got='{extracted_address}'  ev='{evidence_address}'"
        )

        # ── Modus operandi ───────────────────────────────────────────────
        try:
            mo_dict = json.loads(real_crime.get("modus_operandi_perc", "{}"))
            real_weapon = mo_dict.get("mo_arma", "null")
        except (json.JSONDecodeError, TypeError):
            real_weapon = "null"

        extracted_weapon_field = (
            extraction.modus_operandi.weapon_mo if extraction.modus_operandi else None
        )
        extracted_weapon = (
            extracted_weapon_field.value
            if extracted_weapon_field and extracted_weapon_field.value is not None
            else "null"
        )
        evidence_weapon = (
            extracted_weapon_field.evidence
            if extracted_weapon_field and extracted_weapon_field.evidence is not None
            else "null"
        )
        match_mo = str(real_weapon).lower() == str(extracted_weapon).lower()
        if match_mo:
            hits_mo += 1
        tqdm.write(
            f"  [{'OK' if match_mo else 'XX'}] weapon_mo  "
            f"expected='{real_weapon}'  got='{extracted_weapon}'  ev='{evidence_weapon}'"
        )

        # ── Number of suspects ───────────────────────────────────────────
        real_num = str(real_crime.get("num_perceived_suspects", "nan")).replace(
            ".0", ""
        )
        extracted_num_field = (
            extraction.facts.num_perceived_suspects if extraction.facts else None
        )
        extracted_num = (
            str(extracted_num_field.value)
            if extracted_num_field and extracted_num_field.value is not None
            else str(len(extraction.suspects))
        )
        evidence_num = (
            str(extracted_num_field.evidence)
            if extracted_num_field and extracted_num_field.evidence is not None
            else "null"
        )
        match_num = str(real_num) == extracted_num or (
            real_num == "nan" and extracted_num in ("0", "null")
        )
        if match_num:
            hits_num_suspects += 1
        tqdm.write(
            f"  [{'OK' if match_num else 'XX'}] num_suspects  "
            f"expected='{real_num}'  got='{extracted_num}'  ev='{evidence_num}'"
        )

        # ── Suspect attributes ───────────────────────────────────────────
        suspects_extracted = [s.to_row() for s in extraction.suspects]
        total_expected_attrs += len(real_suspects) * len(SUSPECT_ALIASES)

        if len(real_suspects) == 0:
            status = "OK" if len(suspects_extracted) == 0 else "XX"
            tqdm.write(f"  [{status}] suspects (expected 0)")
            continue

        if len(suspects_extracted) == 0:
            tqdm.write(f"  [XX] suspects - expected {len(real_suspects)}, got 0")
            continue

        if len(real_suspects) == 1 and len(suspects_extracted) == 1:
            real_s = real_suspects.iloc[0]
            llm_s = suspects_extracted[0]
            for alias in SUSPECT_ALIASES:
                score = _safe_comp(
                    str(real_s.get(alias, "null")),
                    str(llm_s.get(alias, "null")),
                )
                hits_suspect_attrs += score
                ev = llm_s.get(f"evidence_{alias}", "null")
                tqdm.write(
                    f"    [{'OK' if score > 0 else 'XX'}] {alias}  "
                    f"expected='{real_s.get(alias, 'null')}'  "
                    f"got='{llm_s.get(alias, 'null')}'  ev='{ev}'"
                )
        else:
            # Hungarian algorithm for optimal suspect matching
            cost_matrix = np.zeros((len(real_suspects), len(suspects_extracted)))
            real_list = list(real_suspects.iterrows())
            for i, (_, real_s) in enumerate(real_list):
                for j, llm_s in enumerate(suspects_extracted):
                    score = sum(
                        _safe_comp(
                            str(real_s.get(a, "null")),
                            str(llm_s.get(a, "null")),
                        )
                        for a in SUSPECT_ALIASES
                    )
                    cost_matrix[i, j] = len(SUSPECT_ALIASES) - score

            row_ind, col_ind = linear_sum_assignment(cost_matrix)
            for r_idx, c_idx in zip(row_ind, col_ind, strict=True):
                _, real_s = real_list[r_idx]
                llm_s = suspects_extracted[c_idx]
                tqdm.write(
                    f"  --- Match: ground-truth #{r_idx + 1} <-> "
                    f"extracted #{c_idx + 1} ---"
                )
                for alias in SUSPECT_ALIASES:
                    score = _safe_comp(
                        str(real_s.get(alias, "null")),
                        str(llm_s.get(alias, "null")),
                    )
                    hits_suspect_attrs += score
                    ev = llm_s.get(f"evidence_{alias}", "null")
                    tqdm.write(
                        f"    [{'OK' if score > 0 else 'XX'}] {alias}  "
                        f"expected='{real_s.get(alias, 'null')}'  "
                        f"got='{llm_s.get(alias, 'null')}'  ev='{ev}'"
                    )

    # ── Final report ──────────────────────────────────────────────────────
    acc_date = hits_date / total if total else 0.0
    acc_address = hits_address / total if total else 0.0
    acc_mo = hits_mo / total if total else 0.0
    acc_num = hits_num_suspects / total if total else 0.0
    acc_attrs = (
        hits_suspect_attrs / total_expected_attrs if total_expected_attrs else 0.0
    )

    print(
        f"\n{'=' * 60}\n"
        f"  NER ACCURACY REPORT ({total} complaints)\n"
        f"{'=' * 60}\n"
        f"  Datetime accuracy        : {acc_date:.0%}\n"
        f"  Address accuracy         : {acc_address:.0%}\n"
        f"  Modus operandi accuracy  : {acc_mo:.0%}\n"
        f"  Num suspects accuracy    : {acc_num:.0%}\n"
        f"  Suspect attributes acc.  : {acc_attrs:.0%}\n"
        f"{'=' * 60}"
    )

    assert acc_date > 0.0, f"Datetime accuracy too low: {acc_date:.0%}"
    assert acc_address > 0.0, f"Address accuracy too low: {acc_address:.0%}"
    assert acc_mo > 0.0, f"Modus operandi accuracy too low: {acc_mo:.0%}"
    assert acc_attrs > 0.0, f"Suspect attribute accuracy too low: {acc_attrs:.0%}"


if __name__ == "__main__":
    test_ner_full_accuracy()
