"""
Visual Judge for classifying regions based on image content.

This module provides a simple visual classifier that learns from examples
to classify regions (like checkboxes) into categories. It uses basic image
metrics rather than neural networks for fast, interpretable results.
"""

import base64
import hashlib
import io
import json
import logging
import shutil
from collections import namedtuple
from collections.abc import Iterable as IterableABC
from dataclasses import dataclass
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Iterable,
    List,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    Union,
    cast,
)

import numpy as np
from PIL import Image

from natural_pdf.exceptions import NaturalPDFError

logger = logging.getLogger(__name__)


class SupportsRender(Protocol):
    def render(self, crop: bool = True) -> "Image.Image": ...


if TYPE_CHECKING:
    from natural_pdf.elements.region import Region

# Return types
Decision = namedtuple("Decision", ["label", "score"])
PickResult = namedtuple("PickResult", ["region", "index", "label", "score"])


class JudgeError(NaturalPDFError):
    """Raised when Judge operations fail."""

    pass


@dataclass
class PreviewRegion:
    image: Image.Image

    def render(self, crop: bool = True) -> Image.Image:
        return self.image


class Judge:
    """
    Visual classifier for regions using simple image metrics.

    Requires class labels to be specified. For binary classification,
    requires at least one example of each class before making decisions.

    Examples:
        Checkbox detection:
        ```python
        judge = Judge("checkboxes", labels=["unchecked", "checked"])
        judge.add(empty_box, "unchecked")
        judge.add(marked_box, "checked")

        result = judge.decide(new_box)
        if result.label == "checked":
            print("Box is checked!")
        ```

        Signature detection:
        ```python
        judge = Judge("signatures", labels=["unsigned", "signed"])
        judge.add(blank_area, "unsigned")
        judge.add(signature_area, "signed")

        result = judge.decide(new_region)
        print(f"Classification: {result.label} (confidence: {result.score})")
        ```
    """

    def __init__(
        self,
        name: str,
        labels: List[str],
        base_dir: Optional[Union[str, Path]] = None,
        target_prior: Optional[float] = None,
    ):
        """
        Initialize a Judge for visual classification.

        Args:
            name: Name for this judge (used for folder name)
            labels: Class labels (required, typically 2 for binary classification)
            base_dir: Base directory for storage. Defaults to current directory
            target_prior: Target prior probability for the FIRST label in the labels list.
                         - 0.5 (default) = neutral, treats both classes equally
                         - >0.5 = favors labels[0]
                         - <0.5 = favors labels[1]
                         Example: Judge("cb", ["checked", "unchecked"], target_prior=0.6)
                         favors detecting "checked" checkboxes.
        """
        if not labels or len(labels) != 2:
            raise JudgeError("Judge requires exactly 2 class labels (binary classification only)")

        self.name = name
        self.labels = labels
        self.target_prior = float(target_prior) if target_prior is not None else 0.5

        # Set up directory structure
        self.base_dir: Path = Path(base_dir) if base_dir is not None else Path.cwd()
        self.root_dir: Path = self.base_dir / name
        self.root_dir.mkdir(exist_ok=True)

        # Create label directories
        for label in self.labels:
            (self.root_dir / label).mkdir(exist_ok=True)
        (self.root_dir / "unlabeled").mkdir(exist_ok=True)
        (self.root_dir / "_removed").mkdir(exist_ok=True)

        # Config file
        self.config_path: Path = self.root_dir / "judge.json"

        # Load existing config or initialize
        self.thresholds: Dict[str, Dict[str, Any]] = {}
        self.metrics_info: Dict[str, Dict[str, float]] = {}
        if self.config_path.exists():
            self._load_config()

    def add(self, region: SupportsRender, label: Optional[str] = None) -> None:
        """
        Add a region to the judge's dataset.

        Args:
            region: Region object to add
            label: Class label. If None, added to unlabeled for later teaching

        Raises:
            JudgeError: If label is not in allowed labels
        """
        if label is not None and label not in self.labels:
            raise JudgeError(f"Label '{label}' not in allowed labels: {self.labels}")

        # Render region to image
        try:
            img = region.render(crop=True)
            if img is None:
                raise JudgeError("Region.render returned no image.")
            if not isinstance(img, Image.Image):
                img = Image.fromarray(np.asarray(img))
        except Exception as e:
            raise JudgeError(f"Failed to render region: {e}")

        # Convert to RGB if needed
        if img.mode != "RGB":
            img = img.convert("RGB")

        # Generate hash from image content
        img_array = np.array(img)
        img_hash = hashlib.md5(img_array.tobytes()).hexdigest()[:12]

        # Determine target directory
        target_dir = self.root_dir / (label if label else "unlabeled")
        target_path = target_dir / f"{img_hash}.png"

        # Check if hash already exists anywhere
        existing_locations = []
        for check_label in self.labels + ["unlabeled", "_removed"]:
            check_path = self.root_dir / check_label / f"{img_hash}.png"
            if check_path.exists():
                existing_locations.append(check_label)

        if existing_locations:
            logger.warning(f"Duplicate image detected (hash: {img_hash})")
            logger.warning(f"Already exists in: {', '.join(existing_locations)}")
            print(f"⚠️  Duplicate image - already exists in: {', '.join(existing_locations)}")
            return

        # Save image
        img.save(target_path)
        logger.debug(f"Added image {img_hash} to {label if label else 'unlabeled'}")

    def teach(self, labels: Optional[List[str]] = None, review: bool = False) -> None:
        """
        Interactive teaching interface using IPython widgets.

        Args:
            labels: Labels to use for teaching. Defaults to self.labels
            review: If True, review already labeled images for re-classification
        """
        # Check for IPython environment
        try:
            import ipywidgets as widgets
            from IPython.display import clear_output, display
        except ImportError:
            raise JudgeError(
                "Teaching requires IPython and ipywidgets. Use 'pip install ipywidgets'"
            )

        labels = labels or self.labels

        # Get images to review
        if review:
            # Get all labeled images for review
            files_to_review = []
            for label in self.labels:
                label_dir = self.root_dir / label
                for img_path in sorted(label_dir.glob("*.png")):
                    files_to_review.append((img_path, label))

            if not files_to_review:
                print("No labeled images to review")
                return

            # Shuffle for review
            import random

            random.shuffle(files_to_review)
            review_files = [f[0] for f in files_to_review]
            original_labels = {str(f[0]): f[1] for f in files_to_review}
        else:
            # Get unlabeled images
            unlabeled_dir = self.root_dir / "unlabeled"
            review_files = sorted(unlabeled_dir.glob("*.png"))
            original_labels = {}

            if not review_files:
                print("No unlabeled images to teach")
                return

        # State for teaching
        self._teaching_state = {
            "current_index": 0,
            "labeled_count": 0,
            "removed_count": 0,
            "files": review_files,
            "labels": labels,
            "review_mode": review,
            "original_labels": original_labels,
        }

        # Create widgets
        image_widget = widgets.Image()
        status_label = widgets.Label()

        # Create buttons for labeling
        button_layout = widgets.Layout(width="auto", margin="5px")

        btn_prev = widgets.Button(description="↑ Previous", layout=button_layout)
        btn_class1 = widgets.Button(
            description=f"← {labels[0]}", layout=button_layout, button_style="primary"
        )
        btn_class2 = widgets.Button(
            description=f"→ {labels[1]}", layout=button_layout, button_style="success"
        )
        btn_skip = widgets.Button(description="↓ Skip", layout=button_layout)
        btn_remove = widgets.Button(
            description="✗ Remove", layout=button_layout, button_style="danger"
        )

        button_box = widgets.HBox([btn_prev, btn_class1, btn_class2, btn_skip, btn_remove])

        # Keyboard shortcuts info
        info_label = widgets.Label(
            value="Keys: ↑ prev | ← "
            + labels[0]
            + " | → "
            + labels[1]
            + " | ↓ skip | Delete remove"
        )

        def update_display():
            """Update the displayed image and status."""
            state = self._teaching_state
            if 0 <= state["current_index"] < len(state["files"]):
                img_path = state["files"][state["current_index"]]
                with open(img_path, "rb") as f:
                    image_widget.value = f.read()

                # Build status text
                status_text = f"Image {state['current_index'] + 1} of {len(state['files'])}"
                if state["review_mode"]:
                    current_label = state["original_labels"].get(str(img_path), "unknown")
                    status_text += f" (Currently: {current_label})"
                status_text += f" | Labeled: {state['labeled_count']}"
                if state["removed_count"] > 0:
                    status_text += f" | Removed: {state['removed_count']}"

                status_label.value = status_text

                # Update button states
                btn_prev.disabled = state["current_index"] == 0
            else:
                status_label.value = "Teaching complete!"
                # Hide the image widget instead of showing broken image
                image_widget.layout.display = "none"
                # Disable all buttons
                btn_prev.disabled = True
                btn_class1.disabled = True
                btn_class2.disabled = True
                btn_skip.disabled = True

                # Auto-retrain
                if state["labeled_count"] > 0 or state["removed_count"] > 0:
                    clear_output(wait=True)
                    print("Teaching complete!")
                    print(f"Labeled: {state['labeled_count']} images")
                    if state["removed_count"] > 0:
                        print(f"Removed: {state['removed_count']} images")

                    if state["labeled_count"] > 0:
                        print("\nRetraining with new examples...")
                        self._retrain()
                        print("✓ Training complete! Judge is ready to use.")
                else:
                    print("No changes made.")

        def move_file_to_class(class_index):
            """Move current file to specified class."""
            state = self._teaching_state
            if state["current_index"] >= len(state["files"]):
                return

            current_file = state["files"][state["current_index"]]
            target_dir = self.root_dir / labels[class_index]
            shutil.move(str(current_file), str(target_dir / current_file.name))
            state["labeled_count"] += 1
            state["current_index"] += 1
            update_display()

        # Button callbacks
        def on_prev(b):
            state = self._teaching_state
            if state["current_index"] > 0:
                state["current_index"] -= 1
                update_display()

        def on_class1(b):
            move_file_to_class(0)

        def on_class2(b):
            move_file_to_class(1)

        def on_skip(b):
            state = self._teaching_state
            state["current_index"] += 1
            update_display()

        def on_remove(b):
            state = self._teaching_state
            if state["current_index"] >= len(state["files"]):
                return

            current_file = state["files"][state["current_index"]]
            target_dir = self.root_dir / "_removed"
            shutil.move(str(current_file), str(target_dir / current_file.name))
            state["removed_count"] += 1
            state["current_index"] += 1
            update_display()

        # Connect buttons
        btn_prev.on_click(on_prev)
        btn_class1.on_click(on_class1)
        btn_class2.on_click(on_class2)
        btn_skip.on_click(on_skip)
        btn_remove.on_click(on_remove)

        # Create output widget for keyboard handling
        output = widgets.Output()

        # Keyboard event handler
        def on_key(event):
            """Handle keyboard events."""
            if event["type"] != "keydown":
                return

            key = event["key"]

            if key == "ArrowUp":
                on_prev(None)
            elif key == "ArrowLeft":
                on_class1(None)
            elif key == "ArrowRight":
                on_class2(None)
            elif key == "ArrowDown":
                on_skip(None)
            elif key in ["Delete", "Backspace"]:
                on_remove(None)

        # Display everything
        display(status_label)
        display(image_widget)
        display(button_box)
        display(info_label)
        display(output)

        # Show first image
        update_display()

        # Try to set up keyboard handling (may not work in all environments)
        try:
            from ipyevents import Event  # type: ignore[import-untyped]

            event_handler = Event(source=output, watched_events=["keydown"])
            event_handler.on_dom_event(on_key)
        except:
            # If ipyevents not available, just use buttons
            print("Note: Install ipyevents for keyboard shortcuts: pip install ipyevents")

    def decide(
        self, regions: Union[SupportsRender, Iterable[SupportsRender]]
    ) -> Union[Decision, List[Decision]]:
        """
        Classify one or more regions.

        Args:
            regions: Single region or list of regions to classify

        Returns:
            Decision or list of Decisions with label and score

        Raises:
            JudgeError: If not enough training examples
        """
        # Check if we have examples
        for label in self.labels:
            label_dir = self.root_dir / label
            if not any(label_dir.glob("*.png")):
                raise JudgeError(f"Need at least one example of class '{label}' before deciding")

        # Ensure thresholds are current
        if not self.thresholds:
            self._retrain()

        # Normalize to list of regions
        if isinstance(regions, IterableABC) and not isinstance(regions, (str, bytes)):
            region_list: List[SupportsRender] = list(regions)
            single_input = False
        else:
            region_list = [cast(SupportsRender, regions)]
            single_input = True

        results: List[Decision] = []
        for region in region_list:
            # Extract metrics
            metrics = self._extract_metrics(region)

            # Apply thresholds with soft voting
            votes: Dict[str, float] = {label: 0.0 for label in self.labels}
            total_weight = 0.0

            for metric_name, value in metrics.items():
                if metric_name in self.thresholds:
                    metric_info = self.thresholds[metric_name]
                    weight = metric_info["accuracy"]  # This is now Youden's J

                    # For binary classification
                    label1, label2 = self.labels
                    threshold1, direction1 = metric_info["thresholds"][label1]

                    # Get standard deviations for soft voting
                    stats = self.metrics_info.get(metric_name, {})
                    s1 = stats.get(f"std_{label1}", 0.0)
                    s2 = stats.get(f"std_{label2}", 0.0)
                    scale1 = s1 if s1 > 1e-6 else 1.0
                    scale2 = s2 if s2 > 1e-6 else 1.0

                    # Calculate signed margin (positive favors label1, negative favors label2)
                    if direction1 == "higher":
                        margin = (value - threshold1) / (scale1 if value >= threshold1 else scale2)
                    else:
                        margin = (threshold1 - value) / (scale1 if value <= threshold1 else scale2)

                    # Clip margin to avoid single metric dominating
                    margin = np.clip(margin, -6, 6)

                    # Soft votes using sigmoid
                    p1 = 1.0 / (1.0 + np.exp(-margin))
                    p2 = 1.0 - p1

                    votes[label1] += weight * p1
                    votes[label2] += weight * p2
                    total_weight += weight

            # Normalize votes
            if total_weight > 0:
                for label in votes:
                    votes[label] /= total_weight
            else:
                # Fallback: uniform votes so prior still works
                for label in votes:
                    votes[label] = 0.5
                total_weight = 1.0

            # Apply prior bias correction
            def _logit(p, eps=1e-6):
                p = max(eps, min(1 - eps, p))
                return np.log(p / (1 - p))

            def _sigmoid(x):
                if x >= 0:
                    z = np.exp(-x)
                    return 1.0 / (1.0 + z)
                else:
                    z = np.exp(x)
                    return z / (1.0 + z)

            # Estimate priors from training counts
            counts = self._get_training_counts()
            label1, label2 = self.labels
            n1 = counts.get(label1, 0)
            n2 = counts.get(label2, 0)
            total = max(1, n1 + n2)

            if n1 > 0 and n2 > 0:  # Only apply bias if we have examples of both classes
                emp_prior1 = n1 / total
                emp_prior2 = n2 / total

                # Target prior (0.5/0.5 neutralizes imbalance)
                target_prior1 = self.target_prior
                target_prior2 = 1.0 - self.target_prior

                # Calculate bias
                bias1 = _logit(target_prior1) - _logit(emp_prior1)
                bias2 = _logit(target_prior2) - _logit(emp_prior2)

                # Apply bias in logit space
                v1 = _sigmoid(_logit(votes[label1]) + bias1)
                v2 = _sigmoid(_logit(votes[label2]) + bias2)

                # Renormalize
                s = v1 + v2
                votes[label1] = v1 / s
                votes[label2] = v2 / s

            # Find best label
            best_label = max(votes.items(), key=lambda x: x[1])
            results.append(Decision(label=best_label[0], score=best_label[1]))

        return results[0] if single_input else results

    def pick(
        self,
        target_label: str,
        regions: Iterable[SupportsRender],
        labels: Optional[Sequence[str]] = None,
    ) -> PickResult:
        """
        Pick which region best matches the target label.

        Args:
            target_label: The class label to look for
            regions: List of regions to choose from
            labels: Optional human-friendly labels for each region

        Returns:
            PickResult with winning region, index, label (if provided), and score

        Raises:
            JudgeError: If target_label not in allowed labels
        """
        if target_label not in self.labels:
            raise JudgeError(f"Target label '{target_label}' not in allowed labels: {self.labels}")

        # Classify all regions
        region_list = list(regions)
        decisions_result = self.decide(region_list)
        decisions = decisions_result if isinstance(decisions_result, list) else [decisions_result]

        # Find best match for target label
        best_index = -1
        best_score = -1.0

        for i, decision in enumerate(decisions):
            if decision.label == target_label and decision.score > best_score:
                best_score = decision.score
                best_index = i

        if best_index == -1:
            # No region matched the target label
            raise JudgeError(f"No region classified as '{target_label}'")

        # Build result
        region = region_list[best_index]
        label_list = list(labels) if labels is not None else None
        label = (
            label_list[best_index]
            if label_list is not None and best_index < len(label_list)
            else None
        )

        return PickResult(region=region, index=best_index, label=label, score=best_score)

    def count(self, target_label: str, regions: Iterable[SupportsRender]) -> int:
        """
        Count how many regions match the target label.

        Args:
            target_label: The class label to count
            regions: List of regions to check

        Returns:
            Number of regions classified as target_label
        """
        decisions_result = self.decide(regions)
        decisions = decisions_result if isinstance(decisions_result, list) else [decisions_result]
        return sum(1 for decision in decisions if decision.label == target_label)

    def info(self) -> None:
        """
        Show configuration and training information for this Judge.
        """
        print(f"Judge: {self.name}")
        print(f"Labels: {self.labels}")
        if self.target_prior != 0.5:
            print(
                f"Target prior: {self.target_prior:.2f} (favors '{self.labels[0]}')"
                if self.target_prior > 0.5
                else f"Target prior: {self.target_prior:.2f} (favors '{self.labels[1]}')"
            )

        # Get training counts
        counts = self._get_training_counts()
        print("\nTraining examples:")
        for label in self.labels:
            count = counts.get(label, 0)
            print(f"  {label}: {count}")

        if counts.get("unlabeled", 0) > 0:
            print(f"  unlabeled: {counts['unlabeled']}")

        # Show actual imbalance
        labeled_counts = [counts.get(label, 0) for label in self.labels]
        if all(c > 0 for c in labeled_counts):
            max_count = max(labeled_counts)
            min_count = min(labeled_counts)
            if max_count != min_count:
                # Find which is which
                majority_label: Optional[str] = None
                minority_label: Optional[str] = None
                for i, label in enumerate(self.labels):
                    if counts.get(label, 0) == max_count:
                        majority_label = label
                    if counts.get(label, 0) == min_count:
                        minority_label = label

                ratio = max_count / min_count
                if majority_label is not None and minority_label is not None:
                    print(
                        f"\nClass imbalance: {majority_label}:{minority_label} = {max_count}:{min_count} ({ratio:.1f}:1)"
                    )
                    print("  Using Youden's J weights with soft voting and prior correction")

    def inspect(self, preview: bool = True) -> None:
        """
        Inspect all stored examples, showing their true labels and predicted labels/scores.
        Useful for debugging classification issues.

        Args:
            preview: If True (default), display images inline in HTML tables (requires IPython/Jupyter).
                     If False, use text-only output.
        """
        if not self.thresholds:
            print("No trained model yet. Add examples and the model will auto-train.")
            return

        if not preview:
            # Show basic info first
            self.info()
            print("-" * 80)

            print("\nThresholds learned:")
            for metric, info in self.thresholds.items():
                weight = info["accuracy"]  # This is now Youden's J
                selection_acc = info.get(
                    "selection_accuracy", info["accuracy"]
                )  # Fallback for old models
                print(f"  {metric}: weight={weight:.3f} (selection_accuracy={selection_acc:.3f})")
                for label, (threshold, direction) in info["thresholds"].items():
                    print(f"    {label}: {direction} than {threshold:.3f}")

                # Show metric distribution info if available
                if metric in self.metrics_info:
                    metric_stats = self.metrics_info[metric]
                    for label in self.labels:
                        mean_key = f"mean_{label}"
                        std_key = f"std_{label}"
                        if mean_key in metric_stats:
                            print(
                                f"    {label} distribution: mean={metric_stats[mean_key]:.3f}, std={metric_stats[std_key]:.3f}"
                            )

        HTML = None
        display = None

        if preview:
            # HTML preview mode
            try:
                from IPython.display import HTML as _HTML
                from IPython.display import display as _display

                HTML = _HTML
                display = _display
            except ImportError:
                print("Preview mode requires IPython/Jupyter. Falling back to text mode.")
                preview = False
        if preview and (HTML is None or display is None):
            preview = False

        if preview:
            # Build HTML tables for everything
            assert HTML is not None and display is not None
            html_parts = []
            html_parts.append("<style>")
            html_parts.append("table { border-collapse: collapse; margin: 20px 0; }")
            html_parts.append("th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }")
            html_parts.append("th { background-color: #f2f2f2; font-weight: bold; }")
            html_parts.append("img { max-width: 60px; max-height: 60px; }")
            html_parts.append(".correct { color: green; }")
            html_parts.append(".incorrect { color: red; }")
            html_parts.append(".metrics { font-size: 0.9em; color: #666; }")
            html_parts.append("h3 { margin-top: 30px; }")
            html_parts.append(".imbalance-warning { background-color: #fff3cd; color: #856404; }")
            html_parts.append("</style>")

            # Configuration table
            html_parts.append("<h3>Judge Configuration</h3>")
            html_parts.append("<table>")
            html_parts.append("<tr><th>Property</th><th>Value</th></tr>")
            html_parts.append(f"<tr><td>Name</td><td>{self.name}</td></tr>")
            html_parts.append(f"<tr><td>Labels</td><td>{', '.join(self.labels)}</td></tr>")
            html_parts.append(f"<tr><td>Target Prior</td><td>{self.target_prior:.2f}")
            if self.target_prior != 0.5:
                html_parts.append(
                    f" (favors '{self.labels[0] if self.target_prior > 0.5 else self.labels[1]}')"
                )
            html_parts.append("</td></tr>")
            html_parts.append("</table>")

            # Training counts table
            counts = self._get_training_counts()
            html_parts.append("<h3>Training Examples</h3>")
            html_parts.append("<table>")
            html_parts.append("<tr><th>Class</th><th>Count</th></tr>")

            # Check for imbalance
            labeled_counts = [counts.get(label, 0) for label in self.labels]
            is_imbalanced = False
            ratio: Optional[float] = None
            if all(c > 0 for c in labeled_counts):
                max_count = max(labeled_counts)
                min_count = min(labeled_counts)
                if max_count != min_count:
                    ratio = max_count / min_count
                    is_imbalanced = ratio > 1.5

            for label in self.labels:
                count = counts.get(label, 0)
                row_class = ""
                if is_imbalanced:
                    if count == max(labeled_counts):
                        row_class = ' class="imbalance-warning"'
                html_parts.append(f"<tr{row_class}><td>{label}</td><td>{count}</td></tr>")

            if counts.get("unlabeled", 0) > 0:
                html_parts.append(f"<tr><td>unlabeled</td><td>{counts['unlabeled']}</td></tr>")

            html_parts.append("</table>")

            if is_imbalanced and ratio is not None:
                html_parts.append(
                    f"<p><em>Class imbalance detected ({ratio:.1f}:1). Using Youden's J weights with prior correction.</em></p>"
                )

            # Thresholds table
            html_parts.append("<h3>Learned Thresholds</h3>")
            html_parts.append("<table>")
            html_parts.append(
                "<tr><th>Metric</th><th>Weight (Youden's J)</th><th>Selection Accuracy</th><th>Threshold Details</th></tr>"
            )

            for metric, info in self.thresholds.items():
                weight = info["accuracy"]  # This is Youden's J
                selection_acc = info.get("selection_accuracy", weight)

                # Build threshold details
                details = []
                for label, (threshold, direction) in info["thresholds"].items():
                    details.append(f"<br>{label}: {direction} than {threshold:.3f}")

                # Add distribution info if available
                if metric in self.metrics_info:
                    metric_stats = self.metrics_info[metric]
                    details.append("<br><em>Distributions:</em>")
                    for label in self.labels:
                        mean_key = f"mean_{label}"
                        std_key = f"std_{label}"
                        if mean_key in metric_stats:
                            details.append(
                                f"<br>&nbsp;&nbsp;{label}: μ={metric_stats[mean_key]:.1f}, σ={metric_stats[std_key]:.1f}"
                            )

                html_parts.append("<tr>")
                html_parts.append(f"<td>{metric}</td>")
                html_parts.append(f"<td>{weight:.3f}</td>")
                html_parts.append(f"<td>{selection_acc:.3f}</td>")
                html_parts.append(f"<td>{''.join(details)}</td>")
                html_parts.append("</tr>")

            html_parts.append("</table>")

            all_correct = 0
            all_total = 0

            # First show labeled examples
            for true_label in self.labels:
                label_dir = self.root_dir / true_label
                examples = list(label_dir.glob("*.png"))

                if not examples:
                    continue

                html_parts.append(
                    f"<h3>Predictions: {true_label.upper()} ({len(examples)} total)</h3>"
                )
                html_parts.append("<table>")
                html_parts.append(
                    "<tr><th>Image</th><th>Status</th><th>Predicted</th><th>Score</th><th>Key Metrics</th></tr>"
                )

                correct = 0

                for img_path in sorted(examples)[:20]:  # Show max 20 per class in preview
                    # Load image
                    img = Image.open(img_path)
                    preview_region = PreviewRegion(img)

                    # Get prediction
                    decision = cast(Decision, self.decide(preview_region))
                    is_correct = decision.label == true_label
                    if is_correct:
                        correct += 1

                    # Extract metrics
                    metrics = self._extract_metrics(preview_region)

                    # Convert image to base64
                    buffered = io.BytesIO()
                    img.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()

                    # Build row
                    status_class = "correct" if is_correct else "incorrect"
                    status_symbol = "✓" if is_correct else "✗"

                    # Format key metrics
                    metric_strs = []
                    for metric, value in sorted(metrics.items()):
                        if metric in self.thresholds:
                            metric_strs.append(f"{metric}={value:.1f}")
                    metrics_html = "<br>".join(metric_strs[:3])

                    html_parts.append("<tr>")
                    html_parts.append(f'<td><img src="data:image/png;base64,{img_str}" /></td>')
                    html_parts.append(f'<td class="{status_class}">{status_symbol}</td>')
                    html_parts.append(f"<td>{decision.label}</td>")
                    html_parts.append(f"<td>{decision.score:.3f}</td>")
                    html_parts.append(f'<td class="metrics">{metrics_html}</td>')
                    html_parts.append("</tr>")

                html_parts.append("</table>")

                accuracy = correct / len(examples) if examples else 0
                all_correct += correct
                all_total += len(examples)

                if len(examples) > 20:
                    html_parts.append(f"<p><em>... and {len(examples) - 20} more</em></p>")
                html_parts.append(
                    f"<p>Accuracy for {true_label}: <strong>{accuracy:.1%}</strong> ({correct}/{len(examples)})</p>"
                )

            if all_total > 0:
                overall_accuracy = all_correct / all_total
                html_parts.append(
                    f"<h3>Overall accuracy: {overall_accuracy:.1%} ({all_correct}/{all_total})</h3>"
                )

            # Now show unlabeled examples with predictions
            unlabeled_dir = self.root_dir / "unlabeled"
            unlabeled_examples = list(unlabeled_dir.glob("*.png"))

            if unlabeled_examples:
                html_parts.append(
                    f"<h3>Predictions: UNLABELED ({len(unlabeled_examples)} total)</h3>"
                )
                html_parts.append("<table>")
                html_parts.append(
                    "<tr><th>Image</th><th>Predicted</th><th>Score</th><th>Key Metrics</th></tr>"
                )

                for img_path in sorted(unlabeled_examples)[:20]:  # Show max 20
                    # Load image
                    img = Image.open(img_path)
                    preview_region = PreviewRegion(img)

                    # Get prediction
                    decision = cast(Decision, self.decide(preview_region))

                    # Extract metrics
                    metrics = self._extract_metrics(preview_region)

                    # Convert image to base64
                    buffered = io.BytesIO()
                    img.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()

                    # Format key metrics
                    metric_strs = []
                    for metric, value in sorted(metrics.items()):
                        if metric in self.thresholds:
                            metric_strs.append(f"{metric}={value:.1f}")
                    metrics_html = "<br>".join(metric_strs[:3])

                    html_parts.append("<tr>")
                    html_parts.append(f'<td><img src="data:image/png;base64,{img_str}" /></td>')
                    html_parts.append(f"<td>{decision.label}</td>")
                    html_parts.append(f"<td>{decision.score:.3f}</td>")
                    html_parts.append(f'<td class="metrics">{metrics_html}</td>')
                    html_parts.append("</tr>")

                html_parts.append("</table>")

                if len(unlabeled_examples) > 20:
                    html_parts.append(
                        f"<p><em>... and {len(unlabeled_examples) - 20} more</em></p>"
                    )

            # Display HTML
            display(HTML("".join(html_parts)))

        else:
            # Text mode (original)
            print("\nPredictions on training data:")
            print("-" * 80)

            # Test each labeled example
            all_correct = 0
            all_total = 0

            for true_label in self.labels:
                label_dir = self.root_dir / true_label
                examples = list(label_dir.glob("*.png"))

                if not examples:
                    continue

                print(f"\n{true_label.upper()} examples ({len(examples)} total):")
                correct = 0

                for img_path in sorted(examples)[:10]:  # Show max 10 per class
                    # Load image and create mock region
                    img = Image.open(img_path)
                    preview_region = PreviewRegion(img)

                    # Get prediction
                    decision = cast(Decision, self.decide(preview_region))
                    is_correct = decision.label == true_label
                    if is_correct:
                        correct += 1

                    # Extract metrics for this example
                    metrics = self._extract_metrics(preview_region)

                    # Show result
                    status = "✓" if is_correct else "✗"
                    print(
                        f"  {status} {img_path.name}: predicted={decision.label} (score={decision.score:.3f})"
                    )

                    # Show key metric values
                    metric_strs = []
                    for metric, value in sorted(metrics.items()):
                        if metric in self.thresholds:
                            metric_strs.append(f"{metric}={value:.2f}")
                    if metric_strs:
                        print(f"     Metrics: {', '.join(metric_strs[:3])}")

                accuracy = correct / len(examples) if examples else 0
                all_correct += correct
                all_total += len(examples)

                if len(examples) > 10:
                    print(f"  ... and {len(examples) - 10} more")
                print(f"  Accuracy for {true_label}: {accuracy:.1%} ({correct}/{len(examples)})")

            if all_total > 0:
                overall_accuracy = all_correct / all_total
                print(f"\nOverall accuracy: {overall_accuracy:.1%} ({all_correct}/{all_total})")

            # Show unlabeled examples with predictions
            unlabeled_dir = self.root_dir / "unlabeled"
            unlabeled_examples = list(unlabeled_dir.glob("*.png"))

            if unlabeled_examples:
                print(f"\nUNLABELED examples ({len(unlabeled_examples)} total) - predictions:")

                for img_path in sorted(unlabeled_examples)[:10]:  # Show max 10
                    # Load image and create preview region
                    img = Image.open(img_path)
                    preview_region = PreviewRegion(img)

                    # Get prediction
                    decision = cast(Decision, self.decide(preview_region))

                    # Extract metrics
                    metrics = self._extract_metrics(preview_region)

                    print(
                        f"  {img_path.name}: predicted={decision.label} (score={decision.score:.3f})"
                    )

                    # Show key metric values
                    metric_strs = []
                    for metric, value in sorted(metrics.items()):
                        if metric in self.thresholds:
                            metric_strs.append(f"{metric}={value:.2f}")
                    if metric_strs:
                        print(f"     Metrics: {', '.join(metric_strs[:3])}")

                if len(unlabeled_examples) > 10:
                    print(f"  ... and {len(unlabeled_examples) - 10} more")

    def lookup(self, region: SupportsRender) -> Optional[Tuple[str, Image.Image]]:
        """
        Look up a region and return its hash and image if found in training data.

        Args:
            region: Region to look up

        Returns:
            Tuple of (hash, image) if found, None if not found
        """
        try:
            # Generate hash for the region
            img = region.render(crop=True)
            if img is None:
                raise JudgeError("Region.render returned no image")
            if not isinstance(img, Image.Image):
                img = Image.fromarray(np.asarray(img))
            if img.mode != "RGB":
                img = img.convert("RGB")
            img_array = np.array(img)
            img_hash = hashlib.md5(img_array.tobytes()).hexdigest()[:12]

            # Look for the image in all directories
            for subdir in ["checked", "unchecked", "unlabeled", "_removed"]:
                if subdir == "checked" or subdir == "unchecked":
                    # Only look in valid label directories
                    if subdir not in self.labels:
                        continue

                img_path = self.root_dir / subdir / f"{img_hash}.png"
                if img_path.exists():
                    stored_img = Image.open(img_path)
                    logger.debug(f"Found region in '{subdir}' with hash {img_hash}")
                    return (img_hash, stored_img)

            logger.debug(f"Region not found in training data (hash: {img_hash})")
            return None

        except Exception as e:
            logger.error(f"Failed to lookup region: {e}")
            return None

    def show(self, max_per_class: int = 10, size: Tuple[int, int] = (100, 100)) -> None:
        """
        Display a grid showing examples from each category.

        Args:
            max_per_class: Maximum number of examples to show per class
            size: Size of each image in pixels (width, height)
        """
        try:
            import ipywidgets as widgets  # type: ignore[import-untyped]
            from IPython.display import display
            from PIL import Image as PILImage
        except ImportError:
            print("Show requires IPython and ipywidgets")
            return

        # Collect images from each category
        categories = {}
        total_counts = {}
        for label in self.labels:
            label_dir = self.root_dir / label
            all_images = list(label_dir.glob("*.png"))
            total_counts[label] = len(all_images)
            images = sorted(all_images)[:max_per_class]
            if images:
                categories[label] = images

        # Add unlabeled if any
        unlabeled_dir = self.root_dir / "unlabeled"
        all_unlabeled = list(unlabeled_dir.glob("*.png"))
        total_counts["unlabeled"] = len(all_unlabeled)
        unlabeled = sorted(all_unlabeled)[:max_per_class]
        if unlabeled:
            categories["unlabeled"] = unlabeled

        if not categories:
            print("No images to show")
            return

        # Create grid layout
        rows = []

        # Check for class imbalance
        labeled_counts = {k: v for k, v in total_counts.items() if k != "unlabeled"}
        if labeled_counts and len(labeled_counts) >= 2:
            max_count = max(labeled_counts.values())
            min_count = min(labeled_counts.values())
            if min_count > 0 and max_count / min_count > 3:
                warning = widgets.HTML(
                    f'<div style="background: #fff3cd; padding: 10px; margin: 10px 0; border: 1px solid #ffeeba; border-radius: 4px;">'
                    f"<strong>⚠️ Class imbalance detected:</strong> {labeled_counts}<br>"
                    f"Consider adding more examples of the minority class for better accuracy."
                    f"</div>"
                )
                rows.append(warning)

        for category, image_paths in categories.items():
            # Category header showing total count
            shown = len(image_paths)
            total = total_counts[category]
            header_text = f"<h3>{category}"
            if shown < total:
                header_text += f" ({shown} of {total} shown)"
            else:
                header_text += f" ({total} total)"
            header_text += "</h3>"
            header = widgets.HTML(header_text)

            # Image row
            image_widgets = []
            for img_path in image_paths:
                # Load and resize image
                img = PILImage.open(img_path)
                img.thumbnail(size, PILImage.Resampling.LANCZOS)

                # Convert to bytes for display
                import io

                img_bytes = io.BytesIO()
                img.save(img_bytes, format="PNG")
                img_bytes.seek(0)

                # Create image widget
                img_widget = widgets.Image(value=img_bytes.read(), width=size[0], height=size[1])
                image_widgets.append(img_widget)

            # Create horizontal box for this category
            category_box = widgets.VBox([header, widgets.HBox(image_widgets)])
            rows.append(category_box)

        # Display all categories
        display(widgets.VBox(rows))

    def forget(self, region: Optional[SupportsRender] = None, delete: bool = False) -> None:
        """
        Clear training data, delete all files, or move a specific region to unlabeled.

        Args:
            region: If provided, move this specific region to unlabeled
            delete: If True, permanently delete all files
        """
        # Handle specific region case
        if region is not None:
            # Get hash of the region
            try:
                img = region.render(crop=True)
                if img is None:
                    raise ValueError("Region.render returned no image")
                if not isinstance(img, Image.Image):
                    img = Image.fromarray(np.asarray(img))
                if img.mode != "RGB":
                    img = img.convert("RGB")
                img_array = np.array(img)
                img_hash = hashlib.md5(img_array.tobytes()).hexdigest()[:12]
            except Exception as e:
                logger.error(f"Failed to hash region: {e}")
                return

            # Find and move the image
            moved = False
            for label in self.labels + ["_removed"]:
                source_path = self.root_dir / label / f"{img_hash}.png"
                if source_path.exists():
                    target_path = self.root_dir / "unlabeled" / f"{img_hash}.png"
                    shutil.move(str(source_path), str(target_path))
                    print(f"Moved region from '{label}' to 'unlabeled'")
                    moved = True
                    break

            if not moved:
                print("Region not found in training data")
            return

        # Handle delete or clear training
        if delete:
            # Delete entire directory
            if self.root_dir.exists():
                shutil.rmtree(self.root_dir)
                print(f"Deleted all data for judge '{self.name}'")
            else:
                print(f"No data found for judge '{self.name}'")

            # Reset internal state
            self.thresholds = {}  # type: Dict[str, Dict[str, Any]]
            self.metrics_info = {}  # type: Dict[str, Dict[str, float]]

            # Recreate directory structure
            self.root_dir.mkdir(exist_ok=True)
            for label in self.labels:
                (self.root_dir / label).mkdir(exist_ok=True)
            (self.root_dir / "unlabeled").mkdir(exist_ok=True)
            (self.root_dir / "_removed").mkdir(exist_ok=True)

        else:
            # Just clear training (move everything to unlabeled)
            moved_count = 0

            # Move all labeled images back to unlabeled
            unlabeled_dir = self.root_dir / "unlabeled"
            for label in self.labels:
                label_dir = self.root_dir / label
                if label_dir.exists():
                    for img_path in label_dir.glob("*.png"):
                        shutil.move(str(img_path), str(unlabeled_dir / img_path.name))
                        moved_count += 1

            # Clear thresholds
            self.thresholds = {}  # type: Dict[str, Dict[str, Any]]
            self.metrics_info = {}  # type: Dict[str, Dict[str, float]]

            # Remove saved config
            if self.config_path.exists():
                self.config_path.unlink()

            print(f"Moved {moved_count} labeled images back to unlabeled.")
            print("Training data cleared. Judge is now untrained.")

    def save(self, path: Optional[Union[str, Path]] = None) -> None:
        """
        Save the judge configuration (auto-retrains first).

        Args:
            path: Optional path to save to. Defaults to judge.json in root directory
        """
        # Retrain with current examples
        self._retrain()

        # Save config
        save_path = Path(path) if path else self.config_path

        config = {
            "name": self.name,
            "labels": self.labels,
            "target_prior": self.target_prior,
            "thresholds": self.thresholds,
            "metrics_info": self.metrics_info,
            "training_counts": self._get_training_counts(),
        }

        with open(save_path, "w") as f:
            json.dump(config, f, indent=2)

        logger.info(f"Saved judge to {save_path}")

    @classmethod
    def load(cls, path: Union[str, Path]) -> "Judge":
        """
        Load a judge from a saved configuration.

        Args:
            path: Path to the saved judge.json file or the judge directory

        Returns:
            Loaded Judge instance
        """
        path = Path(path)

        # If path is a directory, look for judge.json inside
        if path.is_dir():
            config_path = path / "judge.json"
            base_dir = path.parent
            name = path.name
        else:
            config_path = path
            base_dir = path.parent.parent if path.parent.name != "." else path.parent
            # Try to infer name from path
            name = None

        with open(config_path, "r") as f:
            config = json.load(f)

        # Use saved name if we couldn't infer it
        if name is None:
            name = config["name"]

        # Create judge with saved config
        judge = cls(
            name,
            labels=config["labels"],
            base_dir=base_dir,
            target_prior=config.get("target_prior", 0.5),
        )  # Default to 0.5 for old configs
        judge.thresholds = cast(Dict[str, Dict[str, Any]], config.get("thresholds", {}))
        judge.metrics_info = cast(Dict[str, Dict[str, float]], config.get("metrics_info", {}))

        return judge

    # Private methods

    def _extract_metrics(self, region: SupportsRender) -> Dict[str, float]:
        """Extract image metrics from a region."""
        try:
            img = region.render(crop=True)
            if img is None:
                raise JudgeError("Region.render returned no image")
            if not isinstance(img, Image.Image):
                img_array = np.asarray(img)
                img = Image.fromarray(img_array)

            # Convert to grayscale for analysis
            gray = np.array(img.convert("L"))

            metrics = {}

            # 1. Center darkness
            h, w = gray.shape
            cy, cx = h // 2, w // 2
            center_size = min(5, h // 4, w // 4)  # Adaptive center size
            center = gray[
                max(0, cy - center_size) : min(h, cy + center_size + 1),
                max(0, cx - center_size) : min(w, cx + center_size + 1),
            ]
            metrics["center_darkness"] = 255 - np.mean(center)

            # 2. Overall darkness (ink density)
            metrics["ink_density"] = 255 - np.mean(gray)

            # 3. Dark pixel ratio
            metrics["dark_pixel_ratio"] = np.sum(gray < 200) / gray.size

            # 4. Standard deviation (complexity)
            metrics["std_dev"] = np.std(gray)

            # 5. Edge vs center ratio
            edge_size = max(2, min(h // 10, w // 10))
            edge_mask = np.zeros_like(gray, dtype=bool)
            edge_mask[:edge_size, :] = True
            edge_mask[-edge_size:, :] = True
            edge_mask[:, :edge_size] = True
            edge_mask[:, -edge_size:] = True

            edge_mean = np.mean(gray[edge_mask]) if np.any(edge_mask) else 255
            center_mean = np.mean(center)
            metrics["edge_center_ratio"] = edge_mean / (center_mean + 1)

            # 6. Diagonal density (for X patterns)
            if h > 10 and w > 10:
                diag_mask = np.zeros_like(gray, dtype=bool)
                for i in range(min(h, w)):
                    if i < h and i < w:
                        diag_mask[i, i] = True
                        diag_mask[i, w - 1 - i] = True
                metrics["diagonal_density"] = 255 - np.mean(gray[diag_mask])
            else:
                metrics["diagonal_density"] = metrics["ink_density"]

            return metrics

        except Exception as e:
            raise JudgeError(f"Failed to extract metrics: {e}")

    def _retrain(self) -> None:
        """Retrain thresholds from current examples."""
        # Collect all examples
        examples: Dict[str, List[Dict[str, float]]] = {label: [] for label in self.labels}

        for label in self.labels:
            label_dir = self.root_dir / label
            for img_path in label_dir.glob("*.png"):
                img = Image.open(img_path)
                preview_region = PreviewRegion(img)
                metrics = self._extract_metrics(preview_region)
                examples[label].append(metrics)

        # Check we have examples
        for label, exs in examples.items():
            if not exs:
                logger.warning(f"No examples for class '{label}'")
                return

        # Check for class imbalance
        example_counts = {label: len(exs) for label, exs in examples.items()}
        max_count = max(example_counts.values())
        min_count = min(example_counts.values())

        imbalance_ratio = max_count / min_count if min_count > 0 else float("inf")
        is_imbalanced = imbalance_ratio > 1.5  # Consider imbalanced if more than 1.5x difference

        if is_imbalanced:
            logger.info(
                f"Class imbalance detected: {example_counts} (ratio {imbalance_ratio:.1f}:1)"
            )
            logger.info("Using balanced accuracy for threshold selection")

        # Find best thresholds for each metric
        self.thresholds = {}  # type: Dict[str, Dict[str, Any]]
        self.metrics_info = {}  # type: Dict[str, Dict[str, float]]
        metric_candidates: List[Dict[str, Any]] = []

        all_metrics = set()
        for exs in examples.values():
            for ex in exs:
                all_metrics.update(ex.keys())

        for metric in all_metrics:
            # Get all values for this metric
            values_by_label = {}
            for label, exs in examples.items():
                values_by_label[label] = [ex.get(metric, 0) for ex in exs]

            # Find threshold that best separates classes (for binary)
            if len(self.labels) == 2:
                label1, label2 = self.labels
                vals1 = values_by_label[label1]
                vals2 = values_by_label[label2]

                # Try different thresholds
                all_vals = vals1 + vals2
                best_threshold = None
                best_accuracy = 0
                best_direction = None

                for threshold in np.percentile(all_vals, [10, 20, 30, 40, 50, 60, 70, 80, 90]):
                    # Test both directions
                    for direction in ["higher", "lower"]:
                        if direction == "higher":
                            correct1 = sum(1 for v in vals1 if v > threshold)
                            correct2 = sum(1 for v in vals2 if v <= threshold)
                        else:
                            correct1 = sum(1 for v in vals1 if v < threshold)
                            correct2 = sum(1 for v in vals2 if v >= threshold)

                        # Always use balanced accuracy for threshold selection
                        # This finds fair thresholds regardless of class imbalance
                        acc1 = correct1 / len(vals1) if len(vals1) > 0 else 0
                        acc2 = correct2 / len(vals2) if len(vals2) > 0 else 0
                        accuracy = (acc1 + acc2) / 2

                        if accuracy > best_accuracy:
                            best_accuracy = accuracy
                            best_threshold = threshold
                            best_direction = direction

                if best_threshold is None or best_direction is None:
                    continue

                # Calculate Youden's J statistic for weight (TPR - FPR)
                if best_direction == "higher":
                    tp = sum(1 for v in vals1 if v > best_threshold)
                    fn = len(vals1) - tp
                    tn = sum(1 for v in vals2 if v <= best_threshold)
                    fp = len(vals2) - tn
                else:
                    tp = sum(1 for v in vals1 if v < best_threshold)
                    fn = len(vals1) - tp
                    tn = sum(1 for v in vals2 if v >= best_threshold)
                    fp = len(vals2) - tn

                tpr = tp / len(vals1) if len(vals1) > 0 else 0
                fpr = fp / len(vals2) if len(vals2) > 0 else 0
                youden_j = max(0.0, min(1.0, tpr - fpr))

                # Store all candidates
                metric_candidates.append(
                    {
                        "metric": metric,
                        "youden_j": youden_j,
                        "selection_accuracy": best_accuracy,
                        "threshold": best_threshold,
                        "direction": best_direction,
                        "label1": label1,
                        "label2": label2,
                        "stats": {
                            "mean_" + label1: np.mean(vals1),
                            "mean_" + label2: np.mean(vals2),
                            "std_" + label1: np.std(vals1),
                            "std_" + label2: np.std(vals2),
                        },
                    }
                )

        # Sort by selection accuracy
        metric_candidates.sort(key=lambda x: x["selection_accuracy"], reverse=True)

        # Use relaxed cutoff when imbalanced
        keep_cutoff = 0.55 if is_imbalanced else 0.60

        # Keep metrics that pass cutoff, or top 3 if none pass
        kept_metrics = [m for m in metric_candidates if m["selection_accuracy"] > keep_cutoff]
        if not kept_metrics and metric_candidates:
            # Keep top 3 metrics even if they don't pass cutoff
            kept_metrics = metric_candidates[:3]
            logger.warning(
                f"No metrics passed cutoff {keep_cutoff}, keeping top {len(kept_metrics)} metrics"
            )

        # Store selected metrics
        for candidate in kept_metrics:
            metric = candidate["metric"]
            label1 = candidate["label1"]
            label2 = candidate["label2"]
            self.thresholds[metric] = {
                "accuracy": candidate["youden_j"],  # Use Youden's J as weight
                "selection_accuracy": candidate["selection_accuracy"],
                "thresholds": {
                    label1: (candidate["threshold"], candidate["direction"]),
                    label2: (
                        candidate["threshold"],
                        "lower" if candidate["direction"] == "higher" else "higher",
                    ),
                },
            }
            self.metrics_info[metric] = candidate["stats"]

    def _load_config(self) -> None:
        """Load configuration from file."""
        try:
            with open(self.config_path, "r") as f:
                config = json.load(f)

            self.thresholds = config.get("thresholds", {})
            self.metrics_info = config.get("metrics_info", {})

            # Verify labels match
            if config.get("labels") != self.labels:
                logger.warning(
                    f"Saved labels {config.get('labels')} don't match current {self.labels}"
                )

        except Exception as e:
            logger.warning(f"Failed to load config: {e}")

    def _get_training_counts(self) -> Dict[str, int]:
        """Get count of examples per class."""
        counts: Dict[str, int] = {}
        for label in self.labels:
            label_dir = self.root_dir / label
            counts[label] = len(list(label_dir.glob("*.png")))
        counts["unlabeled"] = len(list((self.root_dir / "unlabeled").glob("*.png")))
        return counts
