"""PV-PRISM: complete reference-guided, multi-scale defect segmentation."""
from collections import OrderedDict
import copy
import hashlib
import json
import os
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parent
sys.dont_write_bytecode = True
os.environ.setdefault("MKL_THREADING_LAYER", "GNU")
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("HF_HUB_OFFLINE", "1")
sys.path.insert(0, str(ROOT))

import torch
import torch._dynamo
if torch.__version__.split("+")[0] != "2.7.1":
    raise RuntimeError("Use the verified PyTorch 2.7.1 environment; see launch.py and requirements.txt")
from torch.nn import functional as F
from prism.pipeline.common import Dataset, flips
from prism.pipeline.evidence import Evidence
from prism.pipeline.components import CachedReferenceSegmenter, LocalExpert, FeatureCache, RegionVerifier, probabilities
from prism.models.encoder import FrozenEncoder
from prism.models.reference_bank import NormalBank
from prism.training.source_pretraining import load_heads
from prism.data.geometry import native_logits
from prism.data.geometry import read_image, make_canvas
from prism.data.loading import load_record
from prism.models.reference_segmenter import ReferenceSegmenter


def read(path):
    return json.loads(Path(path).read_text())


def sha(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for block in iter(lambda: f.read(8 << 20), b""):
            h.update(block)
    return h.hexdigest()


def asset(relative):
    p = (ROOT / relative).resolve()
    if not p.is_relative_to(ROOT) or not p.is_file():
        raise ValueError(f"Missing or external asset: {relative}")
    return p


class SupportDataset(Dataset):
    """Query images and annotated training supports have separate interfaces."""
    def __init__(self, dataset, config):
        self.name, self.spec = dataset, config["specification"]
        self.train = []
        for entry in config["supports"]:
            row = dict(entry)
            row["image_path"] = str(asset(row["image_path"]))
            row["mask_path"] = str(asset(row["mask_path"]))
            self.train.append(row)
        self.dev = []
        self.by_id = {r["sample_id"]: r for r in self.train}
        self.image_cache, self.mask_cache = OrderedDict(), OrderedDict()
        self.limit, self.checked = 48, set()
        self.support_ids = set(self.by_id)
        self.known_groups = config["known_image_groups"]

    def target(self, row):
        if row["sample_id"] not in self.support_ids:
            raise ValueError("Query labels are inaccessible to PV-PRISM inference")
        return super().target(row)

    def query(self, path):
        path = Path(path).resolve()
        digest = sha(path)
        return dict(sample_id="query-" + digest, image_path=str(path), sha256=digest,
                    group_id=self.known_groups.get(digest, "external-" + digest))


class ReferenceEvidence(Evidence):
    """Uses the unchanged registered evidence computation with local assets."""
    def __init__(self, dataset, device, registry):
        cfg = registry["datasets"][dataset]
        self.data, self.device = SupportDataset(dataset, cfg), device
        config = copy.deepcopy(registry["segmenter_config"])
        for key in ("weights", "metadata"):
            config["encoder"][key] = str(asset(config["encoder"][key]))
        config["encoder"]["vendor"] = str((ROOT / config["encoder"]["vendor"]).resolve())
        config["encoder"]["global_long_side"] = 1024
        config["pv"]["classes"] = cfg["specification"]["classes"][1:]
        self.core = CachedReferenceSegmenter(FrozenEncoder(config["encoder"]), config).to(device)
        load_heads(self.core, asset(cfg["weights"]["segmenter"]))
        self.core.eval().requires_grad_(False)
        obj = torch.load(asset(registry["normal_bank"]), map_location="cpu", weights_only=True)
        self.normal = NormalBank(**obj).validate(8).to(device)
        self.features = FeatureCache(96)
        self.expert = LocalExpert(len(self.data.spec["classes"]), False).to(device).eval().requires_grad_(False)
        cp = torch.load(asset(cfg["weights"]["expert"]), map_location="cpu", weights_only=True)
        self.expert.load_state_dict(cp["model"], strict=True)
        self.expert_calls, self.support_items = 0, {}
        self.support_selection = cfg["support_selection"]

    def supports(self, c, query_group):
        rows = [self.data.by_id[i] for i in self.support_selection[str(c)]
                if self.data.by_id[i]["group_id"] != query_group][:5]
        if len(rows) != 5 or len({r["group_id"] for r in rows}) != 5:
            raise ValueError("Five distinct non-query training support groups are required")
        out = []
        for row in rows:
            key = row["sample_id"], c
            if key not in self.support_items:
                self.support_items[key] = self.data.item(row, c, device=self.device)
            out.append(self.support_items[key])
        return out


class FrozenSupportSegmenter(ReferenceSegmenter):
    """Memoize immutable support prototypes for the contact refinement branch."""
    def __init__(self, encoder, config):
        super().__init__(encoder, config)
        self.memo = {}

    def _support(self, item, bank, preencoded=None):
        if self.training or any(p.requires_grad for p in self.parameters()) or preencoded is not None:
            raise ValueError("Support memoization requires a completely frozen model")
        key = item["image_id"]
        if key not in self.memo:
            self.memo[key] = super()._support(item, bank)
        return self.memo[key]


class DevelopmentEvidence:
    """The original two-channel task preserves canonical crack and any heads."""
    def __init__(self, cfg, registry, device):
        self.device, self.config = device, cfg
        config = copy.deepcopy(cfg["segmenter_config"])
        for key in ("weights", "metadata"):
            config["encoder"][key] = str(asset(config["encoder"][key]))
        config["encoder"]["vendor"] = str(ROOT / config["encoder"]["vendor"])
        encoder = FrozenEncoder(config["encoder"])
        self.canonical = ReferenceSegmenter(encoder, copy.deepcopy(config)).to(device).eval().requires_grad_(False)
        load_heads(self.canonical, asset(cfg["weights"]["segmenter"]))
        config["encoder"]["global_long_side"] = 1024
        self.high = FrozenSupportSegmenter(encoder, config).to(device).eval().requires_grad_(False)
        load_heads(self.high, asset(cfg["weights"]["segmenter"]))
        self.normal = NormalBank(**torch.load(asset(registry["normal_bank"]), map_location="cpu", weights_only=True)).validate(8).to(device)
        self.expert = LocalExpert(2, False).to(device).eval().requires_grad_(False)
        self.expert.load_state_dict(torch.load(asset(cfg["weights"]["expert"]), map_location="cpu", weights_only=True)["model"], strict=True)
        self.supports = {}
        for category, records in cfg["supports"].items():
            self.supports[category] = []
            for record in records:
                row = dict(record)
                row["image_path"], row["mask_path"] = str(asset(row["image_path"])), str(asset(row["mask_path"]))
                self.supports[category].append(load_record(row, category, device))

    def query(self, path, view=0):
        digest = sha(path)
        group = self.config["known_image_groups"].get(digest, "external-" + digest)
        if group in {s["group_id"] for rr in self.supports.values() for s in rr}:
            raise ValueError("A fixed training support cannot be its own query group")
        row = dict(image_id="anonymous", group_id=group, image_path=str(path))
        item = load_record(row, device=self.device, augmentation=dict(horizontal=bool(view & 1), vertical=bool(view & 2), gain=1., gamma=1.))
        if set(item) != {"image", "geometry_valid", "geometry", "group_id", "image_id"}:
            raise ValueError("Query loader exposed annotation")
        return item

    def expert_evidence(self, path, view):
        image, bits = read_image(path)
        image = flips(image, view)
        h, w = image.shape[-2:]
        scale = 768 / max(h, w)
        nh, nw = max(1, round(h * scale)), max(1, round(w * scale))
        x = F.interpolate(image[None], (nh, nw), mode="bilinear", align_corners=False)[0]
        x = F.pad(x, (0, 768 - nw, 0, 768 - nh))
        _, _, geometry = make_canvas(image, bits)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            z = self.expert(x[None].to(self.device))
            z = F.interpolate(z, x.shape[-2:], mode="bilinear", align_corners=False)[:, :1].float()
        content = F.interpolate(z[:, :, :nh, :nw], (geometry.content_h, geometry.content_w), mode="bilinear", align_corners=False)
        return F.pad(content, (0, geometry.canvas_w - geometry.content_w, 0, geometry.canvas_h - geometry.content_h))[0], geometry

    @torch.no_grad()
    def predict(self, path, refiners):
        q = self.query(path)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            out = self.canonical(q["image"], q["geometry_valid"], [self.supports[c] for c in ("contact", "crack")], self.normal, subbank_predictions=False)
        protected = native_logits(out["final"].float(), q["geometry"])[0]
        any_logit = native_logits(out["any"].float(), q["geometry"])[0]
        members = []
        for view in (0, 1, 2, 3):
            item = self.query(path, view)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                out = self.high(item["image"], item["geometry_valid"], [self.supports["contact"]], self.normal, subbank_predictions=False)
            evidence = torch.cat((out["final"][:, :1].float(), out["absolute"][:, :1].float(), out["any"].float(), out["gate"][:, :1].float()), 1)[0]
            independent, geometry = self.expert_evidence(path, view)
            if geometry != item["geometry"]:
                raise ValueError("Expert alignment changed")
            evidence = torch.cat((evidence, independent), 0)
            for refiner in refiners:
                delta = refiner.canvas_delta(item["image"], evidence, item["geometry_valid"])
                z = native_logits((evidence[:1] + delta)[None], geometry)[0, 0]
                members.append(flips(z.float(), view))
        logits = torch.stack((torch.stack(members).mean(0), protected[1]))
        if not torch.isfinite(logits).all() or not torch.isfinite(any_logit).all():
            raise FloatingPointError("Nonfinite prediction")
        return logits.sigmoid().cpu().numpy(), any_logit.sigmoid().cpu().numpy(), q


class PVPRISM:
    """One released architecture: high resolution, four views, two refiners."""
    def __init__(self, dataset, device="cuda:0", memory_fraction=.5):
        registry = read(ROOT / "configs/model.json")
        if dataset not in registry["datasets"]:
            raise ValueError(dataset)
        self.dataset, self.device = dataset, torch.device(device)
        self.config = registry["datasets"][dataset]
        if self.device.type != "cuda" or not torch.cuda.is_available():
            raise ValueError("The verified full model requires CUDA BF16; CPU parity is not claimed")
        if not 0 < memory_fraction <= 1:
            raise ValueError("memory_fraction must be in (0,1]")
        torch.set_num_threads(2)
        torch.manual_seed(17007)
        torch.cuda.set_device(self.device)
        torch.cuda.set_per_process_memory_fraction(memory_fraction, self.device)
        torch.backends.cudnn.benchmark = False
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        self.engine = (DevelopmentEvidence(self.config, registry, self.device) if dataset == "development"
                       else ReferenceEvidence(dataset, self.device, registry))
        self.last_any_probability = None
        self.refiners = []
        for relative in self.config["weights"]["refiners"]:
            model = RegionVerifier().to(self.device).eval().requires_grad_(False)
            checkpoint = torch.load(asset(relative), map_location="cpu", weights_only=True)
            model.load_state_dict(checkpoint["editor"], strict=True)
            self.refiners.append(model)

    @torch.no_grad()
    def predict(self, image):
        if self.dataset == "development":
            p, self.last_any_probability, q = self.engine.predict(image, self.refiners)
            return p, dict(model="PV-PRISM", dataset=self.dataset, classes=["contact", "crack"],
                           image_sha256=sha(image), group_id=q["group_id"], geometry=q["geometry"].to_dict(),
                           contact_global_resolution=1024, protected_global_resolution=768,
                           native_window=512, views=["I", "H", "V", "HV"], refiners=2,
                           query_gt_used=False, teacher_and_references_required=True,
                           inverse_map="native continuous logits -> independent sigmoid -> threshold 0.5",
                           any_output="independent canonical head, not a union of contact/crack")
        row = self.engine.data.query(image)
        classes = list(range(1, len(self.engine.data.spec["classes"])))
        total = None
        for view in (0, 1, 2, 3):
            q, evidence = self.engine.encode(row, classes, view)
            for refiner in self.refiners:
                logits = []
                for c, e in evidence.items():
                    delta = refiner.canvas_delta(q["image"], e, q["geometry_valid"])
                    z = native_logits((e[:1] + delta)[None], q["geometry"])[0]
                    logits.append(flips(z, view))
                z = torch.cat(logits, 0).float() / 8
                total = z if total is None else total + z
        p = probabilities(total).cpu().numpy()
        meta = dict(model="PV-PRISM", dataset=self.dataset, image_sha256=row["sha256"],
                    group_id=row["group_id"], classes=self.engine.data.spec["classes"],
                    global_resolution=1024, native_window=512, views=["I", "H", "V", "HV"],
                    refiners=2, geometry=q["geometry"].to_dict(), query_gt_used=False,
                    teacher_and_references_required=True,
                    inverse_map="native continuous logits -> softmax([0,foreground]) -> argmax")
        return p, meta