"""
AngioAI - stenosis detection: 3 x YOLOv8-seg ensemble with weighted mask fusion.

Validated on ARCADE stenosis (threshold/config chosen on VAL, reported on held-out TEST, 300 imgs):
    F1@IoU0.3 = 0.637   precision 0.549   recall 0.759   AP30 0.707   F1@IoU0.5 = 0.467   0.80 FP/img
    (see D:\\AngioAI - Commercial\\stenosis\\results\\ensemble\\)

Fusion (same as eval_ensemble_stenosis.py):
  - predictions from all members clustered by box IoU >= CLUSTER_IOU (one prediction per model per cluster)
  - cluster score = sum_m(w_m * s_m) / sum(all w)   -> a lesion seen by only one model is down-weighted
  - cluster mask  = weighted pixel vote >= 0.5

Usage:
    ens = StenosisEnsemble()                 # loads weights from ./weights/stenosis or ANGIOAI_WEIGHTS
    dets = ens.predict(frame_rgb, score_thr=0.20)
    # -> [{"bbox":[x1,y1,x2,y2], "score":float, "mask":uint8[H,W], "n_models":int, "models":[...]}, ...]
"""
import os
from pathlib import Path

import numpy as np

WEIGHTS_DIR = Path(os.environ.get("ANGIOAI_WEIGHTS", Path(__file__).resolve().parent / "weights")) / "stenosis"

# name: (file, imgsz used in training, fusion weight = val AP@IoU0.3 of the single model)
MEMBERS = {
    "yolov8x_seg":   ("yolov8x_seg.pt",   640, 0.6407),
    "yolov8l_sgd":   ("yolov8l_sgd.pt",   512, 0.6173),
    "yolov8l_adamw": ("yolov8l_adamw.pt", 512, 0.6107),
}
CLUSTER_IOU = 0.20
DEFAULT_SCORE_THR = 0.20      # chosen on ARCADE val (max F1@IoU0.3)
MEMBER_CONF = 0.01            # keep low-confidence member preds; the fused score does the filtering
MAX_DET = 30
HF_REPO = os.environ.get("ANGIOAI_HF_REPO", "MuhammadAdil63/angio-ai-checkpoints")


def _resolve(fname: str) -> str:
    local = WEIGHTS_DIR / fname
    if local.exists():
        return str(local)
    from huggingface_hub import hf_hub_download          # fallback for cloud deployment
    return hf_hub_download(repo_id=HF_REPO, filename=f"stenosis/{fname}", repo_type="model",
                           token=os.environ.get("HF_TOKEN"))


def _box_iou(a, b):
    ix = max(0, min(a[2], b[2]) - max(a[0], b[0])); iy = max(0, min(a[3], b[3]) - max(a[1], b[1]))
    i = ix * iy
    u = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - i
    return i / u if u else 0.0


def _mask_box(m):
    ys, xs = np.nonzero(m)
    return (int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1) if len(ys) else None


class StenosisEnsemble:
    def __init__(self, device=None, members=None):
        from ultralytics import YOLO
        import torch
        self.device = device if device is not None else (0 if torch.cuda.is_available() else "cpu")
        self.members = members or MEMBERS
        self.models = {n: YOLO(_resolve(f)) for n, (f, _, _) in self.members.items()}
        self.W = sum(w for _, _, w in self.members.values())

    def _member_preds(self, frame_rgb):
        H, W = frame_rgb.shape[:2]
        bgr = np.ascontiguousarray(frame_rgb[..., ::-1])      # ultralytics expects BGR numpy
        out = {}
        for n, model in self.models.items():
            _, imgsz, _ = self.members[n]
            r = model.predict(bgr, imgsz=imgsz, conf=MEMBER_CONF, max_det=MAX_DET,
                              retina_masks=True, verbose=False, device=self.device)[0]
            preds = []
            if r.masks is not None:
                for m, s in zip(r.masks.data.cpu().numpy() > 0.5, r.boxes.conf.cpu().numpy()):
                    if m.shape != (H, W):
                        import cv2
                        m = cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST) > 0
                    b = _mask_box(m)
                    if b:
                        preds.append((float(s), b, m))
            out[n] = preds
        return out

    def fuse(self, member_preds):
        cand = sorted(((s, b, m, n) for n, ps in member_preds.items() for s, b, m in ps),
                      key=lambda x: -(x[0] * self.members[x[3]][2]))
        clusters = []
        for s, b, m, n in cand:
            for c in clusters:
                if n not in c["m"] and _box_iou(b, c["ref"]) >= CLUSTER_IOU:
                    c["m"][n] = (s, m); break
            else:
                clusters.append({"ref": b, "m": {n: (s, m)}})
        dets = []
        for c in clusters:
            score = sum(self.members[n][2] * s for n, (s, _) in c["m"].items()) / self.W
            acc = None; tot = 0.0
            for n, (s, m) in c["m"].items():
                w = self.members[n][2] * s
                acc = m * w if acc is None else acc + m * w
                tot += w
            mask = acc >= 0.5 * tot
            b = _mask_box(mask)
            if b:
                dets.append({"bbox": [b[0], b[1], b[2] - 1, b[3] - 1], "score": float(score),
                             "mask": mask.astype(np.uint8), "n_models": len(c["m"]),
                             "models": sorted(c["m"])})
        return sorted(dets, key=lambda d: -d["score"])

    def predict(self, frame_rgb, score_thr=DEFAULT_SCORE_THR):
        return [d for d in self.fuse(self._member_preds(frame_rgb)) if d["score"] >= score_thr]