"""Verify the full released PV-PRISM against registered native predictions."""
import argparse
import gc
import json
import os
from pathlib import Path
import shutil
import site
import sys
import tempfile
import time

ROOT = Path(__file__).resolve().parent
sys.dont_write_bytecode = True


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--forbid-original", type=Path, required=True)
    args = parser.parse_args()
    target = ROOT / "validation/result.json"
    if target.exists():
        parser.error("Existing validation receipt must not be overwritten")
    old = args.forbid_original.resolve()
    # Installed third-party libraries are runtime dependencies, not experiment assets.
    libraries = [Path(p).resolve() for p in site.getsitepackages()]
    blocked = []
    def audit(event, values):
        if event == "open" and isinstance(values[0], (str, bytes, os.PathLike)):
            path = Path(os.fsdecode(values[0])).resolve()
            if path.is_relative_to(old) and not path.is_relative_to(ROOT) and not any(path.is_relative_to(p) for p in libraries):
                blocked.append(str(path))
                raise PermissionError("Original workspace read blocked: " + str(path))
    sys.addaudithook(audit)
    from runtime import PVPRISM, read
    import numpy as np
    import torch
    config = read(ROOT / "configs/model.json")
    result = dict(status="running", model="PV-PRISM", device=args.device, cases=[],
                  all_original_project_reads_denied=True, query_labels_shipped=False,
                  allowed_third_party_library_roots=[str(p) for p in libraries],
                  atol=2e-5, rtol=2e-5, started=time.time())
    try:
        for ds in ("development", "csb", "msps"):
            fixture = config["datasets"][ds]["fixture"]
            model = PVPRISM(ds, args.device)
            begin = time.monotonic()
            actual, meta = model.predict(ROOT / fixture["image"])
            if ds == "development":
                parts = []
                for name in ("contact", "crack"):
                    with np.load(ROOT / fixture["expected"][name], allow_pickle=False) as z:
                        parts.append(z["probability"])
                        any_expected = z["any_probability"]
                expected = np.stack(parts)
                any_actual = model.last_any_probability.copy()
            else:
                with np.load(ROOT / fixture["expected"], allow_pickle=False) as z:
                    expected = z["probabilities"]
            with tempfile.TemporaryDirectory(prefix="prism_anonymous_") as temporary:
                anonymous = Path(temporary) / "unlabelled_query.png"
                shutil.copy2(ROOT / fixture["image"], anonymous)
                renamed, _ = model.predict(anonymous)
            case = dict(dataset=ds, passed=bool(np.allclose(actual, expected, atol=2e-5, rtol=2e-5)),
                        max_probability_error=float(np.max(np.abs(actual-expected))),
                        argmax_equal=bool(np.array_equal(actual.argmax(0), expected.argmax(0))),
                        anonymous_name_equal=bool(np.array_equal(actual, renamed)),
                        shape=list(actual.shape), query_gt_used=meta["query_gt_used"],
                        elapsed_seconds=time.monotonic()-begin)
            if ds == "development":
                from unittest.mock import patch
                with patch("prism.data.loading.read_labels", side_effect=AssertionError("Query GT access")):
                    query = model.engine.query(ROOT / fixture["image"])
                case["query_gt_reader_denied"] = "target" not in query
                case["threshold_equal"] = bool(np.array_equal(actual >= .5, expected >= .5))
                case["any_bitwise_equal"] = bool(np.array_equal(any_actual[0], any_expected))
                case["any_anonymous_equal"] = bool(np.array_equal(any_actual, model.last_any_probability))
                case["passed"] = case["passed"] and case["any_bitwise_equal"] and case["any_anonymous_equal"] and case["threshold_equal"]
            else:
                query = model.engine.data.query(ROOT / fixture["image"])
                try:
                    model.engine.data.target(query)
                except ValueError:
                    case["query_gt_reader_denied"] = True
                else:
                    case["query_gt_reader_denied"] = False
            result["cases"].append(case)
            print(json.dumps(case), flush=True)
            if not case["passed"] or not case["anonymous_name_equal"] or not case["query_gt_reader_denied"]:
                raise AssertionError(case)
            del model
            gc.collect()
            torch.cuda.empty_cache()
        if blocked:
            raise AssertionError("Original workspace reads were attempted")
        result.update(status="passed", blocked_access_attempts=blocked, finished=time.time())
    except BaseException as error:
        result.update(status="failed", error=repr(error), blocked_access_attempts=blocked, finished=time.time())
        target.parent.mkdir(exist_ok=True)
        target.write_text(json.dumps(result, ensure_ascii=False, indent=2)+"\n")
        raise
    target.parent.mkdir(exist_ok=True)
    target.write_text(json.dumps(result, ensure_ascii=False, indent=2)+"\n")
    print(json.dumps(dict(status=result["status"], cases=len(result["cases"]), receipt=str(target))), flush=True)


if __name__ == "__main__":
    main()