package_code_imports/v4p4_scripts/audit_v4p4_finalization.py

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from pathlib import Path
from typing import Any


SCRIPT_DIR = Path(__file__).resolve().parent
V3P5_SCRIPT_DIR = SCRIPT_DIR.parents[1] / "v3p5_static" / "scripts"
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))
if str(V3P5_SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(V3P5_SCRIPT_DIR))

from smoke_v3p4_architecture import load_json
from train_v4p4_cloud import DEFAULT_STAGE0_DIR, DEFAULT_TENSOR_DIR, build_color_map


AUTHORITATIVE_SCRIPTS = [
    "model_v4p4.py",
    "loss_v4p4.py",
    "config_v4p4_train.py",
    "train_v4p4_cloud.py",
    "hf_submit_v4p4_world_model_train_job.py",
    "hf_submit_v4p4_eval_export_job.py",
    "export_v4p4_visit_outputs.py",
    "smoke_v4p4_architecture.py",
    "rollout_v4p4_closed_loop.py",
    "evaluate_v4p4_independent.py",
    "evaluate_v4p4_closed_loop_rollout.py",
    "evaluate_v4p4_final_metric_tables.py",
    "evaluate_v4p4_world_model.py",
    "audit_v4p4_metric_coverage.py",
    "audit_v4p4_finalization.py",
]


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


def json_ready(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {str(k): json_ready(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_ready(v) for v in value]
    return value


def read_text(name: str) -> str:
    return (SCRIPT_DIR / name).read_text(encoding="utf-8")


def read_optional_text(name: str) -> str:
    path = SCRIPT_DIR / name
    return path.read_text(encoding="utf-8") if path.exists() else ""


def read_compat_script_text(name: str, legacy_dir: str) -> str:
    candidates = [
        SCRIPT_DIR / name,
        SCRIPT_DIR.parents[1] / legacy_dir / "scripts" / name,
    ]
    for path in candidates:
        if path.exists():
            return path.read_text(encoding="utf-8")
    raise FileNotFoundError(
        f"Could not find {name}; checked: " + ", ".join(str(path) for path in candidates)
    )


def audit_code_version() -> dict[str, Any]:
    rows = []
    for name in AUTHORITATIVE_SCRIPTS:
        path = SCRIPT_DIR / name
        rows.append(
            {
                "script": name,
                "exists": path.exists(),
                "sha256": sha256_file(path) if path.exists() else "",
                "bytes": path.stat().st_size if path.exists() else 0,
            }
        )
    py_files = sorted(p.name for p in SCRIPT_DIR.glob("*.py"))
    return {
        "status": "pass" if all(row["exists"] for row in rows) else "fail",
        "authoritative_scripts": rows,
        "python_files_in_script_dir": py_files,
        "legacy_policy": "Use v4p4 scripts for v4p4 training/evaluation; v4p3/v3p5 imports are compatibility helpers only.",
    }


def audit_export_sources() -> dict[str, Any]:
    src = read_text("export_v4p4_visit_outputs.py")
    checks = {
        "pre_visit_uses_history_deployable": 'out["z_history"], torch.sigmoid(out["deployable_risk_logits_pre"]), out["pwe_log_lambda_pre"]' in src,
        "contact_uses_contact_deployable": 'out["z_contact"], torch.sigmoid(out["deployable_risk_logits_contact"]), out["pwe_log_lambda_contact"]' in src,
        "field_end_uses_post_deployable": 'out["z_post"], torch.sigmoid(out["deployable_risk_logits_post"]), out["pwe_log_lambda_post"]' in src,
        "export_disables_reconstruction": "compute_reconstruction=False" in src,
        "export_disables_teacher_forcing": "compute_teacher_forced=False" in src,
        "export_does_not_read_event_logits": 'out["event_logits"]' not in src and "out['event_logits']" not in src,
    }
    return {
        "status": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "expected_mode_sources": {
            "pre_visit": {"embedding": "z_history", "risk": "deployable_risk_logits_pre", "pwe": "pwe_log_lambda_pre"},
            "contact": {"embedding": "z_contact", "risk": "deployable_risk_logits_contact", "pwe": "pwe_log_lambda_contact"},
            "field_end": {"embedding": "z_post", "risk": "deployable_risk_logits_post", "pwe": "pwe_log_lambda_post"},
        },
    }


def audit_model_factorization() -> dict[str, Any]:
    src = read_text("model_v4p4.py")
    checks = {
        "active_state_head_uses_time_context": 'out["active_state_logits"] = self.active_state_head(time_ctx)' in src,
        "downstream_heads_condition_on_active_context": "obs_ctx = self._condition_on_next_active(time_ctx, active_ids)" in src,
        "rollout_downstream_conditions_on_sampled_active": "active_ids = torch.where(next_contact, active_state" in src
        and "obs_ctx = self._condition_on_next_active(time_ctx, active_ids)" in src,
        "event_generation_head_dedicated": "self.event_generation_head" in src and "generation_event_logits = self.event_generation_head(obs_ctx)" in src,
        "year_uses_absolute_elapsed_time": "start_year + torch.floor(next_time / 365.25).long()" in src,
        "event_logits_removed": '"event_logits"' not in re.sub(r"event_logits_absent", "", src),
    }
    return {
        "status": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "factorization": "p(K,dt|z_post) -> p(S_next|z_post,dt) -> p(M_next,X_next,Y_next|z_post,dt,S_next)",
    }


def audit_pwe_math() -> dict[str, Any]:
    src = read_compat_script_text("loss_v4p3.py", "v4p3_l1l2")
    v4p4_src = read_text("loss_v4p4.py")
    checks = {
        "exact_time_density_log_survival_plus_log_lambda": "event_loglik = log_survival_total + torch.log(lambda_k_star.clamp(min=1.0e-8))" in src,
        "admin_censor_survival_only": "torch.where(is_censored, log_survival_total, event_loglik)" in src,
        "v4p4_scores_pre_contact_post_pwe": all(token in v4p4_src for token in ["pwe_log_lambda_pre", "pwe_log_lambda_contact", "pwe_log_lambda_post"]),
        "no_interval_bin_probability_term_in_nll": "log1mexp_neg(" not in src.split("def pwe_target_from_terminal", 1)[0].split("def log1mexp_neg", 1)[-1],
    }
    return {
        "status": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "math": "Exact-time piecewise exponential competing-risk density: log S(t) + log lambda_k(t); admin censoring uses log S(t).",
    }


def audit_color_order(tensor_dir: Path) -> dict[str, Any]:
    meta = load_json(tensor_dir / "tensor_metadata.json")
    vocab = load_json(tensor_dir / "cat_value_vocab.json")
    color_field_index, pairs = build_color_map(meta, vocab)
    id_to_label = {int(v): k for k, v in vocab.items()}
    rows = []
    for value_id, cls in pairs:
        token = id_to_label.get(int(value_id), "")
        rows.append({"value_id": int(value_id), "token": token, "ordinal_class": int(cls)})
    rows_sorted = sorted(rows, key=lambda x: (x["ordinal_class"], x["value_id"]))
    observed_order = [row["token"].split("=", 1)[-1] for row in rows_sorted]
    checks = {
        "color_field_present": color_field_index >= 0,
        "five_color_values_found": len(rows_sorted) == 5,
        "green_blue_yellow_orange_red": observed_order == ["绿色", "蓝色", "黄色", "橙色", "红色"] or observed_order == ["绿", "蓝", "黄", "橙", "红"],
    }
    return {
        "status": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "color_field_index": int(color_field_index),
        "official_order_implemented": "green < blue < yellow < orange < red",
        "value_to_class": rows_sorted,
        "evidence": "train_v4p4_cloud.build_color_map plus tensor cat_value_vocab.json",
    }


def audit_training_defaults() -> dict[str, Any]:
    train_src = read_text("train_v4p4_cloud.py")
    submit_src = read_optional_text("hf_submit_v4p4_world_model_train_job.py")
    checks = {
        "train_default_stage2a": 'parser.add_argument("--stage", choices=["stage1", "stage2a"], default="stage2a")' in train_src,
        "submit_script_present": bool(submit_src),
        "submit_default_stage2a": 'parser.add_argument("--stage", choices=["stage1", "stage2a"], default="stage2a")' in submit_src,
        "metrics_are_sparse": "return_metrics=need_step_metrics" in train_src,
        "loss_can_skip_metrics": "return_metrics: bool = False" in read_text("loss_v4p4.py"),
        "permutation_sampler_present": "class TrainIndexSampler" in train_src and "self.rng.permutation" in train_src,
        "formal_eval_frequency": '--eval-every", type=int, default=500' in train_src and '--eval-every", type=int, default=500' in submit_src,
        "no_replacement_submit_default": "parser.set_defaults(sample_with_replacement=False)" in submit_src,
    }
    return {
        "status": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "recommended_cloud_run": "stage2a/joint_one_step, bf16, fused AdamW, no replacement sampler, log_every=25, eval_every=500, save_every=1000.",
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Write SCTM-v4.4 finalization audit files.")
    parser.add_argument("--tensor-dir", type=Path, default=DEFAULT_TENSOR_DIR)
    parser.add_argument("--stage0-dir", type=Path, default=DEFAULT_STAGE0_DIR)
    parser.add_argument("--out-dir", type=Path, required=True)
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    audits = {
        "code_version_audit": audit_code_version(),
        "export_source_audit": audit_export_sources(),
        "model_factorization_audit": audit_model_factorization(),
        "pwe_math_audit": audit_pwe_math(),
        "color_signal_order_audit": audit_color_order(args.tensor_dir),
        "training_engineering_audit": audit_training_defaults(),
    }
    for name, payload in audits.items():
        (args.out_dir / f"{name}.json").write_text(json.dumps(json_ready(payload), ensure_ascii=False, indent=2), encoding="utf-8")

    summary = {
        "status": "pass" if all(payload.get("status") == "pass" for payload in audits.values()) else "fail",
        "tensor_dir": str(args.tensor_dir),
        "stage0_dir": str(args.stage0_dir),
        "audit_files": {name: str(args.out_dir / f"{name}.json") for name in audits},
    }
    (args.out_dir / "finalization_audit_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()