"""Explicit adapters for additional paper methods; no automatic disease votes.

SAM2 follows the author's training architecture, not the incompatible SAM1
inference_clr.py. Automatic points come from DeepLab, never a reference mask.
SCOLD architecture adapted from enalis/scold, CC BY 4.0, Nguyen Quoc et al.
"""
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import tarfile
import time
import urllib.error
import urllib.request

import cv2
import numpy as np
from PIL import Image
import torch
from torch import nn

SAM2_FILE = "01_models/03_trained_models/severity_rust.pt"
SAM3_REPO = "facebook/sam3"
SCOLD_REPO = "enalis/scold"
ROBERTA_REPO = "FacebookAI/roberta-base"
REVISIONS = {
    ROBERTA_REPO: "e2da8e2f811d1448a5b465c236feacd80ffbac7b",
    SAM3_REPO: "3c879f39826c281e95690f02c7821c4de09afae7",
    SCOLD_REPO: "f0a65acc70c028510e0fd87ce9295ccfee1c97b3",
}
IJ_VERSION = "1.54p"
IJ_SHA256 = "2e1a09961dfb41cee66ddc821b2577a41a072566ce45a49bae69267099741e20"


def download_native_source(url, path):
    """Retry transient upstream failures; checksum validation stays with callers."""
    for attempt in range(3):
        try:
            with urllib.request.urlopen(url, timeout=60) as source:
                path.write_bytes(source.read())
            return
        except urllib.error.HTTPError as exc:
            if exc.code not in {429, 500, 502, 503, 504} or attempt == 2:
                raise
            time.sleep(2 ** (attempt + 1))


def prepare_native_runtime():
    """Install isolated R dependencies on Linux once; never alter system R."""
    if os.name == "nt" or os.getenv("AGROMIND_INSTALL_NATIVE", "1") != "1":
        return
    folder = Path.home() / ".cache" / "agromind" / "native"
    folder.mkdir(parents=True, exist_ok=True)
    prefix, mamba = folder / "r-runtime", folder / "micromamba"
    marker = prefix / "agromind-pliman-2.1.0.ready"
    if not marker.exists():
        if not mamba.exists():
            archive = folder / "micromamba-2.9.0.tar.bz2"
            download_native_source("https://conda.anaconda.org/conda-forge/linux-64/micromamba-2.9.0-0.tar.bz2", archive)
            if hashlib.sha256(archive.read_bytes()).hexdigest() != "8761c382127e6363bd9e0a2451aa3ef90d071a79133f736e2f759a3bf13040dd":
                raise ValueError("Native runtime installer checksum mismatch")
            with tarfile.open(archive) as source:
                mamba.write_bytes(source.extractfile("bin/micromamba").read())
            mamba.chmod(0o755)
        if hashlib.sha256(mamba.read_bytes()).hexdigest() != "366cd9cd8be14df1ab8ed50352a82111082a36686b2d389fdb79a92c3fafb3e3":
            raise ValueError("Native runtime executable checksum mismatch")
        env = {**os.environ, "MAMBA_ROOT_PREFIX": str(folder / "mamba-cache")}
        args = [str(mamba), "create", "-y", "-p", str(prefix), "-c", "conda-forge", "-c", "bioconda",
                "--strict-channel-priority", "r-base=4.4", "bioconductor-ebimage=4.48.0", "r-sf", "r-stars", "r-terra",
                "r-doparallel", "r-foreach", "r-rcpparmadillo", "cxx-compiler", "make"]
        log = folder / "setup.log"
        package = folder / "pliman_2.1.0.tar.gz"
        download_native_source("https://cloud.r-project.org/src/contrib/Archive/pliman/pliman_2.1.0.tar.gz", package)
        if hashlib.sha256(package.read_bytes()).hexdigest() != "e73b2aeb56e97ee54a4b0b73840d6f70c65bfd96301cd07c4064a2b3bc3afb5d":
            raise ValueError("pliman source checksum mismatch")
        with log.open("w", encoding="utf-8") as output:
            subprocess.run(args, env=env, stdout=output, stderr=subprocess.STDOUT, check=True, timeout=900)
            subprocess.run([str(mamba), "run", "-p", str(prefix), "R", "CMD", "INSTALL", str(package)],
                           env=env, stdout=output, stderr=subprocess.STDOUT, check=True, timeout=600)
        subprocess.run([str(prefix / "bin" / "Rscript"), "-e", "library(pliman); library(EBImage)"], check=True, timeout=60)
        marker.write_text("pliman 2.1.0; EBImage 4.48.0; source checksums verified\n")
    os.environ["AGROMIND_RSCRIPT"] = str(prefix / "bin" / "Rscript")


def seed_points(mask, limit=5):
    """One interior point in each largest component; never read expert labels."""
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8)
    points = []
    for i in sorted(range(1, count), key=lambda i: -stats[i, cv2.CC_STAT_AREA])[:limit]:
        distance = cv2.distanceTransform((labels == i).astype(np.uint8), cv2.DIST_L2, 5)
        y, x = np.unravel_index(distance.argmax(), distance.shape)
        points.append([int(x), int(y)])
    return points


def parse_points(text, size):
    values = json.loads(text)
    if not isinstance(values, list) or not 1 <= len(values) <= 5:
        raise ValueError("Provide 1–5 symptom points as [[x,y], ...] in crop pixels.")
    a = np.asarray(values, dtype=float)
    if a.shape != (len(values), 2) or not np.isfinite(a).all() or (a < 0).any() or (a >= size).any():
        raise ValueError("Every symptom point must lie inside the selected crop.")
    return a.astype(np.float32)


class ImageEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        import timm
        self.swin = timm.create_model("swin_base_patch4_window7_224.ms_in22k", pretrained=False, features_only=True)
        self.fc1 = nn.Linear(1024 * 7 * 7, 512)

    def forward(self, x):
        return self.fc1(self.swin(x)[-1].flatten(1))


class RobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        from transformers import RobertaModel
        self.roberta = RobertaModel(config)
        self.projection = nn.Linear(config.hidden_size, 512)

    def forward(self, tokens):
        x = self.roberta(**tokens).last_hidden_state
        # Author uses an unmasked mean. Process prompts individually so batch
        # padding cannot change embeddings compared with the published example.
        return self.projection(x[:, 0] + x.mean(dim=1))


class Scold(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.image_encoder = ImageEncoder()
        self.text_encoder = RobertaEncoder(config)
        self.t_prime = nn.Parameter(torch.tensor(np.log(.07), dtype=torch.float32))
        self.b = nn.Parameter(torch.zeros(()))


def load(kind, device, download, revision, provenance, clr_repo):
    if kind == "sam2":
        from sam2.build_sam import build_sam2
        from sam2.sam2_image_predictor import SAM2ImagePredictor
        from unittest.mock import patch
        # SAM2 eagerly moves RoPE caches to CUDA, then deep-copies them when
        # cloning attention layers. ZeroGPU cannot deepcopy its virtual CUDA
        # storage during startup. Construct entirely on CPU, then move weights.
        # This short scope runs during serialized startup, never during requests.
        with patch("torch.cuda.is_available", return_value=False), torch.device("cpu"):
            model = build_sam2("configs/sam2.1/sam2.1_hiera_t.yaml", None, device="cpu", apply_postprocessing=False)
        model.load_state_dict(torch.load(download(clr_repo, SAM2_FILE), map_location="cpu", weights_only=True), strict=True)
        model.eval().to(device)
        meta = provenance(clr_repo, SAM2_FILE)
        meta.update(architecture="SAM2.1 Hiera tiny", prompt_contract="Independent positive-point prompts, candidate 0 as in training; union for display",
                    adaptation="Automatic points, if selected, come from DeepLab. Training used reference-derived points; those are prohibited in benchmark inference.")
        return model, SAM2ImagePredictor(model), meta
    if kind == "sam3":
        from transformers import Sam3Model, Sam3Processor
        dtype_name = os.getenv("AGROMIND_SAM3_DTYPE", "bfloat16" if device == "cuda" else "float32")
        if dtype_name not in {"float32", "bfloat16"}:
            raise ValueError("AGROMIND_SAM3_DTYPE must be float32 or bfloat16")
        model = Sam3Model.from_pretrained(SAM3_REPO, revision=revision(SAM3_REPO),
                    use_safetensors=True, trust_remote_code=False,
                    dtype=getattr(torch, dtype_name)).eval().to(device)
        return model, Sam3Processor.from_pretrained(SAM3_REPO, revision=revision(SAM3_REPO)), {**provenance(SAM3_REPO), "dtype": dtype_name}
    if kind == "scold":
        from transformers import RobertaConfig, RobertaTokenizer
        rev = revision(ROBERTA_REPO)
        config = RobertaConfig.from_pretrained(ROBERTA_REPO, revision=rev)
        model = Scold(config)
        model.load_state_dict(torch.load(download(SCOLD_REPO, "scold.pth"), map_location="cpu", weights_only=True), strict=True)
        meta = provenance(SCOLD_REPO, "scold.pth")
        meta.update(tokenizer_revision=rev, license="CC BY 4.0", credit="Khang Nguyen Quoc, Lan Le Thi Thu, Luyl-Da Quach",
                    source="https://huggingface.co/enalis/scold", adaptation="Trusted local architecture, strict numeric weight loading; no pretrained-backbone substitution")
        return model.eval().to(device), RobertaTokenizer.from_pretrained(ROBERTA_REPO, revision=rev), meta
    raise ValueError(kind)


def predict(kind, image, model, proc, meta, device, prompts=None, points=None, text="orange spot"):
    if kind == "sam2":
        if points is None or len(points) == 0:
            raise ValueError("SAM2 needs symptom points. No seed is unavailable evidence, not zero rust.")
        proc.set_image(np.asarray(image).copy())
        masks, scores, _ = proc.predict(point_coords=np.asarray(points, np.float32)[:, None, :],
                point_labels=np.ones((len(points), 1), np.int32), multimask_output=True)
        masks = np.asarray(masks)
        # SAM2 squeezes a single prompt batch; preserve independent prompts.
        candidates = masks[None] if masks.ndim == 3 else masks
        mask = np.any(candidates[:, 0].astype(bool), axis=0)
        if mask.shape != (image.height, image.width):
            raise ValueError("SAM2 returned a different pixel grid.")
        return mask, {**meta, "points_xy": np.asarray(points).tolist(), "mask_index": 0,
                      "predicted_mask_quality": np.asarray(scores).reshape(-1, 3)[:, 0].tolist(), "not_disease_probability": True}
    if kind == "sam3":
        inputs = proc(images=image, text=text, return_tensors="pt").to(device, dtype=model.dtype)
        output = model(**inputs)
        result = proc.post_process_instance_segmentation(output, threshold=.5, mask_threshold=.5,
                            target_sizes=inputs["original_sizes"].tolist())[0]
        masks = result["masks"].detach().cpu().numpy()
        mask = np.any(masks > 0, axis=0) if len(masks) else np.zeros((image.height, image.width), bool)
        return mask, {**meta, "text_prompt": text, "instance_count": len(masks), "score_threshold": .5,
                      "mask_threshold": .5, "warning": "Generic text-prompt segmentation, not a rust diagnosis. Zero masks does not mean healthy."}
    if kind == "scold":
        a = np.asarray(image.resize((224, 224), Image.Resampling.BILINEAR)).astype(np.float32) / 255
        x = torch.from_numpy(a.transpose(2, 0, 1)).unsqueeze(0).to(device)
        image_emb = nn.functional.normalize(model.image_encoder(x), dim=-1)
        ranking = []
        for category, text in prompts:
            tokens = proc(text, return_tensors="pt", truncation=True).to(device)
            text_emb = nn.functional.normalize(model.text_encoder(tokens), dim=-1)
            ranking.append({"category": category, "prompt": text, "cosine_similarity": float((image_emb @ text_emb.T).item())})
        return {"ranking": sorted(ranking, key=lambda r: -r["cosine_similarity"]), "provenance": meta,
                "warning": "SCOLD symptom similarity is not a diagnostic vote or calibrated probability."}, None
    raise ValueError(kind)


def imagej_jar():
    folder = Path(os.getenv("AGROMIND_TOOL_CACHE", str(Path.home() / ".cache" / "agromind")))
    folder.mkdir(parents=True, exist_ok=True)
    jar = folder / f"ij-{IJ_VERSION}.jar"
    if not jar.exists():
        with urllib.request.urlopen(f"https://repo.maven.apache.org/maven2/net/imagej/ij/{IJ_VERSION}/ij-{IJ_VERSION}.jar", timeout=60) as source:
            data = source.read(5_000_000)
        if hashlib.sha256(data).hexdigest() != IJ_SHA256:
            raise ValueError("ImageJ download checksum differs from the reviewed artifact.")
        jar.write_bytes(data)
    if hashlib.sha256(jar.read_bytes()).hexdigest() != IJ_SHA256:
        raise ValueError("ImageJ cached JAR checksum mismatch.")
    return jar


def classical(kind, image, leaf, hsb="0,50,60,255,20,255", index="NGRDI"):
    """Invoke actual ImageJ/pliman runtimes, never relabel a Python approximation."""
    base = Path(__file__).resolve().parent
    with tempfile.TemporaryDirectory(prefix="agromind_method_") as tmp:
        folder = Path(tmp)
        image.save(folder / "image.png")
        Image.fromarray(leaf.astype(np.uint8) * 255).save(folder / "leaf.png")
        if kind == "imagej":
            values = [int(v) for v in hsb.split(",")]
            if len(values) != 6 or any(not 0 <= v <= 255 for v in values) or any(values[i] > values[i+1] for i in (0, 2, 4)):
                raise ValueError("HSB needs six integers in 0–255: min,max for hue, saturation, brightness.")
            java = shutil.which("java")
            if not java:
                raise RuntimeError("ImageJ requires Java 11+; runtime is unavailable.")
            args = [java, "-Djava.awt.headless=true", "--class-path", str(imagej_jar()), str(base / "ImageJMask.java"),
                    str(folder / "image.png"), str(folder / "mask.png"), *map(str, values)]
            meta = {"library": "ImageJ", "version": IJ_VERSION, "jar_sha256": IJ_SHA256, "hsb_ranges": values}
        elif kind == "pliman":
            if index not in {"R", "G", "B", "NGRDI"}:
                raise ValueError("Unsupported pliman index.")
            rscript = os.getenv("AGROMIND_RSCRIPT") or shutil.which("Rscript")
            if not rscript:
                raise RuntimeError("pliman requires R and its installed package; runtime is unavailable.")
            args = [rscript, str(base / "pliman_mask.R"), str(folder / "image.png"), str(folder / "leaf.png"), str(folder / "mask.png"), index]
            meta = {"library": "pliman", "index": index, "threshold": "Otsu"}
        else:
            raise ValueError(kind)
        process = subprocess.run(args, capture_output=True, text=True, timeout=120, check=False)
        if process.returncode or not (folder / "mask.png").exists():
            raise RuntimeError(f"{kind} failed: {process.stderr[-2000:]}")
        with Image.open(folder / "mask.png") as im:
            mask = np.asarray(im.convert("L")) > 0
        if mask.shape != leaf.shape:
            raise ValueError("Classical method returned a different pixel grid.")
        meta.update(runtime_output=process.stdout[-1000:],
            warning="Configurable colour baseline, not disease recognition. Original paper threshold parameters were not published; this is an explicit adaptation.")
        return mask & leaf, meta