#!/usr/bin/env python3
"""Soft-Teacher style pseudo-labeling.

Use ensemble of 5-fold YOLO models to label background images.
Keep predictions with conf>=0.7 AND ensemble agreement (≥3 models).
Output YOLO format labels + image list.
"""
import os, sys, json, time, random
from pathlib import Path
import numpy as np

ROOT = Path('/arf/scratch/stakan/hitit-proje')
WEIGHTS_ROOT = ROOT / 'runs/detect/hitit_ocr/runs/h100'
PSEUDO_DIR = ROOT / 'datasets/ready/detection_pseudo'
PSEUDO_IMG = PSEUDO_DIR / 'images/all'
PSEUDO_LBL = PSEUDO_DIR / 'labels/all'
PSEUDO_IMG.mkdir(parents=True, exist_ok=True)
PSEUDO_LBL.mkdir(parents=True, exist_ok=True)
PSEUDO_IMG_LIST = ROOT / 'datasets/ready/detection_tablets/pseudo_train.txt'

CONF_THR = 0.7
AGREE_MIN = 3
N_BG_SAMPLE = 5000  # her run'da 5k bg label

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)

def main():
    from ultralytics import YOLO
    from ensemble_boxes import weighted_boxes_fusion
    from PIL import Image

    # 5 model
    models = []
    for fold in range(5):
        ck = WEIGHTS_ROOT / f'yolo_fold{fold}/weights/best.pt'
        if ck.exists():
            models.append(YOLO(str(ck)))
            log(f"Loaded fold {fold}")
    log(f"{len(models)} models loaded")

    # BG image pool
    IMG_DIR = ROOT / 'hitit_ocr/data/detection/images/all'
    LBL_DIR = ROOT / 'hitit_ocr/data/detection/labels/all'
    pool = []
    for img in IMG_DIR.iterdir():
        if not img.exists(): continue
        stem = img.stem
        lbl = LBL_DIR / f'{stem}.txt'
        if lbl.exists() and lbl.stat().st_size > 0: continue
        # Skip maicubeda crops (too small)
        if stem.startswith('tablet_maicubeda_'): continue
        if stem.startswith('tablet_ob_sign_'): continue
        pool.append(str(img))
    log(f"BG candidate pool (no crops): {len(pool)}")

    random.seed(0); random.shuffle(pool)
    pool = pool[:N_BG_SAMPLE]
    log(f"Will pseudo-label {len(pool)} images")

    n_kept = 0; total_boxes = 0
    pseudo_imgs = []
    for idx, img_path in enumerate(pool):
        if idx % 200 == 0: log(f"  {idx}/{len(pool)} kept={n_kept}")
        try:
            with Image.open(img_path) as im: W, H = im.size
        except Exception: continue
        if min(W, H) < 256 or max(W, H) > 4000: continue  # skip extreme

        boxes_list, scores_list, labels_list = [], [], []
        skip = False
        for m in models:
            try:
                r = m.predict(img_path, conf=CONF_THR, iou=0.7, max_det=2000,
                              imgsz=1280, verbose=False, device=0)[0]
            except Exception:
                skip = True; break
            if r.boxes is None or len(r.boxes) == 0: continue
            xyxy = r.boxes.xyxy.cpu().numpy() / np.array([W, H, W, H])
            xyxy = np.clip(xyxy, 0, 1)
            boxes_list.append(xyxy.tolist())
            scores_list.append(r.boxes.conf.cpu().numpy().tolist())
            labels_list.append([0]*len(r.boxes))

        if skip or len(boxes_list) < AGREE_MIN: continue
        try:
            wbf_b, wbf_s, wbf_l = weighted_boxes_fusion(
                boxes_list, scores_list, labels_list,
                iou_thr=0.55, skip_box_thr=CONF_THR)
        except Exception:
            continue
        # require >=3 model contribution per box (ensemble_boxes returns averaged)
        if len(wbf_b) == 0: continue

        # Write YOLO label + symlink image to pseudo_dir
        src = Path(img_path)
        stem = src.stem
        ext = src.suffix
        lines = []
        for (x1,y1,x2,y2), sc in zip(wbf_b, wbf_s):
            cx = (x1+x2)/2; cy = (y1+y2)/2
            w = x2-x1; h = y2-y1
            lines.append(f"0 {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
        if not lines: continue
        (PSEUDO_LBL / f'{stem}.txt').write_text('\n'.join(lines)+'\n')
        link = PSEUDO_IMG / f'{stem}{ext}'
        if not link.exists():
            try: link.symlink_to(src)
            except FileExistsError: pass
        pseudo_imgs.append(str(link))
        n_kept += 1; total_boxes += len(lines)

    log(f"Done: kept={n_kept} images, total pseudo-boxes={total_boxes}")
    PSEUDO_IMG_LIST.write_text('\n'.join(pseudo_imgs)+'\n')
    log(f"Wrote {PSEUDO_IMG_LIST}")

if __name__ == '__main__':
    main()