"""Step 4: Build unified YOLO 2-class dataset from synthetic + real data.

Reads the augmented COCO annotations (from synthesize_checks.py) and builds
a YOLO-format dataset with train/val/test splits.

Train and val images are tiled into overlapping crops at TILE_SIZE so that
small checkboxes (~10-16pt) appear at near-native pixel size in each tile,
rather than being shrunk to <15px when the full page is downscaled to model
input size.  Test images are kept as full pages to evaluate real-world
inference conditions (where SAHI tiling handles the same problem).

Critical constraint: test set contains ONLY real (non-synthetic) checked
examples for honest evaluation.

Split strategy:
  - All synthetic data → train (80%) + val (20%)
  - Real data → split by document:
    - ~60% of real docs → train (mixed with synthetic)
    - ~15% of real docs → val (mixed with synthetic)
    - ~25% of real docs → test (REAL ONLY)

Output:
  datasets/checkbox/yolo_2class/
  ├── dataset.yaml
  ├── train/
  │   ├── images/
  │   └── labels/
  ├── val/
  │   ├── images/
  │   └── labels/
  └── test/
      ├── images/
      └── labels/

Usage:
    python -m datasets.checkbox.tools.build_yolo_dataset
    python datasets/checkbox/tools/build_yolo_dataset.py [--oversample 6]
"""

import argparse
import json
import logging
import os
import random
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import List, Tuple

from PIL import Image
from tqdm import tqdm

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent  # datasets/checkbox/
COMMONFORMS_DIR = ROOT / "commonforms"
AUGMENTED_DIR = COMMONFORMS_DIR / "augmented"
STAGING_DIR = ROOT / "images" / "staging"
YOLO_2CLASS_DIR = ROOT / "yolo_2class"

# Class mapping (same as existing pipeline)
COCO_TO_YOLO_CLASS = {1: 0, 2: 1}  # COCO: 1=checked, 2=unchecked → YOLO: 0, 1
YOLO_CLASS_NAMES = ["checkbox_checked", "checkbox_unchecked"]

# Split fractions for real documents
REAL_TRAIN_FRAC = 0.60
REAL_VAL_FRAC = 0.15
# Remaining ~25% goes to test

# Split fractions for synthetic data
SYNTH_TRAIN_FRAC = 0.80
# Remaining 20% goes to val

DEFAULT_OVERSAMPLE = 6  # Oversample checked images in train split
MAX_ANNOTS_PER_IMAGE = 500  # High cap for CUDA; was 30 for MPS compatibility

# Tiling parameters — match model input size so checkboxes stay at native pixel size
TILE_SIZE = 1216  # Matches model imgsz
TILE_OVERLAP = 0.2  # 20% overlap between tiles
MIN_BOX_VISIBLE = 0.4  # Keep annotation if ≥40% of its area is inside the tile
MIN_BOX_PX = 4  # Drop annotations smaller than 4px in the tile


def _generate_tiles(
    img_w: int, img_h: int, tile_size: int, overlap: float,
) -> List[Tuple[int, int, int, int]]:
    """Generate overlapping tile grid for an image.

    Returns list of (x0, y0, x1, y1) tile coordinates.
    Images smaller than tile_size in both dimensions return a single tile.
    """
    if img_w <= tile_size and img_h <= tile_size:
        return [(0, 0, img_w, img_h)]

    stride = max(1, int(tile_size * (1 - overlap)))
    tiles = []
    y_positions = list(range(0, img_h, stride))
    x_positions = list(range(0, img_w, stride))
    for y in y_positions:
        for x in x_positions:
            x1 = min(x + tile_size, img_w)
            y1 = min(y + tile_size, img_h)
            x0 = max(0, x1 - tile_size)
            y0 = max(0, y1 - tile_size)
            tiles.append((x0, y0, x1, y1))

    # Deduplicate (edge tiles can overlap with snapped-back tiles)
    return list(dict.fromkeys(tiles))


def _clip_annotations_to_tile(
    annotations: list, tile: Tuple[int, int, int, int],
    img_w: int, img_h: int,
) -> list:
    """Clip COCO annotations to a tile, keeping boxes with sufficient overlap.

    Args:
        annotations: COCO annotations with 'bbox' [x, y, w, h] in pixel coords.
        tile: (tx0, ty0, tx1, ty1) tile bounds.
        img_w, img_h: Original image dimensions.

    Returns:
        List of annotations remapped to tile-local coordinates.
    """
    tx0, ty0, tx1, ty1 = tile
    tile_w = tx1 - tx0
    tile_h = ty1 - ty0
    clipped = []

    for annot in annotations:
        if COCO_TO_YOLO_CLASS.get(annot["category_id"]) is None:
            continue

        bx, by, bw, bh = annot["bbox"]
        # Box bounds in image coords
        ax0 = bx
        ay0 = by
        ax1 = bx + bw
        ay1 = by + bh

        # Intersection with tile
        ix0 = max(ax0, tx0)
        iy0 = max(ay0, ty0)
        ix1 = min(ax1, tx1)
        iy1 = min(ay1, ty1)

        if ix0 >= ix1 or iy0 >= iy1:
            continue  # No overlap

        # Check visibility fraction
        box_area = max(bw * bh, 1e-6)
        visible_area = (ix1 - ix0) * (iy1 - iy0)
        if visible_area / box_area < MIN_BOX_VISIBLE:
            continue

        # Remap to tile-local coords (clipped to tile bounds)
        local_x = ix0 - tx0
        local_y = iy0 - ty0
        local_w = ix1 - ix0
        local_h = iy1 - iy0

        # Skip tiny boxes
        if local_w < MIN_BOX_PX or local_h < MIN_BOX_PX:
            continue

        clipped.append({
            **annot,
            "bbox": [local_x, local_y, local_w, local_h],
            "_tile_img_w": tile_w,
            "_tile_img_h": tile_h,
        })

    return clipped


def _split_list(items: list, fractions: list, seed: int = 42) -> list:
    """Split a list into N groups according to fractions.

    Args:
        items: List to split.
        fractions: List of fractions (must sum to <= 1.0).

    Returns:
        List of lists, one per fraction + remainder.
    """
    rng = random.Random(seed)
    shuffled = list(items)
    rng.shuffle(shuffled)

    splits = []
    offset = 0
    for frac in fractions:
        n = int(len(shuffled) * frac)
        splits.append(shuffled[offset:offset + n])
        offset += n

    # Remainder
    splits.append(shuffled[offset:])
    return splits


def _convert_bbox_to_yolo(bbox_coco: list, img_w: int, img_h: int) -> str:
    """Convert COCO bbox [x, y, w, h] to YOLO normalized [cx, cy, nw, nh]."""
    bx, by, bw, bh = bbox_coco
    x0 = max(0.0, bx)
    y0 = max(0.0, by)
    x1 = min(float(img_w), bx + bw)
    y1 = min(float(img_h), by + bh)

    cx = ((x0 + x1) / 2) / img_w
    cy = ((y0 + y1) / 2) / img_h
    nw = (x1 - x0) / img_w
    nh = (y1 - y0) / img_h

    return f"{cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}"


def _convert_to_greyscale(src: Path, dst: Path):
    """Convert image to greyscale (3-channel RGB for YOLO compatibility)."""
    with Image.open(src) as img:
        img.convert("L").convert("RGB").save(dst, quality=95)


def _write_label(annotations: list, img_w: int, img_h: int,
                 out_lbl_dir: Path, name: str):
    """Write a YOLO label file, capping at MAX_ANNOTS_PER_IMAGE."""
    # Filter to supported categories first
    supported = [a for a in annotations if COCO_TO_YOLO_CLASS.get(a["category_id"]) is not None]

    # Cap annotations with class-aware sampling: keep all checked (minority),
    # subsample unchecked if needed
    if len(supported) > MAX_ANNOTS_PER_IMAGE:
        checked = [a for a in supported if a["category_id"] == 1]
        unchecked = [a for a in supported if a["category_id"] != 1]
        random.shuffle(unchecked)
        k_unchecked = max(0, MAX_ANNOTS_PER_IMAGE - len(checked))
        supported = checked + unchecked[:k_unchecked]

    lines = []
    for annot in supported:
        cat_id = annot["category_id"]
        yolo_cls = COCO_TO_YOLO_CLASS[cat_id]
        yolo_coords = _convert_bbox_to_yolo(annot["bbox"], img_w, img_h)
        lines.append(f"{yolo_cls} {yolo_coords}")

    label_text = "\n".join(lines) + "\n" if lines else ""
    (out_lbl_dir / f"{name}.txt").write_text(label_text)


def build_yolo_dataset(oversample: int = DEFAULT_OVERSAMPLE, seed: int = 42):
    """Build the unified YOLO 2-class dataset.

    Args:
        oversample: How many times to duplicate images with checked boxes (train only).
        seed: Random seed for splits.
    """
    # Load COCO dataset from synthesis step
    coco_path = COMMONFORMS_DIR / "augmented_coco.json"
    if not coco_path.exists():
        raise FileNotFoundError(
            f"Augmented COCO not found: {coco_path}\n"
            "Run synthesize_checks.py first."
        )

    with open(coco_path) as f:
        coco = json.load(f)

    images = {img["id"]: img for img in coco["images"]}
    annots_by_image = {}
    for annot in coco["annotations"]:
        annots_by_image.setdefault(annot["image_id"], []).append(annot)

    # Separate synthetic vs real images
    synthetic_ids = []
    real_by_doc = {}  # doc_id → [image_id, ...]

    for img_id, img_info in images.items():
        is_synthetic = img_info.get("synthetic", False)
        source = img_info.get("source", "")

        if is_synthetic or source == "CommonForms":
            synthetic_ids.append(img_id)
        else:
            # Real data: group by document
            fname = img_info["file_name"]
            parts = fname.rsplit("_p", 1)
            doc_id = parts[0] if len(parts) == 2 else fname
            real_by_doc.setdefault(doc_id, []).append(img_id)

    logger.info(
        "Dataset: %d synthetic images, %d real images (%d documents)",
        len(synthetic_ids), sum(len(v) for v in real_by_doc.values()),
        len(real_by_doc),
    )

    # Split synthetic data: 80% train, 20% val
    synth_train, synth_val = _split_list(
        synthetic_ids, [SYNTH_TRAIN_FRAC], seed=seed
    )[:2]

    # Split real data by document: 60% train, 15% val, 25% test
    real_docs = list(real_by_doc.keys())
    doc_splits = _split_list(
        real_docs, [REAL_TRAIN_FRAC, REAL_VAL_FRAC], seed=seed + 1
    )
    train_docs, val_docs, test_docs = doc_splits[0], doc_splits[1], doc_splits[2]

    real_train_ids = [img_id for doc in train_docs for img_id in real_by_doc[doc]]
    real_val_ids = [img_id for doc in val_docs for img_id in real_by_doc[doc]]
    real_test_ids = [img_id for doc in test_docs for img_id in real_by_doc[doc]]

    # Build final splits
    splits = {
        "train": synth_train + real_train_ids,
        "val": synth_val + real_val_ids,
        "test": real_test_ids,  # REAL ONLY
    }

    logger.info(
        "Splits — train: %d (synth: %d, real: %d), val: %d (synth: %d, real: %d), test: %d (real only)",
        len(splits["train"]), len(synth_train), len(real_train_ids),
        len(splits["val"]), len(synth_val), len(real_val_ids),
        len(splits["test"]),
    )

    # Verify test set has checked examples
    test_checked = 0
    test_unchecked = 0
    for img_id in splits["test"]:
        for annot in annots_by_image.get(img_id, []):
            if annot["category_id"] == 1:
                test_checked += 1
            else:
                test_unchecked += 1

    logger.info(
        "Test set: %d checked, %d unchecked annotations",
        test_checked, test_unchecked,
    )

    if test_checked == 0:
        logger.warning(
            "WARNING: Test set has no checked examples! "
            "This means evaluation of checked recall won't be possible."
        )

    # Clean dataset directories (preserve runs/ which has training checkpoints)
    YOLO_2CLASS_DIR.mkdir(parents=True, exist_ok=True)
    for split in ("train", "val", "test"):
        p = YOLO_2CLASS_DIR / split
        if p.exists():
            shutil.rmtree(p)
    yaml_p = YOLO_2CLASS_DIR / "dataset.yaml"
    if yaml_p.exists():
        yaml_p.unlink()

    stats = {}

    checked_coco_ids = {
        coco_id for coco_id, yolo_cls in COCO_TO_YOLO_CLASS.items()
        if yolo_cls == 0
    }
    num_workers = min(os.cpu_count() or 4, 8)

    for split_name, img_ids in splits.items():
        img_dir = YOLO_2CLASS_DIR / split_name / "images"
        lbl_dir = YOLO_2CLASS_DIR / split_name / "labels"
        img_dir.mkdir(parents=True, exist_ok=True)
        lbl_dir.mkdir(parents=True, exist_ok=True)

        # Test split: full pages (evaluates real-world inference with SAHI)
        # Train/val splits: tiled into overlapping crops at TILE_SIZE
        use_tiling = split_name in ("train", "val")

        # Collect all work items
        work_items = []
        for img_id in img_ids:
            img_info = images[img_id]
            img_annots = annots_by_image.get(img_id, [])
            fname = img_info["file_name"]

            if img_info.get("synthetic", False) or img_info.get("source") == "CommonForms":
                img_path = AUGMENTED_DIR / fname
            else:
                img_path = STAGING_DIR / fname

            if not img_path.exists():
                continue

            has_checked = any(a["category_id"] in checked_coco_ids for a in img_annots)
            copies = oversample if (split_name == "train" and has_checked and oversample > 1) else 1

            work_items.append((
                img_path, Path(fname).stem, img_annots,
                img_info["width"], img_info["height"], copies,
            ))

        def _process_item_tiled(item):
            src_path, stem, annots, img_w, img_h, copies = item
            with Image.open(src_path) as img:
                grey = img.convert("L").convert("RGB")

            tiles = _generate_tiles(img_w, img_h, TILE_SIZE, TILE_OVERLAP)
            written = 0

            for ti, tile in enumerate(tiles):
                tx0, ty0, tx1, ty1 = tile
                tile_annots = _clip_annotations_to_tile(annots, tile, img_w, img_h)

                # Skip empty tiles (no annotations) — reduces noise
                if not tile_annots:
                    continue

                tile_w = tx1 - tx0
                tile_h = ty1 - ty0
                tile_name = f"{stem}_t{ti}"
                tile_img = grey.crop((tx0, ty0, tx1, ty1))
                tile_img.save(img_dir / f"{tile_name}.jpg", quality=95)

                # Write label using tile dimensions
                _write_label(tile_annots, tile_w, tile_h, lbl_dir, tile_name)
                written += 1

                # Oversampling copies
                for i in range(1, copies):
                    os_name = f"{tile_name}_os{i}"
                    shutil.copy2(
                        str(img_dir / f"{tile_name}.jpg"),
                        str(img_dir / f"{os_name}.jpg"),
                    )
                    _write_label(tile_annots, tile_w, tile_h, lbl_dir, os_name)
                    written += 1

            return written

        def _process_item_full(item):
            src_path, stem, annots, img_w, img_h, copies = item
            base_dst = img_dir / f"{stem}.jpg"
            _convert_to_greyscale(src_path, base_dst)
            _write_label(annots, img_w, img_h, lbl_dir, stem)
            written = 1
            for i in range(1, copies):
                name = f"{stem}_os{i}"
                shutil.copy2(str(base_dst), str(img_dir / f"{name}.jpg"))
                _write_label(annots, img_w, img_h, lbl_dir, name)
                written += 1
            return written

        process_fn = _process_item_tiled if use_tiling else _process_item_full

        count = 0
        with ThreadPoolExecutor(max_workers=num_workers) as pool:
            futures = {pool.submit(process_fn, item): item for item in work_items}
            for future in tqdm(as_completed(futures), total=len(futures),
                               desc=f"Building {split_name}"):
                count += future.result()

        stats[f"{split_name}_images"] = count
        logger.info("  %s: %d images written", split_name, count)

    # Write dataset.yaml with relative paths (portable to Colab)
    yaml_path = YOLO_2CLASS_DIR / "dataset.yaml"
    content = (
        f"path: .\n"
        f"train: train/images\n"
        f"val: val/images\n"
        f"test: test/images\n"
        f"\n"
        f"nc: {len(YOLO_CLASS_NAMES)}\n"
        f"names: {YOLO_CLASS_NAMES}\n"
    )
    yaml_path.write_text(content)
    logger.info("Wrote dataset YAML: %s", yaml_path)

    # Summary
    logger.info("\nDataset summary:")
    for k, v in stats.items():
        logger.info("  %s: %d", k, v)
    logger.info("  Test checked annotations: %d", test_checked)
    logger.info("  Test unchecked annotations: %d", test_unchecked)
    logger.info("\nTo train:")
    logger.info(
        "  yolo train data=%s epochs=200 imgsz=1024",
        yaml_path,
    )

    return stats


def main():
    parser = argparse.ArgumentParser(
        description="Build unified YOLO 2-class dataset"
    )
    parser.add_argument(
        "--oversample", type=int, default=DEFAULT_OVERSAMPLE,
        help=f"Oversample checked images N times in train (default: {DEFAULT_OVERSAMPLE})",
    )
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    build_yolo_dataset(oversample=args.oversample, seed=args.seed)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    main()
