#!/usr/bin/env python3
"""Independently validate the LM04 RAID-extra OOD quality package."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import os
import tempfile
import unicodedata
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import numpy as np


PREDICTION_FIELDS = [
    "sample_index", "id", "domain", "source_model", "label", "generation_sha256",
    "token_count", "token_ids_sha256", "fp32_logit", "fp32_machine_score",
    "public_quantized_logit", "public_quantized_machine_score",
]


class Checks:
    def __init__(self) -> None:
        self.total = 0
        self.failures: list[dict[str, str]] = []
        self.categories: Counter[str] = Counter()

    def check(self, condition: bool, category: str, detail: str) -> None:
        self.total += 1
        self.categories[category] += 1
        if not condition:
            self.failures.append({"category": category, "detail": detail})


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def sha256_text(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(value, handle, indent=2, sort_keys=True, ensure_ascii=False, allow_nan=False)
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def _space(character: str) -> bool:
    return character in " \t\n\r" or unicodedata.category(character) == "Zs"


def _control(character: str) -> bool:
    return character not in "\t\n\r" and unicodedata.category(character).startswith("C")


def _punct(character: str) -> bool:
    code = ord(character)
    return (33 <= code <= 47 or 58 <= code <= 64 or 91 <= code <= 96 or 123 <= code <= 126
            or unicodedata.category(character).startswith("P"))


def _chinese(code: int) -> bool:
    ranges = ((0x4E00, 0x9FFF), (0x3400, 0x4DBF), (0x20000, 0x2A6DF),
              (0x2A700, 0x2B73F), (0x2B740, 0x2B81F), (0x2B820, 0x2CEAF),
              (0xF900, 0xFAFF), (0x2F800, 0x2FA1F))
    return any(start <= code <= end for start, end in ranges)


class IndependentWordPiece:
    """Second implementation used only by the package validator."""

    def __init__(self, vocab_path: Path) -> None:
        self.tokens = vocab_path.read_text(encoding="utf-8").splitlines()
        self.index = {token: i for i, token in enumerate(self.tokens)}

    def _basic(self, text: str) -> list[str]:
        cleaned = []
        for char in text:
            if ord(char) in (0, 0xFFFD) or _control(char):
                continue
            char = " " if _space(char) else char
            cleaned.extend((" ", char, " ") if _chinese(ord(char)) else (char,))
        output = []
        for word in "".join(cleaned).strip().split():
            word = "".join(c for c in unicodedata.normalize("NFD", word.lower())
                           if unicodedata.category(c) != "Mn")
            current = []
            for char in word:
                if _punct(char):
                    if current:
                        output.append("".join(current))
                        current = []
                    output.append(char)
                else:
                    current.append(char)
            if current:
                output.append("".join(current))
        return output

    def _wordpiece(self, token: str) -> list[str]:
        if len(token) > 100:
            return ["[UNK]"]
        output = []
        start = 0
        while start < len(token):
            match = None
            end = len(token)
            while end > start:
                candidate = token[start:end]
                if start:
                    candidate = "##" + candidate
                if candidate in self.index:
                    match = candidate
                    break
                end -= 1
            if match is None:
                return ["[UNK]"]
            output.append(match)
            start = end
        return output

    def encode(self, text: str) -> list[int]:
        pieces = []
        for token in self._basic(text):
            pieces.extend(self._wordpiece(token))
            if len(pieces) >= 510:
                pieces = pieces[:510]
                break
        return [self.index["[CLS]"], *(self.index.get(piece, self.index["[UNK]"]) for piece in pieces), self.index["[SEP]"]]


def rank(protocol: str, sample_id: str) -> str:
    return sha256_text(protocol + "\0" + sample_id)


def allocate(counts: dict[tuple[str, str], int], target: int) -> dict[tuple[str, str], int]:
    total = sum(counts.values())
    quotas = {key: target * value / total for key, value in counts.items()}
    result = {key: int(quotas[key] // 1) for key in counts}
    order = sorted(counts, key=lambda key: (-(quotas[key] - result[key]), key))
    for key in order[: target - sum(result.values())]:
        result[key] += 1
    return result


def reconstruct_sample(dataset: Path, config: dict[str, Any]) -> tuple[list[dict[str, str]], dict[str, Any]]:
    humans = []
    machines: dict[tuple[str, str], list[dict[str, str]]] = defaultdict(list)
    ids = set()
    domains = Counter()
    models = Counter()
    with dataset.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            if row["id"] in ids:
                raise ValueError(f"duplicate id: {row['id']}")
            ids.add(row["id"])
            if row["attack"] != "none":
                raise ValueError("non-clean row in extra_none")
            record = {key: row[key] for key in ("id", "source_id", "domain", "model", "generation")}
            record["rank"] = rank(config["protocol_id"], row["id"])
            domains[row["domain"]] += 1
            models[row["model"]] += 1
            (humans if row["model"] == "human" else machines[(row["domain"], row["model"])]).append(record)
    counts = {key: len(value) for key, value in machines.items()}
    allocations = allocate(counts, int(config["sample"]["machine_target_records"]))
    selected = list(humans)
    for key in sorted(machines):
        selected.extend(sorted(machines[key], key=lambda row: (row["rank"], row["id"]))[:allocations[key]])
    selected.sort(key=lambda row: (row["domain"], row["model"], row["rank"], row["id"]))
    return selected, {
        "records": len(ids), "human_records": len(humans), "machine_records": len(ids) - len(humans),
        "domains": dict(sorted(domains.items())), "models": dict(sorted(models.items())),
        "allocations": {"|".join(key): allocations[key] for key in sorted(allocations)},
        "ids_sha256": sha256_text("\n".join(row["id"] for row in selected) + "\n"),
    }


def auc(labels: list[int], scores: list[float]) -> float:
    positives = sum(labels)
    negatives = len(labels) - positives
    ordered = sorted(zip(scores, labels), key=lambda item: item[0])
    rank_sum = 0.0
    start = 0
    while start < len(ordered):
        end = start + 1
        while end < len(ordered) and ordered[end][0] == ordered[start][0]:
            end += 1
        rank_sum += (((start + 1) + end) / 2.0) * sum(label for _, label in ordered[start:end])
        start = end
    return (rank_sum - positives * (positives + 1) / 2.0) / (positives * negatives)


def fpr(human_scores: list[float], threshold: float) -> float:
    return sum(score >= threshold for score in human_scores) / len(human_scores)


def threshold_search(human_scores: list[float], target: float, epsilon: float) -> tuple[float, float]:
    threshold = sum(human_scores) / len(human_scores)
    step = 0.5
    previous = None
    found = []
    for _ in range(50):
        observed = fpr(human_scores, threshold)
        if abs(observed - target) <= epsilon:
            return threshold, observed
        found.append((threshold, observed))
        distance = target - observed
        if previous is not None and ((distance < 0) != (previous < 0)):
            step *= -0.5
        elif previous is not None and abs(distance) - abs(previous) > 0.01:
            step *= -1
        threshold += step
        previous = distance
    differences = [(target - observed, value) for value, observed in found if observed > 0]
    positive = [(distance, value) for distance, value in differences if distance >= 0]
    threshold = min(positive)[1] if positive else max(differences)[1]
    return threshold, fpr(human_scores, threshold)


def metrics(rows: list[dict[str, str]], field: str, targets: list[float], epsilon: float) -> dict[str, Any]:
    result = {"auroc": auc([int(row["label"]) for row in rows], [float(row[field]) for row in rows])}
    result["tpr_at_fpr"] = {}
    result["thresholds_by_domain"] = {}
    result["true_fpr_by_domain"] = {}
    for target in targets:
        key = str(target)
        result["thresholds_by_domain"][key] = {}
        result["true_fpr_by_domain"][key] = {}
        correct = total = 0
        for domain in sorted({row["domain"] for row in rows}):
            human = [float(row[field]) for row in rows if row["domain"] == domain and row["label"] == "0"]
            machine = [float(row[field]) for row in rows if row["domain"] == domain and row["label"] == "1"]
            value, observed = threshold_search(human, target, epsilon)
            result["thresholds_by_domain"][key][domain] = value
            result["true_fpr_by_domain"][key][domain] = observed
            correct += sum(score >= value for score in machine)
            total += len(machine)
        result["tpr_at_fpr"][key] = correct / total
    return result


def close(left: float, right: float, tolerance: float = 1e-12) -> bool:
    return math.isclose(left, right, rel_tol=tolerance, abs_tol=tolerance)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", required=True, type=Path)
    parser.add_argument("--config", required=True, type=Path)
    parser.add_argument("--result-dir", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    root = args.repo_root.resolve()
    config = json.loads(args.config.read_text(encoding="utf-8"))
    result_dir = args.result_dir.resolve()
    checks = Checks()

    paths = {
        "dataset": root / config["datasets"]["extra_ood"]["path"],
        "fp32": root / config["models"]["fp32"]["path"],
        "public_quantized": root / config["models"]["public_quantized"]["path"],
        "vocab": root / config["tokenizer"]["vocab_path"],
        "official_evaluator": root / config["metrics"]["official_evaluator_snapshot"],
        "metadata_template": root / config["hidden_test_package"]["metadata_template_snapshot"],
    }
    expected = {
        "dataset": config["datasets"]["extra_ood"]["sha256"],
        "fp32": config["models"]["fp32"]["sha256"],
        "public_quantized": config["models"]["public_quantized"]["sha256"],
        "vocab": config["tokenizer"]["vocab_sha256"],
        "official_evaluator": config["metrics"]["official_evaluator_sha256"],
        "metadata_template": config["hidden_test_package"]["metadata_template_sha256"],
    }
    for name, path in paths.items():
        checks.check(path.is_file(), "input_file", f"{name} exists")
        if path.is_file():
            checks.check(sha256_file(path) == expected[name], "input_checksum", name)
    checks.check(paths["dataset"].stat().st_size == int(config["datasets"]["extra_ood"]["bytes"]), "dataset", "bytes")

    selected, source = reconstruct_sample(paths["dataset"], config)
    checks.check(source["records"] == int(config["datasets"]["extra_ood"]["records"]), "dataset", "record count")
    checks.check(len(selected) == int(config["sample"]["target_records"]), "sample", "target count")
    checks.check(source["human_records"] == 4855, "sample", "all human records retained")
    checks.check({row["domain"] for row in selected} == {"code", "czech", "german"}, "sample", "OOD domain coverage")
    checks.check(len({(row["domain"], row["model"]) for row in selected if row["model"] != "human"}) == 33, "sample", "all machine strata covered")

    manifest = json.loads((result_dir / "sample_manifest.json").read_text(encoding="utf-8"))
    checks.check(manifest["selected"]["ordered_ids_sha256"] == source["ids_sha256"], "sample", "ordered id digest")
    checks.check(manifest["selected"]["machine_allocations"] == source["allocations"], "sample", "stratum allocations")
    checks.check(manifest["selected"]["records"] == len(selected), "sample", "manifest count")

    with (result_dir / "selected_samples.csv").open(newline="", encoding="utf-8") as handle:
        selected_rows = list(csv.DictReader(handle))
    with (result_dir / "sample_predictions.csv").open(newline="", encoding="utf-8") as handle:
        predictions = list(csv.DictReader(handle))
    checks.check(len(selected_rows) == len(selected), "rows", "selected sample rows")
    checks.check(len(predictions) == len(selected), "rows", "prediction rows")
    checks.check(list(predictions[0]) == PREDICTION_FIELDS, "schema", "prediction columns")

    tokenizer = IndependentWordPiece(paths["vocab"])
    for index, (source_row, selected_row, prediction) in enumerate(zip(selected, selected_rows, predictions, strict=True)):
        checks.check(selected_row["id"] == source_row["id"] == prediction["id"], "row_identity", f"id {index}")
        checks.check(int(selected_row["sample_index"]) == int(prediction["sample_index"]) == index, "row_identity", f"index {index}")
        expected_label = 0 if source_row["model"] == "human" else 1
        checks.check(int(prediction["label"]) == expected_label, "label", f"label {index}")
        generation_hash = sha256_text(source_row["generation"])
        checks.check(selected_row["generation_sha256"] == prediction["generation_sha256"] == generation_hash, "text", f"generation {index}")
        token_ids = tokenizer.encode(source_row["generation"])
        token_hash = hashlib.sha256(np.asarray(token_ids, dtype="<i8").tobytes()).hexdigest()
        checks.check(int(prediction["token_count"]) == len(token_ids), "tokenizer", f"count {index}")
        checks.check(prediction["token_ids_sha256"] == token_hash, "tokenizer", f"ids {index}")
        fp_logit = float(prediction["fp32_logit"])
        q_logit = float(prediction["public_quantized_logit"])
        fp_score = 1.0 / (1.0 + math.exp(fp_logit))
        q_score = 1.0 / (1.0 + math.exp(q_logit))
        checks.check(close(float(prediction["fp32_machine_score"]), fp_score), "score", f"fp32 polarity {index}")
        checks.check(close(float(prediction["public_quantized_machine_score"]), q_score), "score", f"quant polarity {index}")

    reference = json.loads((root / "results/lm04_raid_extra_ood/tokenizer_reference_equivalence.json").read_text(encoding="utf-8"))
    checks.check(reference["status"] == "PASS" and reference["all_input_tensors_exact"] is True, "reference_tokenizer", "BertTokenizerFast exact equivalence")
    checks.check(reference["reference"]["transformers_version"] == "4.57.6", "reference_tokenizer", "transformers pin")
    checks.check(reference["reference"]["tokenizers_version"] == "0.22.2", "reference_tokenizer", "tokenizers pin")
    checks.check(reference["case_count"] >= 14, "reference_tokenizer", "multilingual and long cases")
    checks.check(any(row["reaches_max_length"] for row in reference["cases"]), "reference_tokenizer", "right truncation case")

    targets = [float(value) for value in config["metrics"]["target_fpr"]]
    epsilon = float(config["metrics"]["epsilon"])
    computed_fp32 = metrics(predictions, "fp32_machine_score", targets, epsilon)
    computed_quant = metrics(predictions, "public_quantized_machine_score", targets, epsilon)
    summary = json.loads((result_dir / "quality_summary.json").read_text(encoding="utf-8"))
    checks.check(summary["status"] == "THRESHOLD_UNDEFINED", "status", "no invented threshold")
    checks.check(summary["measurement_status"] == "PASS", "status", "metric execution")
    checks.check(summary["result_scope"] == "LOCAL_RAID_EXTRA_OOD_DIAGNOSTIC_NOT_OFFICIAL_HIDDEN_TEST", "scope", "not hidden-test claim")
    for variant, computed in (("fp32", computed_fp32), ("public_quantized", computed_quant)):
        checks.check(close(summary[variant]["auroc"], computed["auroc"]), "metric", f"{variant} AUROC")
        for target in targets:
            key = str(target)
            checks.check(close(summary[variant]["tpr_at_fpr"][key], computed["tpr_at_fpr"][key]), "metric", f"{variant} TPR {key}")
            for domain in ("code", "czech", "german"):
                checks.check(close(summary[variant]["thresholds_by_domain"][key][domain], computed["thresholds_by_domain"][key][domain]), "metric", f"{variant} threshold {key} {domain}")
                checks.check(close(summary[variant]["true_fpr_by_domain"][key][domain], computed["true_fpr_by_domain"][key][domain]), "metric", f"{variant} FPR {key} {domain}")
    checks.check(close(summary["public_quantized_minus_fp32"]["auroc"], computed_quant["auroc"] - computed_fp32["auroc"]), "delta", "AUROC")
    for target in targets:
        key = str(target)
        checks.check(close(summary["public_quantized_minus_fp32"]["tpr_at_fpr"][key], computed_quant["tpr_at_fpr"][key] - computed_fp32["tpr_at_fpr"][key]), "delta", f"TPR {key}")
    policy = summary["policy"]
    for name in ("training_or_fine_tuning", "ptq_qat_or_calibration", "model_weight_or_architecture_modification", "quantized_model_generation", "hidden_test_submission"):
        checks.check(policy[name] is False, "policy", name)
    checks.check(policy["same_sample_and_token_tensors_for_pair"] is True, "policy", "pair token identity")

    result = {
        "schema_version": "1.0",
        "model_id": "LM04",
        "protocol_id": config["protocol_id"],
        "status": "PASS" if not checks.failures else "FAIL",
        "checks_total": checks.total,
        "checks_passed": checks.total - len(checks.failures),
        "checks_failed": len(checks.failures),
        "categories": dict(sorted(checks.categories.items())),
        "failures": checks.failures,
        "independent_recomputed_metrics": {"fp32": computed_fp32, "public_quantized": computed_quant},
        "independent_selected_ids_sha256": source["ids_sha256"],
        "evaluator_module_imported": False,
    }
    atomic_json(args.output, result)
    print(json.dumps({key: result[key] for key in ("status", "checks_total", "checks_passed", "checks_failed")}, sort_keys=True))
    return 0 if result["status"] == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())