"""기존 행동 head를 teacher 형태군과 confidence gate로 결합해 실제 exact 회수율을 감사한다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import sys
from typing import Any, Sequence

import torch
from torch.utils.data import DataLoader, TensorDataset

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.behavior_role_head06 import BehaviorRoleHead06
from math_grid_drawer.research.trajectory_sequence import visual_label_family
from scripts.audit_math_ink_06_case_context import _load_model06
from scripts.crohme_lattice_common import writer_fit_validation
from scripts.train_crohme_segmentation_lattice_selector import _samples
from scripts.train_math_ink_06_behavior_role import _materialize_split06


def _parse_args() -> argparse.Namespace:
    """필요 변수: seed별 adapter/behavior head·CROHME 공식 split. 작동 원리: validation-only gate 감사 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Audit Math Ink 0.6 behavior exact gate")
    parser.add_argument("--adapter", type=Path, action="append", required=True)
    parser.add_argument("--behavior-head", type=Path, action="append", required=True)
    parser.add_argument(
        "--train-root", type=Path,
        default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/trainData",
    )
    parser.add_argument(
        "--test-root", type=Path,
        default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/testDataGT",
    )
    parser.add_argument("--profile", default="median_height_32")
    parser.add_argument("--teacher-batch-size", type=int, default=256)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if len(args.adapter) != len(args.behavior_head):
        raise ValueError("adapter와 behavior head 개수는 같아야 합니다.")
    if not args.adapter:
        raise ValueError("한 개 이상의 seed artifact가 필요합니다.")
    return args


def _role_label06(truth_label: str, role_index: int) -> str | None:
    """필요 변수: target visual base·행동 role. 작동 원리: 같은 형태군 안에서만 exact label로 변환한다."""

    base = "x" if truth_label == r"\times" else truth_label.lower()
    if role_index == 0:
        return base
    if role_index == 1:
        return base.upper()
    if role_index == 2 and base == "x":
        return r"\times"
    return None


def behavior_exact_gate_metrics06(
    behavior_logits: torch.Tensor,
    metadata: Sequence[dict[str, Any]],
    *,
    threshold: float,
) -> dict[str, float | int]:
    """필요 변수: role logit·teacher/truth metadata·threshold. 작동 원리: family 확인 뒤 rewrite하고 나머지는 abstain한다."""

    if len(behavior_logits) != len(metadata):
        raise ValueError("behavior logit과 metadata 길이가 다릅니다.")
    probability = behavior_logits.softmax(dim=1)
    confidence, roles = probability.max(dim=1)
    teacher_correct = final_correct = family_correct = rewrites = beneficial = harmful = 0
    abstained = 0
    for index, row in enumerate(metadata):
        truth = str(row["truth_label"])
        teacher = str(row["teacher_label"])
        teacher_hit = teacher == truth
        family_hit = visual_label_family(teacher) == visual_label_family(truth)
        selected = _role_label06(truth, int(roles[index]))
        rewrite = (
            family_hit
            and selected is not None
            and float(confidence[index]) >= threshold
        )
        final = selected if rewrite else teacher
        teacher_correct += int(teacher_hit)
        family_correct += int(family_hit)
        final_correct += int(final == truth)
        rewrites += int(rewrite)
        abstained += int(not rewrite)
        beneficial += int(rewrite and not teacher_hit and final == truth)
        harmful += int(rewrite and teacher_hit and final != truth)
    samples = len(metadata)
    return {
        "threshold": float(threshold),
        "samples": samples,
        "teacher_exact": teacher_correct / max(samples, 1),
        "teacher_visual_family": family_correct / max(samples, 1),
        "final_exact": final_correct / max(samples, 1),
        "gain_pp": (final_correct - teacher_correct) * 100.0 / max(samples, 1),
        "rewrites": rewrites,
        "abstained": abstained,
        "beneficial": beneficial,
        "harmful": harmful,
        "rewrite_precision": beneficial / max(beneficial + harmful, 1),
    }


def _behavior_logits06(
    checkpoint: Path,
    dataset: TensorDataset,
    *,
    device: torch.device,
    batch_size: int,
) -> torch.Tensor:
    """필요 변수: behavior checkpoint·raw dataset. 작동 원리: checkpoint fit 통계로 context를 정규화해 role logit을 반환한다."""

    payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
    model = BehaviorRoleHead06(
        sequence_channels=int(payload.get("sequence_channels", 19)),
        context_features=len(payload["context_features"]),
        hidden=int(payload["hidden"]),
        dropout=float(payload["dropout"]),
    ).to(device)
    model.load_state_dict(payload["state_dict"])
    model.eval()
    mean = payload["context_mean"].float()
    scale = payload["context_scale"].float().clamp_min(1e-5)
    rows = []
    with torch.inference_mode():
        for sequence, context, _target in DataLoader(
            dataset, batch_size=batch_size, shuffle=False,
        ):
            normalized = (context - mean) / scale
            rows.append(model(sequence.to(device), normalized.to(device)).cpu())
    return torch.cat(rows)


def _select_threshold06(
    logits: torch.Tensor,
    metadata: Sequence[dict[str, Any]],
) -> tuple[float, list[dict[str, float | int]]]:
    """필요 변수: validation role logit·metadata. 작동 원리: exact 우선·harm 최소·높은 threshold 순으로 gate를 고정한다."""

    thresholds = tuple(index / 100.0 for index in range(0, 100, 2))
    sweep = [
        behavior_exact_gate_metrics06(logits, metadata, threshold=value)
        for value in thresholds
    ]
    selected = max(
        sweep,
        key=lambda row: (
            float(row["final_exact"]),
            -int(row["harmful"]),
            float(row["threshold"]),
        ),
    )
    return float(selected["threshold"]), sweep


def main() -> None:
    """필요 변수: 공식 train writer-validation과 held-out test. 작동 원리: seed별 threshold를 validation에서 잠그고 test에 한 번 적용한다."""

    args = _parse_args()
    if args.train_root.name.casefold() != "traindata" or args.test_root.name.casefold() != "testdatagt":
        raise ValueError("CROHME2012 공식 trainData/testDataGT 조합만 허용합니다.")
    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA 감사를 요청했지만 사용할 수 없습니다.")
    _fit, validation_samples = writer_fit_validation(args.train_root, args.profile)
    test_samples = _samples(args.test_root, args.profile)
    seed_rows = []
    for adapter_path, head_path in zip(args.adapter, args.behavior_head, strict=True):
        adapter_payload = torch.load(adapter_path, map_location="cpu", weights_only=False)
        base_checkpoint = Path(str(adapter_payload["base_checkpoint"]))
        if not base_checkpoint.is_absolute():
            base_checkpoint = PROJECT_ROOT / base_checkpoint
        engine, adapter = _load_model06(base_checkpoint, adapter_path, device)
        validation, _validation_counts, validation_metadata = _materialize_split06(
            validation_samples, engine, adapter, device=device,
            teacher_batch_size=args.teacher_batch_size, return_metadata=True,
        )
        testing, _test_counts, test_metadata = _materialize_split06(
            test_samples, engine, adapter, device=device,
            teacher_batch_size=args.teacher_batch_size, return_metadata=True,
        )
        validation_logits = _behavior_logits06(
            head_path, validation, device=device, batch_size=args.batch_size,
        )
        selected_threshold, sweep = _select_threshold06(
            validation_logits, validation_metadata,
        )
        test_logits = _behavior_logits06(
            head_path, testing, device=device, batch_size=args.batch_size,
        )
        seed_rows.append({
            "adapter": str(adapter_path),
            "behavior_head": str(head_path),
            "selected_threshold": selected_threshold,
            "validation_selected": behavior_exact_gate_metrics06(
                validation_logits, validation_metadata, threshold=selected_threshold,
            ),
            "validation_sweep": sweep,
            "official_test": behavior_exact_gate_metrics06(
                test_logits, test_metadata, threshold=selected_threshold,
            ),
        })
        del engine, adapter
        if device.type == "cuda":
            torch.cuda.empty_cache()
    metric_names = (
        "teacher_exact", "teacher_visual_family", "final_exact", "gain_pp",
        "rewrite_precision",
    )
    summary = {
        name: {
            "values": [float(row["official_test"][name]) for row in seed_rows],
            "mean": sum(float(row["official_test"][name]) for row in seed_rows) / len(seed_rows),
        }
        for name in metric_names
    }
    report = {
        "experiment": "R-MATH-INK-06-BEHAVIOR-EXACT-GATE-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "split_contract": "CROHME trainData writer-validation threshold; testDataGT one-shot",
        "scope": "truth symbol grouping conditional; c/C, x/X/times, z/Z only",
        "seeds": seed_rows,
        "official_test_summary": summary,
        "track": "R_noncommercial_only",
        "product_validation": False,
        "distillation_allowed": False,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({
        "official_test_summary": summary,
        "thresholds": [row["selected_threshold"] for row in seed_rows],
        "product_validation": False,
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()