from __future__ import annotations

import csv
import hashlib
import json
import math
import os
import resource
import time
import unicodedata
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import skops.io as sio
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
    roc_auc_score,
)
from sklearn.pipeline import FeatureUnion

CAMPAIGN_ID = "banking77-intent-error-predictor-v1"
SOURCE_REVISION = "57ec275d8078af65b7731c2a98be812d844a6d6b"
HUB_REVISION = "90d4e2ee5521c04fc1488f065b8b083658768c57"
SOURCE_ROOT = (
    "https://raw.githubusercontent.com/PolyAI-LDN/task-specific-datasets/"
    f"{SOURCE_REVISION}/banking_data"
)
EXPECTED_SHA256 = {
    "train.csv": "b06e26ac675513959a63135f11b94ea7786ed02da65db93a5650d8838cbc664b",
    "test.csv": "d12d6e3bc4c3103966ae786dc435913c0c563dfa328f5a3646d0e62cfeeb474d",
    "categories.json": (
        "53261da888122daf2d120d925458631d9619e15d82e56052e7a42e535ce32b63"
    ),
}
SEED = 20260811
THREAD_LIMIT = 8
REVIEW_RATE = 0.20


@dataclass(frozen=True)
class Row:
    text: str
    label: str


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


def normalize(text: str) -> str:
    return " ".join(unicodedata.normalize("NFKC", text).casefold().split())


def group_bucket(text: str) -> int:
    return int(hashlib.sha256(normalize(text).encode()).hexdigest()[:8], 16) % 100


def download(cache_dir: Path, filename: str) -> Path:
    path = cache_dir / "banking77" / SOURCE_REVISION / filename
    if not path.exists():
        path.parent.mkdir(parents=True, exist_ok=True)
        request = urllib.request.Request(
            f"{SOURCE_ROOT}/{filename}",
            headers={"User-Agent": "ITheEqualizer-banking77-audit/0.1"},
        )
        temporary = path.with_suffix(path.suffix + ".partial")
        with urllib.request.urlopen(request, timeout=60) as response:
            temporary.write_bytes(response.read())
        temporary.replace(path)
    actual = sha256_file(path)
    expected = EXPECTED_SHA256[filename]
    if actual != expected:
        raise RuntimeError(f"checksum mismatch for {filename}: {actual}")
    return path


def load_rows(path: Path) -> list[Row]:
    with path.open(newline="", encoding="utf-8") as handle:
        rows = [
            Row(text=row["text"], label=row["category"])
            for row in csv.DictReader(handle)
        ]
    if not rows or any(not row.text.strip() or not row.label.strip() for row in rows):
        raise RuntimeError(f"invalid or empty rows in {path.name}")
    return rows


def build_primary() -> tuple[FeatureUnion, LogisticRegression]:
    features = FeatureUnion(
        [
            (
                "word",
                TfidfVectorizer(
                    ngram_range=(1, 2),
                    min_df=2,
                    max_features=24000,
                    sublinear_tf=True,
                    strip_accents="unicode",
                ),
            ),
            (
                "char",
                TfidfVectorizer(
                    analyzer="char_wb",
                    ngram_range=(3, 5),
                    min_df=2,
                    max_features=36000,
                    sublinear_tf=True,
                    strip_accents="unicode",
                ),
            ),
        ]
    )
    model = LogisticRegression(
        C=4.0,
        max_iter=600,
        solver="lbfgs",
        random_state=SEED,
    )
    return features, model


def risk_features(probabilities: np.ndarray, texts: list[str]) -> np.ndarray:
    clipped = np.clip(probabilities, 1e-12, 1.0)
    ordered = np.sort(clipped, axis=1)
    top1 = ordered[:, -1]
    top2 = ordered[:, -2]
    entropy = -(clipped * np.log(clipped)).sum(axis=1) / math.log(clipped.shape[1])
    shapes = np.asarray(
        [
            [
                min(len(text), 512) / 512,
                min(len(text.split()), 100) / 100,
                min(sum(ch.isdigit() for ch in text), 20) / 20,
                float("?" in text),
                float(
                    any(
                        token in normalize(text).split()
                        for token in ("not", "no", "never", "wrong")
                    )
                ),
            ]
            for text in texts
        ],
        dtype=np.float64,
    )
    predicted_one_hot = np.zeros_like(clipped)
    predicted_one_hot[np.arange(len(clipped)), np.argmax(clipped, axis=1)] = 1.0
    return np.column_stack(
        [clipped, predicted_one_hot, top1, top2, top1 - top2, entropy, shapes]
    )


def error_metrics(
    errors: np.ndarray,
    risks: np.ndarray,
    *,
    threshold: float,
) -> dict[str, Any]:
    reviewed = risks >= threshold
    tn, fp, fn, tp = confusion_matrix(errors, reviewed, labels=[0, 1]).ravel()
    precision, recall, f1, _ = precision_recall_fscore_support(
        errors, reviewed, average="binary", zero_division=0
    )
    routed = ~reviewed
    return {
        "error_prevalence": float(errors.mean()),
        "roc_auc": float(roc_auc_score(errors, risks)),
        "pr_auc": float(average_precision_score(errors, risks)),
        "threshold": float(threshold),
        "review_rate": float(reviewed.mean()),
        "error_precision": float(precision),
        "error_recall": float(recall),
        "error_f1": float(f1),
        "confusion": {"tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp)},
        "routed_accuracy": float(1.0 - errors[routed].mean()) if routed.any() else 0.0,
        "coverage": float(routed.mean()),
    }


def review_threshold(risks: np.ndarray) -> float:
    return float(np.quantile(risks, 1.0 - REVIEW_RATE, method="higher"))


def dump_checked(model: Any, path: Path) -> dict[str, Any]:
    sio.dump(model, path)
    untrusted = sio.get_untrusted_types(file=path)
    if untrusted:
        raise RuntimeError(f"non-empty skops untrusted type set: {untrusted}")
    return {
        "path": path.name,
        "bytes": path.stat().st_size,
        "sha256": sha256_file(path),
        "skops_untrusted_types": untrusted,
    }


def fit_once(train_rows: list[Row], test_rows: list[Row]) -> dict[str, Any]:
    primary_train = [row for row in train_rows if group_bucket(row.text) < 60]
    complement_train = [row for row in train_rows if 60 <= group_bucket(row.text) < 80]
    validation = [row for row in train_rows if group_bucket(row.text) >= 80]
    primary_groups = {normalize(row.text) for row in primary_train}
    complement_groups = {normalize(row.text) for row in complement_train}
    validation_groups = {normalize(row.text) for row in validation}
    if (
        primary_groups & complement_groups
        or primary_groups & validation_groups
        or complement_groups & validation_groups
    ):
        raise RuntimeError("group leakage across development partitions")
    lockbox = [
        row
        for row in test_rows
        if normalize(row.text)
        not in primary_groups | complement_groups | validation_groups
    ]
    labels = sorted({row.label for row in train_rows})
    if len(labels) != 77 or any(
        {row.label for row in part} != set(labels)
        for part in (primary_train, complement_train, validation, lockbox)
    ):
        raise RuntimeError("all four partitions must contain all 77 labels")

    features, primary = build_primary()
    x_primary = features.fit_transform([row.text for row in primary_train])
    primary.fit(x_primary, [row.label for row in primary_train])

    def primary_outputs(rows: list[Row]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        probs = primary.predict_proba(features.transform([row.text for row in rows]))
        predictions = primary.classes_[np.argmax(probs, axis=1)]
        truth = np.asarray([row.label for row in rows])
        return probs, predictions, (predictions != truth).astype(np.int64)

    complement_probs, _, complement_errors = primary_outputs(complement_train)
    validation_probs, validation_predictions, validation_errors = primary_outputs(
        validation
    )
    lockbox_probs, lockbox_predictions, lockbox_errors = primary_outputs(lockbox)

    candidate = HistGradientBoostingClassifier(
        learning_rate=0.06,
        max_iter=140,
        max_leaf_nodes=15,
        min_samples_leaf=25,
        l2_regularization=1.0,
        class_weight="balanced",
        random_state=SEED,
    )
    candidate.fit(
        risk_features(complement_probs, [row.text for row in complement_train]),
        complement_errors,
    )
    validation_risk = candidate.predict_proba(
        risk_features(validation_probs, [row.text for row in validation])
    )[:, 1]
    lockbox_risk = candidate.predict_proba(
        risk_features(lockbox_probs, [row.text for row in lockbox])
    )[:, 1]
    validation_margin_risk = (
        1.0
        - np.partition(validation_probs, -2, axis=1)[:, -1]
        + np.partition(validation_probs, -2, axis=1)[:, -2]
    )
    lockbox_margin_risk = (
        1.0
        - np.partition(lockbox_probs, -2, axis=1)[:, -1]
        + np.partition(lockbox_probs, -2, axis=1)[:, -2]
    )
    candidate_threshold = review_threshold(validation_risk)
    margin_threshold = review_threshold(validation_margin_risk)

    return {
        "models": {
            "feature_extractor": features,
            "primary": primary,
            "candidate": candidate,
        },
        "partitions": {
            "primary_train": len(primary_train),
            "complement_train": len(complement_train),
            "validation": len(validation),
            "lockbox": len(lockbox),
            "lockbox_overlap_removed": len(test_rows) - len(lockbox),
        },
        "labels": labels,
        "primary": {
            "validation_accuracy": float(
                accuracy_score(
                    [row.label for row in validation], validation_predictions
                )
            ),
            "validation_macro_f1": float(
                f1_score(
                    [row.label for row in validation],
                    validation_predictions,
                    average="macro",
                )
            ),
            "lockbox_accuracy": float(
                accuracy_score([row.label for row in lockbox], lockbox_predictions)
            ),
            "lockbox_macro_f1": float(
                f1_score(
                    [row.label for row in lockbox], lockbox_predictions, average="macro"
                )
            ),
        },
        "validation": {
            "candidate": error_metrics(
                validation_errors, validation_risk, threshold=candidate_threshold
            ),
            "margin_baseline": error_metrics(
                validation_errors, validation_margin_risk, threshold=margin_threshold
            ),
        },
        "lockbox": {
            "candidate": error_metrics(
                lockbox_errors, lockbox_risk, threshold=candidate_threshold
            ),
            "margin_baseline": error_metrics(
                lockbox_errors, lockbox_margin_risk, threshold=margin_threshold
            ),
        },
        "reproduction_reference": {
            "validation_candidate_risk": validation_risk,
            "validation_candidate_predictions": validation_risk >= candidate_threshold,
            "lockbox_candidate_risk": lockbox_risk,
            "lockbox_candidate_predictions": lockbox_risk >= candidate_threshold,
        },
    }


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--cache-dir", type=Path, required=True)
    parser.add_argument("--ledger", type=Path, required=True)
    parser.add_argument("--state", type=Path, required=True)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=False)
    cache_before = (
        sum(path.stat().st_size for path in args.cache_dir.rglob("*") if path.is_file())
        if args.cache_dir.exists()
        else 0
    )
    for name in (
        "OMP_NUM_THREADS",
        "OPENBLAS_NUM_THREADS",
        "MKL_NUM_THREADS",
        "VECLIB_MAXIMUM_THREADS",
        "NUMEXPR_NUM_THREADS",
    ):
        os.environ[name] = str(THREAD_LIMIT)
    train_path = download(args.cache_dir, "train.csv")
    test_path = download(args.cache_dir, "test.csv")
    categories_path = download(args.cache_dir, "categories.json")
    categories = json.loads(categories_path.read_text(encoding="utf-8"))
    train_rows = load_rows(train_path)
    test_rows = load_rows(test_path)
    if sorted(categories) != sorted({row.label for row in train_rows}):
        raise RuntimeError("category manifest mismatch")

    code_hash = sha256_file(Path(__file__))
    specification = {
        "campaign_id": CAMPAIGN_ID,
        "candidate": "histogram_gradient_boosting_score_error_predictor",
        "consumer": "banking-support intent router with a human review queue",
        "task": "predict whether a fixed BANKING77 primary router is wrong",
        "input_contract": (
            "77 probabilities in sorted Banking77 label order plus bounded "
            "text-shape features"
        ),
        "output_contract": "error probability and advisory review decision",
        "dataset_hub_revision": HUB_REVISION,
        "dataset_source_revision": SOURCE_REVISION,
        "dataset_hashes": EXPECTED_SHA256,
        "preprocessing": (
            "NFKC casefold whitespace normalization for group hashing; TF-IDF "
            "primary; score and text-shape candidate features"
        ),
        "split": (
            "normalized-text SHA-256 grouped 60/20/20 development partitions; "
            "official test untouched lockbox with overlap removal"
        ),
        "primary": (
            "24k word plus 36k character TF-IDF with multinomial logistic "
            "regression C=4"
        ),
        "architecture": (
            "histogram gradient boosting over probabilities, predicted-intent "
            "one-hot, confidence geometry, and five text-shape features"
        ),
        "objective": "balanced binary log loss for primary-router error prediction",
        "hyperparameters": {
            "learning_rate": 0.06,
            "max_iter": 140,
            "max_leaf_nodes": 15,
            "min_samples_leaf": 25,
            "l2_regularization": 1.0,
        },
        "seed": SEED,
        "code_sha256": code_hash,
    }
    spec_hash = hashlib.sha256(
        json.dumps(specification, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    if any(
        json.loads(line).get("experiment_spec_hash") == spec_hash
        for line in args.ledger.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ):
        raise RuntimeError("experiment specification already exists in ledger")

    wall_start = time.perf_counter()
    cpu_start = time.process_time()
    usage_start = resource.getrusage(resource.RUSAGE_SELF)
    result = fit_once(train_rows, test_rows)
    telemetry = {
        "wall_seconds": time.perf_counter() - wall_start,
        "cpu_seconds": time.process_time() - cpu_start,
        "peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
        "minor_page_fault_delta": resource.getrusage(resource.RUSAGE_SELF).ru_minflt
        - usage_start.ru_minflt,
        "thread_limit": THREAD_LIMIT,
        "single_training_process": True,
    }
    primary_bundle = {
        "features": result["models"]["feature_extractor"],
        "classifier": result["models"]["primary"],
        "labels": result["labels"],
    }
    primary_artifact = dump_checked(
        primary_bundle, args.output / "primary_baseline.skops"
    )
    candidate_bundle = {
        "classifier": result["models"]["candidate"],
        "labels": result["labels"],
        "review_rate": REVIEW_RATE,
    }
    candidate_artifact = dump_checked(candidate_bundle, args.output / "model.skops")

    rerun = fit_once(train_rows, test_rows)
    exact_scores = all(
        np.array_equal(
            result["reproduction_reference"][key], rerun["reproduction_reference"][key]
        )
        for key in ("validation_candidate_risk", "lockbox_candidate_risk")
    )
    exact_predictions = all(
        np.array_equal(
            result["reproduction_reference"][key], rerun["reproduction_reference"][key]
        )
        for key in ("validation_candidate_predictions", "lockbox_candidate_predictions")
    )
    loaded = sio.load(args.output / "model.skops", trusted=[])
    serialization_scores_exact = np.array_equal(
        loaded["classifier"].predict_proba(
            risk_features(
                rerun["models"]["primary"].predict_proba(
                    rerun["models"]["feature_extractor"].transform(
                        [
                            row.text
                            for row in [
                                r for r in train_rows if group_bucket(r.text) >= 80
                            ]
                        ]
                    )
                ),
                [
                    row.text
                    for row in [r for r in train_rows if group_bucket(r.text) >= 80]
                ],
            )
        )[:, 1],
        rerun["reproduction_reference"]["validation_candidate_risk"],
    )
    reproduction = {
        "clean_refit_scores_exact": exact_scores,
        "clean_refit_predictions_exact": exact_predictions,
        "serialization_scores_exact": serialization_scores_exact,
    }
    validation_candidate = result["validation"]["candidate"]
    validation_margin = result["validation"]["margin_baseline"]
    lockbox_candidate = result["lockbox"]["candidate"]
    acceptance = {
        "validation_error_recall_margin_delta_at_least_0.02": validation_candidate[
            "error_recall"
        ]
        >= validation_margin["error_recall"] + 0.02,
        "validation_routed_accuracy_gain_at_least_0.03": validation_candidate[
            "routed_accuracy"
        ]
        >= result["primary"]["validation_accuracy"] + 0.03,
        "lockbox_error_recall_at_least_0.50": lockbox_candidate["error_recall"] >= 0.50,
        "artifact_at_most_2_mb": candidate_artifact["bytes"] <= 2_000_000,
        "empty_skops_untrusted_type_set": not candidate_artifact[
            "skops_untrusted_types"
        ],
        "exact_reproduction": all(reproduction.values()),
    }
    decision = (
        "release_work_pending" if all(acceptance.values()) else "measured_gates_failed"
    )
    cache_after = sum(
        path.stat().st_size for path in args.cache_dir.rglob("*") if path.is_file()
    )
    metrics = {
        "record_type": "fit",
        "campaign_id": CAMPAIGN_ID,
        "candidate": specification["candidate"],
        "hypothesis": (
            "learned score-shape and intent features improve error capture over "
            "margin-only triage at equal review rate"
        ),
        "experiment_spec_hash": spec_hash,
        "code_sha256": code_hash,
        "dataset_revision": SOURCE_REVISION,
        "base_model_revision": None,
        "configuration": specification,
        "partitions": result["partitions"],
        "primary": result["primary"],
        "validation": result["validation"],
        "lockbox": result["lockbox"],
        "artifacts": {
            "primary_baseline": primary_artifact,
            "candidate": candidate_artifact,
        },
        "acceptance": acceptance,
        "acceptance_decision": decision,
        "reproduction": reproduction,
        "telemetry": {**telemetry, "cache_growth_bytes": cache_after - cache_before},
    }
    (args.output / "metrics.json").write_text(
        json.dumps(metrics, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    temporary_disk_growth = sum(
        path.stat().st_size for path in args.output.rglob("*") if path.is_file()
    )
    metrics["telemetry"]["temporary_disk_growth_bytes"] = temporary_disk_growth
    (args.output / "metrics.json").write_text(
        json.dumps(metrics, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    ledger_record = {**metrics, "telemetry": metrics["telemetry"]}
    with args.ledger.open("a", encoding="utf-8") as handle:
        handle.write(
            json.dumps(ledger_record, sort_keys=True, separators=(",", ":")) + "\n"
        )
        handle.flush()
        os.fsync(handle.fileno())
    state = {
        "campaign_id": CAMPAIGN_ID,
        "created_at": "2026-08-11T13:25:00Z",
        "status": decision,
        "consumer": specification["consumer"],
        "task": specification["task"],
        "differentiator": (
            "sub-2 MB model-specific selective-routing complement rather than "
            "another primary intent classifier"
        ),
        "trend_snapshot": "campaign/trends/20260811T1300Z-v2.json",
        "preflight": "campaign/trends/20260811T1325Z-banking-error-preflight.json",
        "current_experiment": {
            "candidate": specification["candidate"],
            "experiment_spec_hash": spec_hash,
            "hypothesis": ledger_record["hypothesis"],
        },
        "last_run": metrics,
        "release_blocker": None
        if decision == "release_work_pending"
        else "learned error risk did not clear every predeclared usefulness gate",
        "next_action": (
            "run one error-driven candidate using class-conditional calibration "
            "and confusion-neighborhood features, selected on validation only; "
            "abandon if it still fails to beat margin triage"
        ),
    }
    temporary_state = args.state.with_suffix(".json.tmp")
    temporary_state.write_text(
        json.dumps(state, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    temporary_state.replace(args.state)
    print(
        json.dumps(
            {
                "campaign_id": CAMPAIGN_ID,
                "decision": decision,
                "spec_hash": spec_hash,
                "candidate_artifact": candidate_artifact,
                "validation": result["validation"],
                "lockbox": result["lockbox"],
                "reproduction": reproduction,
                "telemetry": metrics["telemetry"],
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()