#!/usr/bin/env python3
"""Prepare CAMELYON17 pathology task data for ALE submission.

This script must be run on a machine with:
  - OpenSlide system library (libopenslide-dev / openslide-tools)
  - openslide-python (pip install openslide-python)
  - numpy, Pillow

It downloads one CAMELYON17 training slide + its XML annotation from AWS S3,
then generates:
  - input/slide_metadata.json
  - reference/reference_annotations.json (tile-level binary labels)
  - output/.gitkeep

Usage:
    # First download the slide and annotation:
    # aws s3 cp --no-sign-request \
    #   s3://camelyon-dataset/CAMELYON17/images/patient_004_node_4.tif \
    #   input/slide.tif
    # aws s3 cp --no-sign-request \
    #   s3://camelyon-dataset/CAMELYON17/annotations/patient_004_node_4.xml \
    #   raw/annotation.xml

    python3 prepare_data.py
"""

import json
import os
import xml.etree.ElementTree as ET

import numpy as np
from PIL import Image, ImageDraw

try:
    import openslide
except ImportError:
    print("ERROR: openslide-python not installed.")
    print("Run: pip install openslide-python")
    print("Also ensure libopenslide-dev is installed on the system.")
    raise


# Configuration
SLIDE_PATH = "input/slide.tif"
ANNOTATION_XML = "raw/annotation.xml"
TILE_SIZE = 256  # pixels at level 0
ANALYSIS_LEVEL = 0
TISSUE_THRESHOLD_LEVEL = 5  # Low-res level for tissue detection
OVERLAP_THRESHOLD = 0.2  # Fraction of tile area that must overlap with
                         # annotation to be labeled tumor-positive
                         # (matches HealthAgentBench threshold)

# Output paths
METADATA_PATH = "input/slide_metadata.json"
REFERENCE_PATH = "reference/reference_annotations.json"
GITKEEP_PATH = "output/.gitkeep"


def parse_xml_annotations(xml_path):
    """Parse CAMELYON17 XML annotation file.

    Returns a list of polygons, where each polygon is a list of (x, y) tuples
    in level-0 pixel coordinates.
    """
    tree = ET.parse(xml_path)
    root = tree.getroot()

    polygons = []
    # CAMELYON XML format: Annotations > Annotation > Coordinates > Coordinate
    for annotation in root.iter("Annotation"):
        coords = []
        for coord in annotation.iter("Coordinate"):
            x = float(coord.get("X"))
            y = float(coord.get("Y"))
            coords.append((x, y))
        if coords:
            polygons.append(coords)

    return polygons


def rasterize_polygon_overlap(tile_x, tile_y, tile_size, polygon,
                               raster_size=64):
    """Compute exact overlap fraction between a tile and a polygon.

    Rasterizes the polygon onto a raster_size x raster_size grid
    covering the tile area using PIL.ImageDraw, then counts filled
    pixels. This avoids the approximation error of point sampling.

    Returns the fraction of the tile area covered by the polygon
    (0.0 to 1.0).
    """
    # Tile bounds in level-0 coordinates
    tx0 = tile_x * tile_size
    ty0 = tile_y * tile_size

    # Scale polygon coordinates to raster grid
    scale_x = raster_size / tile_size
    scale_y = raster_size / tile_size
    scaled_coords = [
        ((px - tx0) * scale_x, (py - ty0) * scale_y)
        for (px, py) in polygon
    ]

    # Rasterize
    mask = Image.new("L", (raster_size, raster_size), 0)
    draw = ImageDraw.Draw(mask)
    draw.polygon(scaled_coords, fill=255)

    # Count filled pixels
    mask_arr = np.array(mask)
    filled = np.count_nonzero(mask_arr)
    total = raster_size * raster_size
    return filled / total


def detect_tissue_tiles(slide, tile_size, threshold_level=5):
    """Detect tissue-containing tiles using Otsu thresholding on a low-res
    overview.

    Returns a set of (x, y) tile coordinates that contain tissue.
    """
    # Get low-res thumbnail
    level_dims = slide.level_dimensions
    if threshold_level >= len(level_dims):
        threshold_level = len(level_dims) - 1

    thumb_w, thumb_h = level_dims[threshold_level]
    downsample = slide.level_downsamples[threshold_level]

    # Read the entire slide at the low-res level
    thumb = slide.read_region((0, 0), threshold_level, (thumb_w, thumb_h))
    thumb_gray = np.array(thumb.convert("L"))

    # Otsu threshold
    from PIL import ImageFilter

    # Simple Otsu: find threshold that minimizes intra-class variance
    hist, _ = np.histogram(thumb_gray.ravel(), bins=256, range=(0, 256))
    total = thumb_gray.size
    sum_total = np.sum(np.arange(256) * hist)

    best_thresh = 0
    best_var = 0
    sum_bg = 0
    weight_bg = 0

    for t in range(256):
        weight_bg += hist[t]
        if weight_bg == 0:
            continue
        weight_fg = total - weight_bg
        if weight_fg == 0:
            break
        sum_bg += t * hist[t]
        mean_bg = sum_bg / weight_bg
        mean_fg = (sum_total - sum_bg) / weight_fg
        var = weight_bg * weight_fg * (mean_bg - mean_fg) ** 2
        if var > best_var:
            best_var = var
            best_thresh = t

    # Tissue is darker than background (lower values)
    tissue_mask = thumb_gray < best_thresh

    # Map back to tile grid
    grid_w = slide.dimensions[0] // tile_size
    grid_h = slide.dimensions[1] // tile_size
    tissue_tiles = set()

    for ty in range(grid_h):
        for tx in range(grid_w):
            # Map tile center to low-res coordinates
            center_x = (tx * tile_size + tile_size // 2) / downsample
            center_y = (ty * tile_size + tile_size // 2) / downsample

            cx = int(min(center_x, thumb_w - 1))
            cy = int(min(center_y, thumb_h - 1))

            # Check a small region around the center
            r = max(1, int(tile_size / downsample / 4))
            x0 = max(0, cx - r)
            x1 = min(thumb_w, cx + r)
            y0 = max(0, cy - r)
            y1 = min(thumb_h, cy + r)

            region = tissue_mask[y0:y1, x0:x1]
            if region.size > 0 and region.mean() > 0.3:
                tissue_tiles.add((tx, ty))

    return tissue_tiles


def main():
    print("=" * 60)
    print("CAMELYON17 Pathology Task Data Preparation")
    print("=" * 60)

    # Check files exist
    if not os.path.exists(SLIDE_PATH):
        print(f"\nERROR: Slide not found at {SLIDE_PATH}")
        print("Download it first:")
        print(
            "  aws s3 cp --no-sign-request "
            "s3://camelyon-dataset/CAMELYON17/"
            "images/patient_004_node_4.tif input/slide.tif"
        )
        return

    if not os.path.exists(ANNOTATION_XML):
        print(f"\nERROR: Annotation not found at {ANNOTATION_XML}")
        print("Download it first:")
        print(
            "  aws s3 cp --no-sign-request "
            "s3://camelyon-dataset/CAMELYON17/"
            "annotations/patient_004_node_4.xml raw/annotation.xml"
        )
        return

    # Open slide
    print(f"\nOpening slide: {SLIDE_PATH}")
    slide = openslide.OpenSlide(SLIDE_PATH)
    width, height = slide.dimensions
    print(f"  Dimensions (level 0): {width} x {height}")
    print(f"  Levels: {slide.level_count}")
    print(f"  Level dimensions: {slide.level_dimensions}")
    print(f"  Level downsamples: {slide.level_downsamples}")

    # Extract metadata
    mpp_x = float(slide.properties.get("openslide.mpp-x", 0.25))
    mpp_y = float(slide.properties.get("openslide.mpp-y", 0.25))
    print(f"  Microns per pixel: {mpp_x} x {mpp_y}")

    grid_w = width // TILE_SIZE
    grid_h = height // TILE_SIZE
    print(f"  Analysis grid: {grid_w} x {grid_h} tiles "
          f"({TILE_SIZE}px at level {ANALYSIS_LEVEL})")

    # Generate slide_metadata.json
    metadata = {
        "slide_id": "patient_004_node_4",
        "tissue_type": "sentinel_lymph_node",
        "stain": "H&E",
        "source": "CAMELYON17_centre_0",
        "microns_per_pixel": mpp_x,
        "width": width,
        "height": height,
        "level_count": slide.level_count,
        "analysis_tile_size": TILE_SIZE,
        "analysis_level": ANALYSIS_LEVEL,
        "grid_width": grid_w,
        "grid_height": grid_h,
    }

    os.makedirs(os.path.dirname(METADATA_PATH), exist_ok=True)
    with open(METADATA_PATH, "w") as f:
        json.dump(metadata, f, indent=2)
    print(f"\nWrote: {METADATA_PATH}")

    # Parse XML annotations
    print(f"\nParsing annotations: {ANNOTATION_XML}")
    polygons = parse_xml_annotations(ANNOTATION_XML)
    print(f"  Found {len(polygons)} annotation contours")

    # Detect tissue tiles
    print("\nDetecting tissue tiles (Otsu thresholding)...")
    tissue_tiles = detect_tissue_tiles(slide, TILE_SIZE)
    print(f"  Found {len(tissue_tiles)} tissue tiles "
          f"(out of {grid_w * grid_h} total)")

    # Generate tumor tile labels
    # NOTE: We scan ALL tiles in annotation bounding boxes, not just tissue
    # tiles. Small metastatic deposits can fall in regions that Otsu
    # classifies as background at low resolution. Tiles confirmed as
    # tumor-positive are also added to the tissue set.
    print("\nGenerating tumor tile labels...")
    tumor_tiles = set()

    # First, identify all candidate tiles from annotation bounding boxes
    candidate_tiles = set()
    for polygon in polygons:
        poly_xs = [p[0] for p in polygon]
        poly_ys = [p[1] for p in polygon]
        poly_xmin, poly_xmax = min(poly_xs), max(poly_xs)
        poly_ymin, poly_ymax = min(poly_ys), max(poly_ys)

        tx_min = max(0, int(poly_xmin // TILE_SIZE))
        tx_max = min(grid_w - 1, int(poly_xmax // TILE_SIZE))
        ty_min = max(0, int(poly_ymin // TILE_SIZE))
        ty_max = min(grid_h - 1, int(poly_ymax // TILE_SIZE))

        for tx in range(tx_min, tx_max + 1):
            for ty in range(ty_min, ty_max + 1):
                candidate_tiles.add((tx, ty))

    print(f"  Annotation bounding-box candidate tiles: {len(candidate_tiles)}")

    # Check overlap for each candidate tile
    for tx, ty in sorted(candidate_tiles):
        for polygon in polygons:
            poly_xs = [p[0] for p in polygon]
            poly_ys = [p[1] for p in polygon]
            poly_xmin, poly_xmax = min(poly_xs), max(poly_xs)
            poly_ymin, poly_ymax = min(poly_ys), max(poly_ys)

            tile_xmin = tx * TILE_SIZE
            tile_xmax = (tx + 1) * TILE_SIZE
            tile_ymin = ty * TILE_SIZE
            tile_ymax = (ty + 1) * TILE_SIZE

            if (tile_xmax < poly_xmin or tile_xmin > poly_xmax or
                    tile_ymax < poly_ymin or tile_ymin > poly_ymax):
                continue

            overlap = rasterize_polygon_overlap(
                tx, ty, TILE_SIZE, polygon, raster_size=64
            )
            if overlap >= OVERLAP_THRESHOLD:
                tumor_tiles.add((tx, ty))
                # Ensure tumor tiles are also in the tissue set
                tissue_tiles.add((tx, ty))
                break

    print(f"  Found {len(tumor_tiles)} tumor tiles")

    tumor_present = len(tumor_tiles) > 0
    tumor_area_fraction = (
        len(tumor_tiles) / len(tissue_tiles) if tissue_tiles else 0.0
    )

    # Generate reference_annotations.json
    reference = {
        "slide_id": "patient_004_node_4",
        "tumor_present": tumor_present,
        "tumor_tiles": [{"x": t[0], "y": t[1]} for t in sorted(tumor_tiles)],
        "tissue_tiles": [
            {"x": t[0], "y": t[1]} for t in sorted(tissue_tiles)
        ],
        "tumor_tile_count": len(tumor_tiles),
        "total_tissue_tiles": len(tissue_tiles),
        "tumor_area_fraction": round(tumor_area_fraction, 6),
    }

    os.makedirs(os.path.dirname(REFERENCE_PATH), exist_ok=True)
    with open(REFERENCE_PATH, "w") as f:
        json.dump(reference, f, indent=2)
    print(f"\nWrote: {REFERENCE_PATH}")

    # Create output/.gitkeep
    os.makedirs(os.path.dirname(GITKEEP_PATH), exist_ok=True)
    with open(GITKEEP_PATH, "w") as f:
        pass
    print(f"Wrote: {GITKEEP_PATH}")

    slide.close()

    # Summary
    print("\n" + "=" * 60)
    print("PREPARATION COMPLETE")
    print("=" * 60)
    print(f"  Slide: patient_004_node_4.tif")
    print(f"  Grid: {grid_w} x {grid_h} tiles ({TILE_SIZE}px)")
    print(f"  Tissue tiles: {len(tissue_tiles)}")
    print(f"  Tumor tiles: {len(tumor_tiles)}")
    print(f"  Tumor present: {tumor_present}")
    print(f"  Tumor area fraction: {tumor_area_fraction:.4f}")
    print()
    print("Files ready:")
    print(f"  {METADATA_PATH}")
    print(f"  {REFERENCE_PATH}")
    print(f"  {GITKEEP_PATH}")
    print()
    print("Next steps:")
    print("  1. Verify reference_annotations.json looks correct")
    print("  2. Package input/ as input.zip (slide.tif + slide_metadata.json)")
    print("  3. Package reference/ as reference.zip")
    print("  4. Upload to ALE")


if __name__ == "__main__":
    main()