"""Step 3: Synthesize checked checkboxes by pasting templates onto CommonForms pages.

For each CommonForms page:
  1. Load image and checkbox bbox annotations
  2. Randomly select 30-50% of checkboxes to mark as "checked"
  3. For each selected checkbox:
     - Pick a random curated template
     - Resize, rotate, jitter position
     - Create alpha mask via thresholding
     - Paste using Multiply blending
     - Add photometric variation
  4. Save augmented image + COCO annotation with 2 classes

Also includes our existing real annotated data (Gemini pipeline) in the output
so the model sees actual pen/ink textures during training.

Output:
  - commonforms/augmented/{name}.png  (augmented images)
  - commonforms/augmented/{name}.json (per-image annotation in our format)
  - commonforms/augmented_coco.json   (full COCO dataset)

Usage:
    python -m datasets.checkbox.tools.synthesize_checks
    python datasets/checkbox/tools/synthesize_checks.py [--check-fraction 0.4]
"""

import argparse
import json
import logging
import os
import random
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter
from scipy.ndimage import map_coordinates, gaussian_filter
from tqdm import tqdm

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent  # datasets/checkbox/
COMMONFORMS_DIR = ROOT / "commonforms"
CURATED_DIR = ROOT / "templates" / "curated"
AUGMENTED_DIR = COMMONFORMS_DIR / "augmented"

# Real annotated data from Gemini pipeline
STAGING_DIR = ROOT / "images" / "staging"
PAGE_ANNOTS_DIR = ROOT / "raw" / "page_annotations"

# Synthesis parameters
DEFAULT_CHECK_FRACTION_MIN = 0.30
DEFAULT_CHECK_FRACTION_MAX = 0.50
SCALE_JITTER = (0.7, 1.4)  # Marks from 70% to 140% of checkbox size
ROTATION_RANGE = (-15, 15)  # Degrees
POSITION_JITTER_FRAC = 0.15  # Fraction of checkbox size (±15% offset from center)
OPACITY_RANGE = (0.7, 1.0)
ALPHA_THRESHOLD = 200  # Threshold for separating check mark from bg
ALPHA_BLUR_RADIUS = 1.0  # Gaussian blur for soft alpha edges

# Elastic transform parameters
ELASTIC_ALPHA = 15.0   # Displacement intensity (pixels)
ELASTIC_SIGMA = 3.0    # Smoothness of displacement field

# Morphological augmentation
MORPH_PROB = 0.3       # Probability of applying erosion or dilation
MORPH_KERNEL_MIN = 2   # Min kernel size
MORPH_KERNEL_MAX = 3   # Max kernel size

# COCO categories (same as existing pipeline)
COCO_CATEGORIES = [
    {"id": 1, "name": "checkbox_checked", "supercategory": "checkbox"},
    {"id": 2, "name": "checkbox_unchecked", "supercategory": "checkbox"},
]


def load_templates() -> list:
    """Load curated check mark templates from templates/curated/."""
    if not CURATED_DIR.exists():
        raise FileNotFoundError(
            f"Curated templates directory not found: {CURATED_DIR}\n"
            "Run extract_check_templates.py first, then manually curate "
            "the results into templates/curated/"
        )

    templates = []
    for p in sorted(CURATED_DIR.glob("*.png")):
        img = Image.open(p).convert("RGBA")
        templates.append({"path": p, "image": img})

    if not templates:
        raise FileNotFoundError(
            f"No PNG templates found in {CURATED_DIR}\n"
            "Copy curated check mark crops there first."
        )

    logger.info("Loaded %d curated templates", len(templates))
    return templates


def _elastic_transform(img_arr: np.ndarray, alpha: float, sigma: float,
                       rng: random.Random) -> np.ndarray:
    """Apply elastic deformation to a 2D or 3D array.

    Creates random displacement fields, smooths them with Gaussian filter,
    then applies the warping. Makes check marks look more hand-drawn.
    """
    shape = img_arr.shape[:2]
    seed = rng.randint(0, 2**31)
    rs = np.random.RandomState(seed)

    dx = gaussian_filter(rs.rand(*shape) * 2 - 1, sigma) * alpha
    dy = gaussian_filter(rs.rand(*shape) * 2 - 1, sigma) * alpha

    y, x = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing="ij")
    coords_y = np.clip(y + dy, 0, shape[0] - 1)
    coords_x = np.clip(x + dx, 0, shape[1] - 1)

    if img_arr.ndim == 3:
        result = np.zeros_like(img_arr)
        for c in range(img_arr.shape[2]):
            result[:, :, c] = map_coordinates(
                img_arr[:, :, c], [coords_y, coords_x], order=1, mode="reflect"
            )
        return result
    else:
        return map_coordinates(img_arr, [coords_y, coords_x], order=1, mode="reflect")


def _morph_augment(img_arr: np.ndarray, rng: random.Random) -> np.ndarray:
    """Randomly apply erosion or dilation to make marks thicker or thinner.

    Works on a grayscale or RGB array. For RGB, converts to gray for the
    structuring element then applies to all channels.
    """
    import cv2

    k = rng.randint(MORPH_KERNEL_MIN, MORPH_KERNEL_MAX)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))

    if rng.random() < 0.5:
        # Erosion: makes dark marks thicker (erode white → marks grow)
        return cv2.erode(img_arr, kernel, iterations=1)
    else:
        # Dilation: makes dark marks thinner (dilate white → marks shrink)
        return cv2.dilate(img_arr, kernel, iterations=1)


def create_alpha_mask(template: Image.Image) -> Image.Image:
    """Create alpha mask from template: dark pixels = check mark, white = transparent.

    Thresholds at ALPHA_THRESHOLD: pixels darker than threshold become opaque,
    lighter pixels become transparent.  Applies Gaussian blur for soft edges.
    """
    gray = template.convert("L")
    arr = np.array(gray, dtype=np.float32)

    # Invert: dark pixels → high alpha (check mark), light pixels → low alpha
    alpha = 255.0 - arr
    # Threshold: only keep pixels that were substantially dark
    threshold_val = 255.0 - ALPHA_THRESHOLD
    alpha = np.where(alpha > threshold_val, alpha, 0).astype(np.uint8)

    mask = Image.fromarray(alpha, mode="L")
    if ALPHA_BLUR_RADIUS > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(radius=ALPHA_BLUR_RADIUS))

    return mask


def multiply_blend(page_arr: np.ndarray, template_arr: np.ndarray,
                   alpha_arr: np.ndarray, opacity: float,
                   x0: int, y0: int) -> np.ndarray:
    """Apply Multiply blending of template onto page at (x0, y0).

    Multiply: result = (page * template) / 255
    This darkens where the check mark is dark, preserving box lines underneath.

    Args:
        page_arr: Page image as uint8 array (H, W, 3).
        template_arr: Template RGB as uint8 array (th, tw, 3).
        alpha_arr: Alpha mask as uint8 array (th, tw), 0-255.
        opacity: Overall opacity multiplier 0.0-1.0.
        x0, y0: Top-left position on page.

    Returns:
        Modified page array (in-place).
    """
    th, tw = template_arr.shape[:2]
    ph, pw = page_arr.shape[:2]

    # Clamp to page bounds
    sx = max(0, -x0)
    sy = max(0, -y0)
    ex = min(tw, pw - x0)
    ey = min(th, ph - y0)

    if sx >= ex or sy >= ey:
        return page_arr

    # Regions
    page_region = page_arr[y0 + sy:y0 + ey, x0 + sx:x0 + ex].astype(np.float32)
    tmpl_region = template_arr[sy:ey, sx:ex].astype(np.float32)
    alpha_region = alpha_arr[sy:ey, sx:ex].astype(np.float32) / 255.0 * opacity

    # Multiply blend
    blended = (page_region * tmpl_region) / 255.0

    # Composite with alpha
    alpha_3ch = alpha_region[:, :, np.newaxis]
    result = page_region * (1.0 - alpha_3ch) + blended * alpha_3ch

    page_arr[y0 + sy:y0 + ey, x0 + sx:x0 + ex] = np.clip(result, 0, 255).astype(
        np.uint8
    )
    return page_arr


def paste_check(page_img: Image.Image, bbox: list, template: Image.Image,
                rng: random.Random) -> Image.Image:
    """Paste a check mark template onto a checkbox region.

    Args:
        page_img: Page image (will be modified in-place via numpy).
        bbox: Checkbox bbox [x0, y0, x1, y1] in pixel coords.
        template: Curated template image (RGBA).
        rng: Random instance for reproducibility.

    Returns:
        Modified page image.
    """
    x0, y0, x1, y1 = bbox
    cb_w = x1 - x0
    cb_h = y1 - y0

    if cb_w < 3 or cb_h < 3:
        return page_img

    # Random scale jitter
    scale = rng.uniform(*SCALE_JITTER)
    target_w = max(3, int(cb_w * scale))
    target_h = max(3, int(cb_h * scale))

    # Resize template
    tmpl = template.convert("RGBA").resize((target_w, target_h), Image.LANCZOS)

    # Random rotation
    angle = rng.uniform(*ROTATION_RANGE)
    if abs(angle) > 0.5:
        tmpl = tmpl.rotate(angle, resample=Image.BICUBIC, expand=False,
                           fillcolor=(255, 255, 255, 0))

    # Convert to numpy for augmentations
    tmpl_arr = np.array(tmpl.convert("RGB"))

    # --- Elastic transform (subtle warping for hand-drawn look) ---
    tmpl_arr = _elastic_transform(tmpl_arr, ELASTIC_ALPHA, ELASTIC_SIGMA, rng)

    # --- Morphological augmentation (thicker/thinner strokes) ---
    if rng.random() < MORPH_PROB:
        tmpl_arr = _morph_augment(tmpl_arr, rng)

    # Create alpha mask from the augmented template
    alpha_mask = create_alpha_mask(Image.fromarray(tmpl_arr))

    # Position: center in checkbox with proportional jitter
    max_jx = max(1, int(cb_w * POSITION_JITTER_FRAC))
    max_jy = max(1, int(cb_h * POSITION_JITTER_FRAC))
    jx = rng.randint(-max_jx, max_jx)
    jy = rng.randint(-max_jy, max_jy)
    paste_x = x0 + (cb_w - target_w) // 2 + jx
    paste_y = y0 + (cb_h - target_h) // 2 + jy

    # Random opacity
    opacity = rng.uniform(*OPACITY_RANGE)

    # Multiply blending
    page_arr = np.array(page_img)
    alpha_arr = np.array(alpha_mask)

    page_arr = multiply_blend(page_arr, tmpl_arr, alpha_arr, opacity, paste_x, paste_y)

    return Image.fromarray(page_arr)


def synthesize_page(img_path: Path, ann_path: Path, templates: list,
                    rng: random.Random,
                    check_frac_min: float = DEFAULT_CHECK_FRACTION_MIN,
                    check_frac_max: float = DEFAULT_CHECK_FRACTION_MAX):
    """Synthesize checked checkboxes on a single CommonForms page.

    Returns (augmented_image, annotation_dict) or None if no checkboxes.
    """
    with open(ann_path) as f:
        ann = json.load(f)

    checkboxes = ann.get("checkboxes", [])
    if not checkboxes:
        return None

    # Convert to greyscale (3-channel for YOLO compatibility)
    img = Image.open(img_path).convert("L").convert("RGB")

    # Randomly select fraction of checkboxes to mark as checked
    check_frac = rng.uniform(check_frac_min, check_frac_max)
    n_to_check = max(1, int(len(checkboxes) * check_frac))
    check_indices = set(rng.sample(range(len(checkboxes)), min(n_to_check, len(checkboxes))))

    # Paste check marks
    for idx in check_indices:
        cb = checkboxes[idx]
        bbox = cb["bbox"]
        template = rng.choice(templates)["image"]
        img = paste_check(img, bbox, template, rng)

    # Build annotation with 2-class labels
    result_ann = {
        "source": ann.get("source", "CommonForms"),
        "source_index": ann.get("source_index", -1),
        "image_file": ann.get("image_file", img_path.name),
        "image_size": [img.width, img.height],
        "synthetic": True,
        "checkboxes": [],
    }

    for idx, cb in enumerate(checkboxes):
        state = "checked" if idx in check_indices else "unchecked"
        result_ann["checkboxes"].append({
            "bbox": cb["bbox"],
            "state": state,
            "synthetic": idx in check_indices,
        })

    return img, result_ann


def collect_real_data():
    """Collect real annotated data from the Gemini pipeline.

    Returns list of (image_path, annotation_dict) tuples.
    """
    ann_files = sorted(PAGE_ANNOTS_DIR.glob("*.json"))
    ann_files = [f for f in ann_files if not f.name.endswith(".raw.json")]

    results = []
    for ann_file in ann_files:
        parts = ann_file.stem.rsplit("_p", 1)
        if len(parts) != 2:
            continue
        doc_id, page_str = parts
        page_idx = int(page_str)

        with open(ann_file) as f:
            data = json.load(f)

        checkboxes = data.get("checkboxes", [])
        if not checkboxes:
            continue

        img_path = STAGING_DIR / f"{doc_id}_p{page_idx:03d}.png"
        if not img_path.exists():
            continue

        ann = {
            "source": "real",
            "source_doc": doc_id,
            "page_idx": page_idx,
            "image_file": img_path.name,
            "synthetic": False,
            "checkboxes": checkboxes,
        }

        # Get image size
        with Image.open(img_path) as img:
            ann["image_size"] = [img.width, img.height]

        results.append((img_path, ann))

    logger.info("Found %d real annotated pages", len(results))
    return results


def build_coco(all_entries: list) -> dict:
    """Build COCO-format annotation from list of (image_path, annotation_dict).

    Args:
        all_entries: List of (image_path, annotation_dict) tuples.

    Returns:
        COCO dict.
    """
    coco = {
        "info": {
            "description": "Synthetic + real checkbox detection dataset (2-class)",
            "version": "2.0",
            "year": datetime.now().year,
            "date_created": datetime.now().isoformat(),
        },
        "licenses": [],
        "categories": COCO_CATEGORIES,
        "images": [],
        "annotations": [],
    }

    image_id = 0
    annot_id = 0
    for img_path, ann in all_entries:
        img_w, img_h = ann["image_size"]
        image_id += 1

        coco["images"].append({
            "id": image_id,
            "file_name": str(img_path.name),
            "width": img_w,
            "height": img_h,
            "source": ann.get("source", "unknown"),
            "synthetic": ann.get("synthetic", False),
        })

        for cb in ann["checkboxes"]:
            bbox = cb.get("bbox", [])
            if len(bbox) != 4:
                continue

            x0, y0, x1, y1 = bbox

            is_checked = cb.get("state") == "checked"
            cat_id = 1 if is_checked else 2

            annot_id += 1
            coco["annotations"].append({
                "id": annot_id,
                "image_id": image_id,
                "category_id": cat_id,
                "bbox": [x0, y0, x1 - x0, y1 - y0],
                "area": (x1 - x0) * (y1 - y0),
                "iscrowd": 0,
            })

    return coco


def synthesize(check_frac_min: float = DEFAULT_CHECK_FRACTION_MIN,
               check_frac_max: float = DEFAULT_CHECK_FRACTION_MAX,
               seed: int = 42):
    """Run full synthesis pipeline."""
    rng = random.Random(seed)

    # Load templates
    templates = load_templates()

    # Create output dir
    AUGMENTED_DIR.mkdir(parents=True, exist_ok=True)

    # Process CommonForms pages
    cf_ann_dir = COMMONFORMS_DIR / "annotations"
    cf_img_dir = COMMONFORMS_DIR / "images"

    if not cf_ann_dir.exists():
        raise FileNotFoundError(
            f"CommonForms annotations not found: {cf_ann_dir}\n"
            "Run download_commonforms.py first."
        )

    cf_ann_files = sorted(cf_ann_dir.glob("*.json"))

    # Pre-generate per-page seeds for thread safety
    page_seeds = [rng.randint(0, 2**31) for _ in cf_ann_files]
    num_workers = min(os.cpu_count() or 4, 8)

    def _process_page(ann_file: Path, page_seed: int):
        """Process one page: synthesize + save. Returns (path, ann) or None."""
        img_name = ann_file.stem + ".png"
        img_path = cf_img_dir / img_name
        if not img_path.exists():
            return None

        page_rng = random.Random(page_seed)
        result = synthesize_page(
            img_path, ann_file, templates, page_rng,
            check_frac_min, check_frac_max,
        )
        if result is None:
            return None

        aug_img, aug_ann = result

        out_img_path = AUGMENTED_DIR / img_name
        aug_img.save(out_img_path)

        out_ann_path = AUGMENTED_DIR / (ann_file.stem + ".json")
        with open(out_ann_path, "w") as f:
            json.dump(aug_ann, f, indent=2)

        return (out_img_path, aug_ann)

    all_entries = []
    synth_checked = 0
    synth_unchecked = 0

    with ThreadPoolExecutor(max_workers=num_workers) as pool:
        futures = {
            pool.submit(_process_page, ann_file, page_seed): ann_file
            for ann_file, page_seed in zip(cf_ann_files, page_seeds)
        }
        for future in tqdm(as_completed(futures), total=len(futures),
                           desc="Synthesizing pages"):
            result = future.result()
            if result is None:
                continue
            out_img_path, aug_ann = result
            all_entries.append((out_img_path, aug_ann))
            for cb in aug_ann["checkboxes"]:
                if cb["state"] == "checked":
                    synth_checked += 1
                else:
                    synth_unchecked += 1

    logger.info(
        "Synthetic: %d pages, %d checked, %d unchecked",
        len(all_entries), synth_checked, synth_unchecked,
    )

    # Add real annotated data
    real_entries = collect_real_data()
    real_checked = sum(
        1 for _, ann in real_entries
        for cb in ann["checkboxes"] if cb.get("state") == "checked"
    )
    real_unchecked = sum(
        1 for _, ann in real_entries
        for cb in ann["checkboxes"] if cb.get("state") != "checked"
    )
    logger.info(
        "Real: %d pages, %d checked, %d unchecked",
        len(real_entries), real_checked, real_unchecked,
    )

    all_entries.extend(real_entries)

    # Build COCO dataset
    coco = build_coco(all_entries)
    coco_path = COMMONFORMS_DIR / "augmented_coco.json"
    with open(coco_path, "w") as f:
        json.dump(coco, f, indent=2)

    total_images = len(coco["images"])
    total_annots = len(coco["annotations"])
    logger.info(
        "\nTotal dataset: %d images, %d annotations",
        total_images, total_annots,
    )
    logger.info("COCO annotations: %s", coco_path)

    return {
        "synthetic_pages": len(all_entries) - len(real_entries),
        "real_pages": len(real_entries),
        "total_pages": total_images,
        "synth_checked": synth_checked,
        "synth_unchecked": synth_unchecked,
        "real_checked": real_checked,
        "real_unchecked": real_unchecked,
    }


def main():
    parser = argparse.ArgumentParser(
        description="Synthesize checked checkboxes onto CommonForms pages"
    )
    parser.add_argument(
        "--check-frac-min", type=float, default=DEFAULT_CHECK_FRACTION_MIN,
        help=f"Min fraction of checkboxes to check (default: {DEFAULT_CHECK_FRACTION_MIN})",
    )
    parser.add_argument(
        "--check-frac-max", type=float, default=DEFAULT_CHECK_FRACTION_MAX,
        help=f"Max fraction of checkboxes to check (default: {DEFAULT_CHECK_FRACTION_MAX})",
    )
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    synthesize(
        check_frac_min=args.check_frac_min,
        check_frac_max=args.check_frac_max,
        seed=args.seed,
    )


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