report.py

#!/usr/bin/env python3
"""Generate detection rate report by category, generator, and detector."""

import argparse
from collections import defaultdict

import numpy as np
from scipy.stats import ttest_rel

from config import (
    _connect,
    normalize_model,
    MAIN_MODELS,
    MODEL_DISPLAY_NAMES,
    _ADV_VERSIONS,
    _VERSION_DISPLAY,
    DETECTOR_DISPLAY_NAMES,
    DATA_DIR,
)

DB_PATH = f"{DATA_DIR}/images.db"
THRESHOLD = 0.5


def _get_detectors(conn, ignore: list = None) -> list[str]:
    """Get canonical detector names from the database, optionally filtering."""
    detectors = [
        row[0]
        for row in conn.execute(
            "SELECT DISTINCT detector FROM norm_det ORDER BY detector"
        ).fetchall()
    ]
    if ignore:
        detectors = [d for d in detectors if d not in ignore]
    return detectors


def get_dataset_stats(group: int = None):
    """Get dataset statistics: image counts by category and model."""
    conn = _connect()

    # Base images per category
    base_query = """
        SELECT category, COUNT(*) as count
        FROM images
        WHERE valid = 1
    """
    base_params = []
    if group is not None:
        base_query += ' AND "group" = ?'
        base_params.append(group)
    base_query += " GROUP BY category ORDER BY category"

    base_counts = {}
    for row in conn.execute(base_query, base_params).fetchall():
        base_counts[row["category"]] = row["count"]

    # Generated images per category/model
    gen_query = """
        SELECT i.category, g.model, COUNT(*) as count
        FROM generated_images g
        JOIN images i ON g.base_image_id = i.id
        WHERE i.valid = 1
    """
    gen_params = []
    if group is not None:
        gen_query += ' AND i."group" = ?'
        gen_params.append(group)
    gen_query += " GROUP BY i.category, g.model ORDER BY i.category, g.model"

    gen_counts = defaultdict(dict)
    all_models = set()
    for row in conn.execute(gen_query, gen_params).fetchall():
        model = normalize_model(row["model"])
        gen_counts[row["category"]][model] = (
            gen_counts[row["category"]].get(model, 0) + row["count"]
        )
        all_models.add(model)

    # Prompts per category
    prompt_query = """
        SELECT i.category, COUNT(*) as count
        FROM image_prompts p
        JOIN images i ON p.image_id = i.id
        WHERE i.valid = 1
    """
    prompt_params = []
    if group is not None:
        prompt_query += ' AND i."group" = ?'
        prompt_params.append(group)
    prompt_query += " GROUP BY i.category ORDER BY i.category"

    prompt_counts = {}
    for row in conn.execute(prompt_query, prompt_params).fetchall():
        prompt_counts[row["category"]] = row["count"]

    conn.close()
    return base_counts, gen_counts, prompt_counts, sorted(all_models)


def print_dataset_stats(group: int = None, models_per_table: int = 6):
    """Print dataset statistics summary, splitting into multiple tables if needed."""
    base_counts, gen_counts, prompt_counts, all_models = get_dataset_stats(group=group)

    categories = sorted(set(base_counts.keys()) | set(gen_counts.keys()))
    if not categories:
        print("No data found.")
        return

    # Separate models into non-adversarial and adversarial
    non_adv_models = sorted([m for m in all_models if not m.startswith("adversarial-")])
    adv_models = sorted([m for m in all_models if m.startswith("adversarial-")])
    model_order = non_adv_models + adv_models

    # Calculate totals once
    total_base = sum(base_counts.values())
    total_prompts = sum(prompt_counts.values())
    total_by_model = defaultdict(int)
    for category in categories:
        for model in model_order:
            total_by_model[model] += gen_counts.get(category, {}).get(model, 0)

    cat_width = max(16, max(len(c) for c in categories))
    num_width = 6

    # Split models into chunks
    model_chunks = [
        model_order[i : i + models_per_table]
        for i in range(0, len(model_order), models_per_table)
    ]
    if not model_chunks:
        model_chunks = [[]]

    for chunk_idx, models in enumerate(model_chunks):
        # Header - include Base/Prompts only in first table
        if chunk_idx == 0:
            header = f"{'Category':<{cat_width}} | {'Base':>{num_width}} | {'Prompts':>{num_width}}"
        else:
            header = f"{'Category':<{cat_width}}"
        for model in models:
            header += f" | {model}"

        if chunk_idx == 0:
            print(f"\n{'=' * len(header)}")
            print("DATASET SUMMARY (image counts)")
        else:
            print(f"\n{'=' * len(header)}")
            print("DATASET SUMMARY (continued)")
        print(f"{'=' * len(header)}")
        print(header)
        print("-" * len(header))

        for category in categories:
            if chunk_idx == 0:
                base = base_counts.get(category, 0)
                prompts = prompt_counts.get(category, 0)
                row = f"{category:<{cat_width}} | {base:>{num_width}} | {prompts:>{num_width}}"
            else:
                row = f"{category:<{cat_width}}"
            for model in models:
                count = gen_counts.get(category, {}).get(model, 0)
                model_width = len(model)
                row += (
                    f" | {count:>{model_width}}"
                    if count > 0
                    else f" | {'-':>{model_width}}"
                )
            print(row)

        # Totals row
        print("-" * len(header))
        if chunk_idx == 0:
            row = f"{'TOTAL':<{cat_width}} | {total_base:>{num_width}} | {total_prompts:>{num_width}}"
        else:
            row = f"{'TOTAL':<{cat_width}}"
        for model in models:
            count = total_by_model[model]
            model_width = len(model)
            row += (
                f" | {count:>{model_width}}"
                if count > 0
                else f" | {'-':>{model_width}}"
            )
        print(row)


def get_detection_rates(group: int = None, ignore_detectors: list = None):
    """Get detection rates grouped by category, generator, and detector."""

    # Initialize the `norm det` table if it doesn't exist, fusing `truthscan` and `truthscan-noc2pa`
    # detector headers.
    conn = _connect()

    # Get all detectors
    detectors = _get_detectors(conn, ignore=ignore_detectors)

    # Get generated image detection rates
    # Join: detection_results -> image_refs -> generated_images -> images
    generated_query = """
        SELECT
            i.category,
            g.model,
            d.detector,
            COUNT(*) as total,
            SUM(CASE WHEN d.ai_score >= ? THEN 1 ELSE 0 END) as detected
        FROM norm_det d
        JOIN generated_images g ON d.image_ref_id = g.image_ref_id
        JOIN images i ON g.base_image_id = i.id
    """
    gen_params = [THRESHOLD]
    if group is not None:
        generated_query += ' WHERE i."group" = ?'
        gen_params.append(group)
    generated_query += " GROUP BY i.category, g.model, d.detector"

    # Get base image detection rates (false positives)
    # Join: detection_results -> image_refs -> images
    base_query = """
        SELECT
            i.category,
            'base' as generator,
            d.detector,
            COUNT(*) as total,
            SUM(CASE WHEN d.ai_score >= ? THEN 1 ELSE 0 END) as detected
        FROM norm_det d
        JOIN images i ON d.image_ref_id = i.image_ref_id
    """
    base_params = [THRESHOLD]
    if group is not None:
        base_query += ' WHERE i."group" = ?'
        base_params.append(group)
    base_query += " GROUP BY i.category, d.detector"

    # Collect all data
    data = defaultdict(lambda: defaultdict(dict))

    for row in conn.execute(generated_query, gen_params).fetchall():
        category, model, detector, total, detected = row
        model = normalize_model(model)
        # detector already canonical via norm_det
        existing = data[category][model].get(detector)
        if existing:
            total += existing["count"]
            detected += existing["detected"]
        rate = detected / total if total > 0 else 0
        data[category][model][detector] = {
            "rate": rate,
            "count": total,
            "detected": detected,
        }

    for row in conn.execute(base_query, base_params).fetchall():
        category, model, detector, total, detected = row
        # No need to call `normalize_model` since model is 'base' by construction.

        # detector already canonical via norm_det
        rate = detected / total if total > 0 else 0
        data[category][model][detector] = {
            "rate": rate,
            "count": total,
            "detected": detected,
        }

    conn.close()
    return data, detectors


def calculate_auc(y_true, y_scores):
    """Calculate AUC using Mann-Whitney U statistic (handles ties correctly)."""
    y_true = np.array(y_true)
    y_scores = np.array(y_scores)

    if len(y_true) == 0 or len(set(y_true)) < 2:
        return None

    pos_scores = y_scores[y_true == 1]
    neg_scores = y_scores[y_true == 0]

    n_pos = len(pos_scores)
    n_neg = len(neg_scores)

    if n_pos == 0 or n_neg == 0:
        return None

    # Count wins and ties
    wins = 0
    ties = 0
    for p in pos_scores:
        wins += np.sum(neg_scores < p)
        ties += np.sum(neg_scores == p)

    auc = (wins + 0.5 * ties) / (n_pos * n_neg)
    return auc


def get_auc_data(group: int = None, ignore_detectors: list = None):
    """Get AUC for each model/detector combination."""
    conn = _connect()

    # Get all detectors
    detectors = _get_detectors(conn, ignore=ignore_detectors)

    # Get base image scores (label=0, should not be detected)
    # Join: detection_results -> images (via image_ref_id)
    base_scores = defaultdict(list)
    base_query = """
        SELECT d.detector, d.ai_score
        FROM norm_det d
        JOIN images i ON d.image_ref_id = i.image_ref_id
    """
    base_params = []
    if group is not None:
        base_query += ' WHERE i."group" = ?'
        base_params.append(group)
    for row in conn.execute(base_query, base_params).fetchall():
        det = row["detector"]
        base_scores[det].append((0, row["ai_score"]))

    # Get generated image scores by model (label=1, should be detected)
    # Join: detection_results -> generated_images (via image_ref_id)
    generated_scores = defaultdict(lambda: defaultdict(list))
    gen_query = """
        SELECT g.model, d.detector, d.ai_score
        FROM norm_det d
        JOIN generated_images g ON d.image_ref_id = g.image_ref_id
        JOIN images i ON g.base_image_id = i.id
    """
    gen_params = []
    if group is not None:
        gen_query += ' WHERE i."group" = ?'
        gen_params.append(group)
    for row in conn.execute(gen_query, gen_params).fetchall():
        model = normalize_model(row["model"])
        det = row["detector"]
        generated_scores[model][det].append((1, row["ai_score"]))

    conn.close()

    # Calculate AUC for each model/detector
    auc_data = defaultdict(dict)
    all_models = set(generated_scores.keys())

    for model in all_models:
        for detector in detectors:
            if detector not in generated_scores[model]:
                continue
            if detector not in base_scores:
                continue

            # Combine base (negative) and generated (positive) scores
            all_labels = []
            all_scores = []

            for label, score in base_scores[detector]:
                all_labels.append(label)
                all_scores.append(score)

            for label, score in generated_scores[model][detector]:
                all_labels.append(label)
                all_scores.append(score)

            auc = calculate_auc(all_labels, all_scores)
            if auc is not None:
                auc_data[model][detector] = auc

    return auc_data, detectors, all_models


# Map from canonical model name to the short generator name used in adversarial naming
_MODEL_TO_GENERATOR = {
    "gemini-3-pro-image-preview": "gemini",
    "gpt-image-1.5": "openai",
    "grok-imagine-image-beta": "grok",
    "qwen-image-2512": "qwen",
    "seedream-v4.5-1k": "seedream",
}


def _find_adversarial_models(gen_scores, categories):
    """Find MAIN_MODELS that have adversarial data in gen_scores.

    Returns list of (source_model, gen_short, display_name) tuples.
    """
    result = []
    for source_model in MAIN_MODELS:
        gen_short = _MODEL_TO_GENERATOR.get(source_model)
        if not gen_short:
            continue
        has_adv = any(
            f"adversarial-{v}-{gen_short}" in gen_scores.get(cat, {})
            for v in _ADV_VERSIONS
            for cat in categories
        )
        if not has_adv:
            continue
        display = MODEL_DISPLAY_NAMES.get(source_model, source_model)
        result.append((source_model, gen_short, display))
    return result


def print_adversarial_auc(group: int = None, ignore_detectors: list = None):
    """Print AUC comparison: original vs adversarial (mild/strong) for each model."""
    auc_data, detectors, all_models = get_auc_data(
        group=group, ignore_detectors=ignore_detectors
    )

    # Find source models that have adversarial data
    groups = []  # list of (display_name, [(variant_label, model_key), ...])
    for model in MAIN_MODELS:
        gen = _MODEL_TO_GENERATOR.get(model)
        if not gen:
            continue

        adv_models = {}
        for v in _ADV_VERSIONS:
            adv_key = f"adversarial-{v}-{gen}"
            if adv_key in auc_data:
                adv_models[v] = adv_key

        if not adv_models:
            continue

        display = MODEL_DISPLAY_NAMES.get(model, model)
        variants = [(f"{display} (base)", model)]
        for v in _ADV_VERSIONS:
            if v in adv_models:
                variants.append((f"{display} ({_VERSION_DISPLAY[v]})", adv_models[v]))
        groups.append((display, variants))

    if not groups:
        print("\nNo adversarial AUC data available.")
        return

    # Print table
    label_width = 24
    det_width = 12

    header = f"{'Model':<{label_width}}"
    for det in detectors:
        det_display = DETECTOR_DISPLAY_NAMES.get(det, det)
        header += f" | {det_display:^{det_width}}"
    header += f" | {'Avg':^{det_width}}"

    separator = "-" * len(header)

    print(f"\n{'=' * len(header)}")
    print("ADVERSARIAL AUC COMPARISON (base vs mild/strong)")
    print(f"{'=' * len(header)}")
    print(header)
    print(separator)

    for display, variants in groups:
        for label, model_key in variants:
            row = f"{label:<{label_width}}"
            aucs = []
            for det in detectors:
                auc = auc_data.get(model_key, {}).get(det)
                if auc is not None:
                    row += f" | {auc:^{det_width}.3f}"
                    aucs.append(auc)
                else:
                    row += f" | {'--':^{det_width}}"
            # Average column
            if aucs:
                avg = sum(aucs) / len(aucs)
                row += f" | {avg:^{det_width}.3f}"
            else:
                row += f" | {'--':^{det_width}}"
            print(row)
        print(separator)

    print("\n* AUC = 1.0 means perfect separation; 0.5 = random guessing")
    print("* Lower AUC on adversarial variants = successful evasion")


def plot_aggregate_roc(
    output_path: str = "tmp/aggregate_roc.png",
    group: int = None,
    ignore_detectors: list = None,
):
    """Plot ROC curves using aggregate detector scores (avg, min, max).

    For each image, computes the average, minimum, and maximum detector score
    across all detectors. Plots base/mild/strong ROC curves for Grok using each
    aggregate, one subplot per aggregate function.
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib not installed. Install with: pip install matplotlib")
        return

    conn = _connect()

    detectors = _get_detectors(conn, ignore=ignore_detectors)
    det_set = set(detectors)

    # Base image scores: image_ref_id -> {detector: score}
    base_by_image = defaultdict(dict)
    base_query = """
        SELECT i.image_ref_id, d.detector, d.ai_score
        FROM norm_det d
        JOIN images i ON d.image_ref_id = i.image_ref_id
    """
    base_params = []
    if group is not None:
        base_query += ' WHERE i."group" = ?'
        base_params.append(group)
    for row in conn.execute(base_query, base_params).fetchall():
        det = row["detector"]
        if det not in det_set:
            continue
        base_by_image[row["image_ref_id"]][det] = row["ai_score"]

    # Generated image scores: model -> image_ref_id -> {detector: score}
    gen_by_image = defaultdict(lambda: defaultdict(dict))
    gen_query = """
        SELECT g.model, g.image_ref_id, d.detector, d.ai_score
        FROM norm_det d
        JOIN generated_images g ON d.image_ref_id = g.image_ref_id
        JOIN images i ON g.base_image_id = i.id
    """
    gen_params = []
    if group is not None:
        gen_query += ' WHERE i."group" = ?'
        gen_params.append(group)
    for row in conn.execute(gen_query, gen_params).fetchall():
        model = normalize_model(row["model"])
        det = row["detector"]
        if det not in det_set:
            continue
        gen_by_image[model][row["image_ref_id"]][det] = row["ai_score"]

    conn.close()

    agg_funcs = [
        ("Average", np.mean),
        ("Minimum", np.min),
        ("Maximum", np.max),
    ]

    # Compute per-image aggregate scores for base images
    # Only include images that have scores from all detectors
    n_det = len(detectors)
    base_agg = {}  # agg_name -> array of scores
    for agg_name, agg_fn in agg_funcs:
        scores = []
        for img_id, det_scores in base_by_image.items():
            if len(det_scores) >= n_det:
                scores.append(agg_fn(list(det_scores.values())))
        base_agg[agg_name] = np.array(scores)

    # Grok variants: base, mild, strong
    source_model = "grok-imagine-image-beta"
    gen_short = "grok"
    variant_keys = [("Base", source_model)]
    for v in _ADV_VERSIONS:
        variant_keys.append((_VERSION_DISPLAY[v], f"adversarial-{v}-{gen_short}"))

    # Compute per-image aggregate scores for each variant
    gen_agg = {}  # (variant_label, agg_name) -> array of scores
    for variant_label, model_key in variant_keys:
        if model_key not in gen_by_image:
            continue
        for agg_name, agg_fn in agg_funcs:
            scores = []
            for _, det_scores in gen_by_image[model_key].items():
                if len(det_scores) >= n_det:
                    scores.append(agg_fn(list(det_scores.values())))
            gen_agg[(variant_label, agg_name)] = np.array(scores)

    # Plot: one subplot per aggregate function
    variant_colors = {"Base": "#1f77b4", "Mild": "#2ca02c", "Strong": "#ff7f0e"}
    n_cols = len(agg_funcs)
    with plt.rc_context(_PLOT_RC):
        fig, axes = plt.subplots(1, n_cols, figsize=(5.5 * n_cols, 5), squeeze=False)

        for col_idx, (agg_name, _) in enumerate(agg_funcs):
            ax = axes[0][col_idx]
            ax.set_title(f"{agg_name} Detector Score", fontsize=12, fontweight="bold")
            ax.plot([0, 1], [0, 1], "k--", alpha=0.3)

            base = base_agg[agg_name]

            for variant_label, model_key in variant_keys:
                key = (variant_label, agg_name)
                if key not in gen_agg or len(gen_agg[key]) == 0:
                    continue
                gen = gen_agg[key]

                thresholds = np.unique(
                    np.concatenate([[-1e-9], base, gen, [1.0 + 1e-9]])
                )
                tprs = []
                fprs = []
                for thresh in thresholds:
                    tp = np.sum(gen >= thresh)
                    fn = np.sum(gen < thresh)
                    fp = np.sum(base >= thresh)
                    tn = np.sum(base < thresh)
                    tprs.append(tp / (tp + fn) if (tp + fn) > 0 else 0)
                    fprs.append(fp / (fp + tn) if (fp + tn) > 0 else 0)

                all_scores = np.concatenate([base, gen])
                all_labels = np.concatenate([np.zeros(len(base)), np.ones(len(gen))])
                auc = calculate_auc(all_labels.tolist(), all_scores.tolist())

                label = f"{variant_label} (AUC={auc:.3f})" if auc else variant_label
                color = variant_colors.get(variant_label, "gray")
                ax.plot(fprs, tprs, color=color, label=label, linewidth=2)

            ax.legend(loc="lower right", fontsize=10)
            ax.set_xlim([0, 1])
            ax.set_ylim([0, 1])
            ax.set_xlabel("False Positive Rate", fontsize=10)
            ax.set_ylabel("True Positive Rate", fontsize=10)
            ax.grid(True, alpha=0.3)

        fig.suptitle(
            "Grok ROC Curves: Aggregate Detector Scores",
            fontsize=14,
            fontweight="bold",
        )
        fig.tight_layout()
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
    print(f"Aggregate ROC curves saved to {output_path}")


def get_adversarial_comparison(
    group: int = None, by_category: bool = False, ignore_detectors: list = None
):
    """Get adversarial vs original comparison data by generator+model.

    Args:
        group: Filter by group
        by_category: If True, also return per-category breakdowns
        ignore_detectors: Detectors to exclude

    Returns:
        Tuple of (comparisons, detectors) or (comparisons, category_comparisons, detectors)
    """
    conn = _connect()

    # Get all detectors
    detectors = _get_detectors(conn, ignore=ignore_detectors)

    # Query to get detection rates by generator+model (overall)
    query = """
        SELECT
            g.generator,
            g.model,
            d.detector,
            COUNT(*) as total,
            SUM(CASE WHEN d.ai_score >= ? THEN 1 ELSE 0 END) as detected
        FROM norm_det d
        JOIN generated_images g ON d.image_ref_id = g.image_ref_id
        JOIN images i ON g.base_image_id = i.id
    """
    params = [THRESHOLD]
    if group is not None:
        query += ' WHERE i."group" = ?'
        params.append(group)
    query += " GROUP BY g.generator, g.model, d.detector"

    # Collect data keyed by (generator, model)
    data = defaultdict(lambda: defaultdict(dict))
    for row in conn.execute(query, params).fetchall():
        generator, model, detector, total, detected = row
        model = normalize_model(model)
        # detector already canonical via norm_det
        existing = data[(generator, model)].get(detector)
        if existing:
            total += existing["count"]
            detected += existing["detected"]
        rate = detected / total if total > 0 else 0
        data[(generator, model)][detector] = {
            "rate": rate,
            "count": total,
            "detected": detected,
        }

    # If by_category, also query per-category data
    category_data = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))
    if by_category:
        cat_query = """
            SELECT
                i.category,
                g.generator,
                g.model,
                d.detector,
                COUNT(*) as total,
                SUM(CASE WHEN d.ai_score >= ? THEN 1 ELSE 0 END) as detected
            FROM norm_det d
            JOIN generated_images g ON d.image_ref_id = g.image_ref_id
            JOIN images i ON g.base_image_id = i.id
        """
        cat_params = [THRESHOLD]
        if group is not None:
            cat_query += ' WHERE i."group" = ?'
            cat_params.append(group)
        cat_query += " GROUP BY i.category, g.generator, g.model, d.detector"

        for row in conn.execute(cat_query, cat_params).fetchall():
            category, generator, model, detector, total, detected = row
            model = normalize_model(model)
            # detector already canonical via norm_det
            existing = category_data[category][(generator, model)].get(detector)
            if existing:
                total += existing["count"]
                detected += existing["detected"]
            rate = detected / total if total > 0 else 0
            category_data[category][(generator, model)][detector] = {
                "rate": rate,
                "count": total,
                "detected": detected,
            }

    conn.close()

    def build_comparisons(data_dict):
        """Build comparison list from data dictionary."""
        comparisons = []
        for (generator, model), det_data in data_dict.items():
            if generator == "adversarial" and model.startswith("adversarial-"):
                source_gen = model.replace("adversarial-", "", 1)

                # Find the source model in data
                source_entry = None
                for (gen, mod), sdata in data_dict.items():
                    if gen == source_gen:
                        source_entry = (gen, mod, sdata)
                        break

                if source_entry:
                    comparisons.append(
                        {
                            "source_generator": source_entry[0],
                            "source_model": source_entry[1],
                            "source_data": source_entry[2],
                            "adversarial_model": model,
                            "adversarial_data": det_data,
                        }
                    )
        return comparisons

    comparisons = build_comparisons(data)

    if by_category:
        cat_comparisons = {}
        for category, cat_data_dict in category_data.items():
            cat_comparisons[category] = build_comparisons(cat_data_dict)
        return comparisons, cat_comparisons, detectors

    return comparisons, detectors


def _print_comparison_rows(comparisons, detectors, model_width=26, det_width=24):
    """Helper to print comparison rows (no header)."""
    for comp in sorted(comparisons, key=lambda x: x["source_generator"]):
        label = comp["source_generator"]
        row = f"{label:<{model_width}}"

        for det in detectors:
            orig_data = comp["source_data"].get(det)
            adv_data = comp["adversarial_data"].get(det)

            if orig_data and adv_data:
                orig_rate = orig_data["rate"]
                adv_rate = adv_data["rate"]
                delta = adv_rate - orig_rate
                cell = f"{orig_rate:>5.1%} → {adv_rate:>5.1%} ({delta:>+5.1%})"
                row += f" | {cell:^{det_width}}"
            elif orig_data:
                cell = f"{orig_data['rate']:>5.1%} → {'--':>5}"
                row += f" | {cell:^{det_width}}"
            elif adv_data:
                cell = f"{'--':>5} → {adv_data['rate']:>5.1%}"
                row += f" | {cell:^{det_width}}"
            else:
                row += f" | {'--':^{det_width}}"
        print(row)


def print_adversarial_comparison(group: int = None, ignore_detectors: list = None):
    """Print adversarial vs original comparison."""
    comparisons, cat_comparisons, detectors = get_adversarial_comparison(
        group=group, by_category=True, ignore_detectors=ignore_detectors
    )

    if not comparisons:
        print("\nNo adversarial comparisons available.")
        return

    model_width = 26
    det_width = 24

    # Build header
    header = f"{'Category/Source':<{model_width}}"
    for det in detectors:
        header += f" | {det:^{det_width}}"
    separator = "-" * len(header)

    # Print section header
    print(f"\n{'=' * len(header)}")
    print("ADVERSARIAL COMPARISON (Original → Adversarial)")
    print(f"{'=' * len(header)}")
    print(header)
    print(separator)

    # Print per-category comparisons compactly
    first_category = True
    for category in sorted(cat_comparisons.keys()):
        cat_comps = cat_comparisons[category]
        if not cat_comps:
            continue

        # Add blank line between categories (but not before first)
        if not first_category:
            print()
        first_category = False

        # Category label row
        print(f"{category.upper():<{model_width}}" + " |" * len(detectors))
        _print_comparison_rows(cat_comps, detectors, model_width, det_width)

    # Print overall
    print()
    print(separator)
    print(f"{'OVERALL':<{model_width}}" + " |" * len(detectors))
    _print_comparison_rows(comparisons, detectors, model_width, det_width)

    print("\n* Negative delta = adversarial is LESS detected (evasion success)")
    print("* Positive delta = adversarial is MORE detected")


def print_report(group: int = None, ignore_detectors: list = None):
    """Print formatted detection rate report."""
    data, detectors = get_detection_rates(
        group=group, ignore_detectors=ignore_detectors
    )

    # Get all models from data and define display order
    all_models = set()
    for category in data:
        all_models.update(data[category].keys())

    # Define display order: base first, then non-adversarial sorted, then adversarial sorted
    non_adv_models = sorted(
        [m for m in all_models if m != "base" and not m.startswith("adversarial-")]
    )
    adv_models = sorted([m for m in all_models if m.startswith("adversarial-")])
    model_order = ["base"] + non_adv_models + adv_models

    categories = sorted(data.keys())

    # Calculate column widths
    cat_width = max(24, max((len(c) for c in categories), default=12))
    det_width = 18

    # Print header template
    header = f"{'Category':<{cat_width}}"
    for det in detectors:
        header += f" | {det:^{det_width}}"

    separator = "-" * len(header)

    # Track totals
    overall_totals = defaultdict(lambda: {"detected": 0, "count": 0})

    # Reorganize data by model -> category
    data_by_model = defaultdict(lambda: defaultdict(dict))
    for category in data:
        for model in data[category]:
            for det in data[category][model]:
                data_by_model[model][category][det] = data[category][model][det]

    for model in model_order:
        if model not in data_by_model:
            continue

        print(f"\n{'=' * len(header)}")
        print(f"MODEL: {model}")
        print(f"{'=' * len(header)}")
        print(header)
        print(separator)

        model_totals = defaultdict(lambda: {"detected": 0, "count": 0})

        for category in categories:
            if category not in data_by_model[model]:
                continue

            row = f"{category:<{cat_width}}"
            for det in detectors:
                if det in data_by_model[model][category]:
                    d = data_by_model[model][category][det]
                    rate = d["rate"]
                    count = d["count"]
                    detected = d["detected"]
                    row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"

                    # Add to totals
                    model_totals[det]["detected"] += detected
                    model_totals[det]["count"] += count
                    overall_totals[det]["detected"] += detected
                    overall_totals[det]["count"] += count
                else:
                    row += f" | {'--':^18}"
            print(row)

        # Model totals
        print(separator)
        row = f"{'TOTAL':<{cat_width}}"
        for det in detectors:
            if model_totals[det]["count"] > 0:
                rate = model_totals[det]["detected"] / model_totals[det]["count"]
                count = model_totals[det]["count"]
                detected = model_totals[det]["detected"]
                row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
            else:
                row += f" | {'--':^18}"
        print(row)

        # Add base (FPR) reference row for non-base models
        if model != "base" and "base" in data_by_model:
            base_totals = defaultdict(lambda: {"detected": 0, "count": 0})
            for category in categories:
                if category not in data_by_model["base"]:
                    continue
                for det in detectors:
                    if det in data_by_model["base"][category]:
                        d = data_by_model["base"][category][det]
                        base_totals[det]["detected"] += d["detected"]
                        base_totals[det]["count"] += d["count"]
            row = f"{'base (FPR)':<{cat_width}}"
            for det in detectors:
                if base_totals[det]["count"] > 0:
                    rate = base_totals[det]["detected"] / base_totals[det]["count"]
                    count = base_totals[det]["count"]
                    detected = base_totals[det]["detected"]
                    row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
                else:
                    row += f" | {'--':^18}"
            print(row)

    # Overall totals
    model_width = max(24, max((len(m) for m in all_models), default=12))
    totals_header = f"{'Model':<{model_width}}"
    for det in detectors:
        totals_header += f" | {det:^{det_width}}"
    # Add avg/max columns
    totals_header += f" | {'Avg':^6} | {'Max':^6}"
    totals_separator = "-" * len(totals_header)

    print(f"\n{'=' * len(totals_header)}")
    print("OVERALL TOTALS")
    print(f"{'=' * len(totals_header)}")
    print(totals_header)
    print(totals_separator)

    # Calculate overall by model
    overall_by_model = defaultdict(
        lambda: defaultdict(lambda: {"detected": 0, "count": 0})
    )
    for category in categories:
        for model in data[category]:
            for det in data[category][model]:
                d = data[category][model][det]
                overall_by_model[model][det]["detected"] += d["detected"]
                overall_by_model[model][det]["count"] += d["count"]

    # Track totals separately
    non_adv_totals = defaultdict(lambda: {"detected": 0, "count": 0})
    adv_totals = defaultdict(lambda: {"detected": 0, "count": 0})

    # Print non-adversarial models (base + non-adv)
    for model in ["base"] + non_adv_models:
        if model not in overall_by_model:
            continue
        row = f"{model:<{model_width}}"
        model_rates = []
        for det in detectors:
            if (
                det in overall_by_model[model]
                and overall_by_model[model][det]["count"] > 0
            ):
                d = overall_by_model[model][det]
                rate = d["detected"] / d["count"]
                count = d["count"]
                detected = d["detected"]
                row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
                model_rates.append(rate)
                non_adv_totals[det]["detected"] += detected
                non_adv_totals[det]["count"] += count
            else:
                row += f" | {'--':^18}"
        # Add avg/max for this model
        if model_rates:
            avg_rate = sum(model_rates) / len(model_rates)
            max_rate = max(model_rates)
            row += f" | {avg_rate:>6.1%} | {max_rate:>6.1%}"
        else:
            row += f" | {'--':^8} | {'--':^8}"
        print(row)

    # Print non-adversarial subtotal
    print(totals_separator)
    row = f"{'SUBTOTAL (non-adv)':<{model_width}}"
    subtotal_rates = []
    for det in detectors:
        if non_adv_totals[det]["count"] > 0:
            rate = non_adv_totals[det]["detected"] / non_adv_totals[det]["count"]
            count = non_adv_totals[det]["count"]
            detected = non_adv_totals[det]["detected"]
            row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
            subtotal_rates.append(rate)
        else:
            row += f" | {'--':^18}"
    if subtotal_rates:
        row += f" | {sum(subtotal_rates) / len(subtotal_rates):>6.1%} | {max(subtotal_rates):>6.1%}"
    else:
        row += f" | {'--':^8} | {'--':^8}"
    print(row)

    # Print adversarial models if any exist
    if adv_models:
        print(totals_separator)
        for model in adv_models:
            if model not in overall_by_model:
                continue
            row = f"{model:<{model_width}}"
            model_rates = []
            for det in detectors:
                if (
                    det in overall_by_model[model]
                    and overall_by_model[model][det]["count"] > 0
                ):
                    d = overall_by_model[model][det]
                    rate = d["detected"] / d["count"]
                    count = d["count"]
                    detected = d["detected"]
                    row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
                    model_rates.append(rate)
                    adv_totals[det]["detected"] += detected
                    adv_totals[det]["count"] += count
                else:
                    row += f" | {'--':^18}"
            if model_rates:
                row += f" | {sum(model_rates) / len(model_rates):>6.1%} | {max(model_rates):>6.1%}"
            else:
                row += f" | {'--':^8} | {'--':^8}"
            print(row)

        # Print adversarial subtotal
        print(totals_separator)
        row = f"{'SUBTOTAL (adversarial)':<{model_width}}"
        subtotal_rates = []
        for det in detectors:
            if adv_totals[det]["count"] > 0:
                rate = adv_totals[det]["detected"] / adv_totals[det]["count"]
                count = adv_totals[det]["count"]
                detected = adv_totals[det]["detected"]
                row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
                subtotal_rates.append(rate)
            else:
                row += f" | {'--':^18}"
        if subtotal_rates:
            row += f" | {sum(subtotal_rates) / len(subtotal_rates):>6.1%} | {max(subtotal_rates):>6.1%}"
        else:
            row += f" | {'--':^8} | {'--':^8}"
        print(row)

    # Print grand total
    print(totals_separator)
    row = f"{'TOTAL':<{model_width}}"
    total_rates = []
    for det in detectors:
        if overall_totals[det]["count"] > 0:
            rate = overall_totals[det]["detected"] / overall_totals[det]["count"]
            count = overall_totals[det]["count"]
            detected = overall_totals[det]["detected"]
            row += f" | {rate:>6.1%} ({detected:>4}/{count:<4})"
            total_rates.append(rate)
        else:
            row += f" | {'--':^18}"
    if total_rates:
        row += (
            f" | {sum(total_rates) / len(total_rates):>6.1%} | {max(total_rates):>6.1%}"
        )
    else:
        row += f" | {'--':^8} | {'--':^8}"
    print(row)

    print(f"\n* Detection threshold: {THRESHOLD:.0%}")
    print("* Rate = detected/total at threshold")
    print("* 'base' row shows false positive rate (FPR)")


def get_raw_scores(group=None, ignore_detectors: list = None):
    """Get raw detection scores for base and generated images."""
    conn = _connect()

    # Base image scores per (category, detector)
    base_query = """
        SELECT i.category, d.detector, d.ai_score
        FROM norm_det d
        JOIN images i ON d.image_ref_id = i.image_ref_id
    """
    base_params = []
    if group is not None:
        base_query += ' WHERE i."group" = ?'
        base_params.append(group)

    base_scores = defaultdict(lambda: defaultdict(list))
    for row in conn.execute(base_query, base_params).fetchall():
        det = row["detector"]
        base_scores[row["category"]][det].append(row["ai_score"])

    # Generated image scores per (category, model, detector)
    gen_query = """
        SELECT i.category, g.model, d.detector, d.ai_score
        FROM norm_det d
        JOIN generated_images g ON d.image_ref_id = g.image_ref_id
        JOIN images i ON g.base_image_id = i.id
    """
    gen_params = []
    if group is not None:
        gen_query += ' WHERE i."group" = ?'
        gen_params.append(group)

    gen_scores = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    for row in conn.execute(gen_query, gen_params).fetchall():
        model = normalize_model(row["model"])
        det = row["detector"]
        gen_scores[row["category"]][model][det].append(row["ai_score"])

    # Get detector list
    detectors = _get_detectors(conn, ignore=ignore_detectors)

    conn.close()

    return base_scores, gen_scores, detectors


def _fmt_pct(value):
    """Format a 0-1 value as a percentage string with 1 decimal, or empty if None."""
    if value is None:
        return ""
    return f"{value * 100:.1f}"


def print_latex_by_class(group=None, ignore_detectors: list = None):
    """Print LaTeX table for AI detection results by image class."""
    base_scores, gen_scores, detectors = get_raw_scores(
        group, ignore_detectors=ignore_detectors
    )

    categories = sorted(set(base_scores.keys()) | set(gen_scores.keys()))

    # Compute metrics per (category, detector) using only the 5 main models
    metrics = {}
    for category in categories:
        for detector in detectors:
            base = base_scores[category][detector]

            # Combine scores from the 5 main models only
            gen = []
            for model in MAIN_MODELS:
                gen.extend(gen_scores[category][model][detector])

            base_arr = np.array(base) if base else np.array([])
            gen_arr = np.array(gen) if gen else np.array([])

            tp = int(np.sum(gen_arr >= THRESHOLD)) if len(gen_arr) > 0 else 0
            fn = int(np.sum(gen_arr < THRESHOLD)) if len(gen_arr) > 0 else 0
            tn = int(np.sum(base_arr < THRESHOLD)) if len(base_arr) > 0 else 0
            fp = int(np.sum(base_arr >= THRESHOLD)) if len(base_arr) > 0 else 0

            tpr = tp / (tp + fn) if (tp + fn) > 0 else None
            tnr = tn / (tn + fp) if (tn + fp) > 0 else None

            total = tp + tn + fp + fn
            acc = (tp + tn) / total if total > 0 else None

            # AUC
            if len(base_arr) > 0 and len(gen_arr) > 0:
                all_labels = [0] * len(base_arr) + [1] * len(gen_arr)
                all_scores_list = list(base_arr) + list(gen_arr)
                auc = calculate_auc(all_labels, all_scores_list)
            else:
                auc = None

            metrics[(category, detector)] = {
                "tpr": tpr,
                "tnr": tnr,
                "auc": auc,
                "acc": acc,
            }

    # Split categories into groups of 4
    cat_chunks = [categories[i : i + 4] for i in range(0, len(categories), 4)]

    lines = []
    lines.append("\\begin{table*}[htbp]")
    lines.append("\\centering")
    lines.append("")

    for chunk_idx, cats in enumerate(cat_chunks):
        if chunk_idx > 0:
            lines.append("")
            lines.append("\\vspace{1em}")
            lines.append("")

        ncats = len(cats)
        ncols = 1 + 4 * ncats

        # Column spec
        col_parts = ["!{\\vrule width 1.2pt}l"]
        for i in range(ncats):
            if i < ncats - 1:
                col_parts.append("|cccc")
            else:
                col_parts.append("|cccc!{\\vrule width 1.2pt}")
        col_spec = "".join(col_parts)

        lines.append("\\begin{adjustbox}{max width=\\textwidth, center}")
        lines.append(f"\\begin{{tabular}}{{{col_spec}}}")
        lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")

        # Category headers
        cat_header_parts = []
        for i, cat in enumerate(cats):
            escaped = cat.replace("_", "\\_")
            if i < ncats - 1:
                cat_header_parts.append(
                    f"\\multicolumn{{4}}{{c|}}{{\\texttt{{{escaped}}}}}"
                )
            else:
                cat_header_parts.append(
                    f"\\multicolumn{{4}}{{c!{{\\vrule width 1.2pt}}}}{{\\texttt{{{escaped}}}}}"
                )
        lines.append("& " + " \n& ".join(cat_header_parts) + " \\\\")
        lines.append(f"\\cline{{2-{ncols}}}")

        # Sub-header row
        sub = "\\textbf{Detector} "
        for _ in cats:
            sub += "\n& TPR & TNR & AUC & Acc "
        sub += "\\\\"
        lines.append(sub)
        lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")

        # Data rows
        for det in detectors:
            det_display = DETECTOR_DISPLAY_NAMES.get(det, det)
            row = f"{det_display:<15}"
            for cat in cats:
                m = metrics.get((cat, det))
                if m:
                    row += (
                        f" & {_fmt_pct(m['tpr'])}"
                        f" & {_fmt_pct(m['tnr'])}"
                        f" & {_fmt_pct(m['auc'])}"
                        f" & {_fmt_pct(m['acc'])}"
                    )
                else:
                    row += " &  &  &  & "
            row += " \\\\"
            lines.append(row)

        lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")
        lines.append("\\end{tabular}")
        lines.append("\\end{adjustbox}")

    lines.append("")
    lines.append("\\caption{AI Detection Performance Across Image Classes}")
    lines.append("\\label{tab:ai_detection_results_by_class}")
    lines.append("\\end{table*}")

    print("\n".join(lines))


def print_latex_aggregate(group=None, ignore_detectors: list = None):
    """Print LaTeX table for AI detection results aggregated across models."""
    base_scores, gen_scores, detectors = get_raw_scores(
        group, ignore_detectors=ignore_detectors
    )

    categories = sorted(set(base_scores.keys()) | set(gen_scores.keys()))

    # Compute TNR per detector (overall, across all categories)
    tnr_by_detector = {}
    for detector in detectors:
        all_base = []
        for category in categories:
            all_base.extend(base_scores[category][detector])
        if all_base:
            base_arr = np.array(all_base)
            tn = int(np.sum(base_arr < THRESHOLD))
            total = len(base_arr)
            tnr_by_detector[detector] = tn / total if total > 0 else None
        else:
            tnr_by_detector[detector] = None

    # Compute TPR per (model, detector) overall across all categories
    tpr_by_model_detector = defaultdict(dict)
    for model in MAIN_MODELS:
        for detector in detectors:
            all_gen = []
            for category in categories:
                all_gen.extend(gen_scores[category][model][detector])
            if all_gen:
                gen_arr = np.array(all_gen)
                tp = int(np.sum(gen_arr >= THRESHOLD))
                total = len(gen_arr)
                tpr_by_model_detector[model][detector] = (
                    tp / total if total > 0 else None
                )
            else:
                tpr_by_model_detector[model][detector] = None

    # Format LaTeX
    n_models = len(MAIN_MODELS)

    lines = []
    lines.append("\\begin{table*}[htbp]")
    lines.append("\\centering")
    lines.append("")
    lines.append("\\begin{adjustbox}{max width=\\textwidth, center}")
    lines.append("")

    col_spec = (
        f"!{{\\vrule width 1.2pt}}l|c|"
        f"{'c' * n_models}"
        f"!{{\\vrule width 1.2pt}}c!{{\\vrule width 1.2pt}}"
    )
    lines.append(f"\\begin{{tabular}}{{{col_spec}}}")
    lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")

    # Multi-column headers
    lines.append(
        f"& \\multicolumn{{1}}{{c|}}{{}} \n"
        f"& \\multicolumn{{{n_models}}}{{c!{{\\vrule width 1.2pt}}}}"
        f"{{Detection Rate by Model (Percent)}} \n"
        f"& \\multicolumn{{1}}{{c!{{\\vrule width 1.2pt}}}}{{}}\\\\"
    )
    lines.append(f"\\cline{{3-{2 + n_models}}}")

    # Sub-header
    sub = "\\textbf{Detector} \n& True Negative Rate"
    for model in MAIN_MODELS:
        display = MODEL_DISPLAY_NAMES.get(model, model)
        sub += f" & {display}"
    sub += " & Average\\\\"
    lines.append(sub)
    lines.append("")
    lines.append("% \\hline")
    lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")

    # Data rows
    avg_by_model = defaultdict(list)  # model -> list of rates for averaging
    avg_tnr = []

    for det in detectors:
        det_display = DETECTOR_DISPLAY_NAMES.get(det, det)

        tnr = tnr_by_detector.get(det)
        tnr_s = _fmt_pct(tnr)
        if tnr is not None:
            avg_tnr.append(tnr)

        row = f"{det_display:<15} & {tnr_s}"

        # Per-model TPR
        det_rates = []
        for model in MAIN_MODELS:
            tpr = tpr_by_model_detector[model].get(det)
            if tpr is not None:
                row += f" & {_fmt_pct(tpr)}"
                det_rates.append(tpr)
                avg_by_model[model].append(tpr)
            else:
                row += " & "

        # Average across models for this detector
        if det_rates:
            avg = sum(det_rates) / len(det_rates)
            row += f" & {_fmt_pct(avg)}"
        else:
            row += " & "

        row += "\\\\"
        lines.append(row)

    # Average row
    lines.append("\\specialrule{1.2pt}{0pt}{0pt}  % Thick top line")
    avg_row = "Average & "

    # Average TNR
    if avg_tnr:
        avg_row += _fmt_pct(sum(avg_tnr) / len(avg_tnr))

    all_model_avgs = []
    for model in MAIN_MODELS:
        if avg_by_model[model]:
            mavg = sum(avg_by_model[model]) / len(avg_by_model[model])
            avg_row += f" & {_fmt_pct(mavg)}"
            all_model_avgs.append(mavg)
        else:
            avg_row += " & "

    # Grand average
    if all_model_avgs:
        grand_avg = sum(all_model_avgs) / len(all_model_avgs)
        avg_row += f" & {_fmt_pct(grand_avg)}"
    else:
        avg_row += " & "

    avg_row += "\\\\"
    lines.append(avg_row)
    lines.append("\\specialrule{1.2pt}{0pt}{0pt}")
    lines.append("\\end{tabular}")
    lines.append("\\end{adjustbox}")
    lines.append("")
    lines.append("\\caption{AI Detection Performance Across Models}")
    lines.append("\\label{tab:ai_detection_results}")
    lines.append("\\end{table*}")

    print("\n".join(lines))


def print_latex_adversarial_by_class(group=None, ignore_detectors: list = None):
    """Print LaTeX table for adversarial detection rates by image class.

    For each category, shows 6 columns: base/mild/strong detection rates for
    Gemini and Grok (the two models with adversarial data).
    """
    _, gen_scores, detectors = get_raw_scores(group, ignore_detectors=ignore_detectors)

    categories = sorted(gen_scores.keys())

    # The two models with adversarial data
    adv_models = [
        ("gemini-3-pro-image-preview", "gemini", "Gemini"),
        ("grok-imagine-image-beta", "grok", "Grok"),
    ]
    versions = ["base"] + _ADV_VERSIONS
    version_labels = ["Base"] + [_VERSION_DISPLAY[v] for v in _ADV_VERSIONS]

    # Compute detection rate per (category, detector, model_short, version)
    rates = {}
    for category in categories:
        for det in detectors:
            for source_model, gen_short, _ in adv_models:
                # Base detection rate
                base_scores_list = (
                    gen_scores[category].get(source_model, {}).get(det, [])
                )
                if base_scores_list:
                    arr = np.array(base_scores_list)
                    rates[(category, det, gen_short, "base")] = np.mean(
                        arr >= THRESHOLD
                    )

                # Adversarial versions
                for v in _ADV_VERSIONS:
                    adv_key = f"adversarial-{v}-{gen_short}"
                    adv_scores_list = gen_scores[category].get(adv_key, {}).get(det, [])
                    if adv_scores_list:
                        arr = np.array(adv_scores_list)
                        rates[(category, det, gen_short, v)] = np.mean(arr >= THRESHOLD)

    # Split categories into groups of 4
    cat_chunks = [categories[i : i + 4] for i in range(0, len(categories), 4)]
    n_adv = len(adv_models)  # 2
    n_ver = len(versions)  # 3
    cols_per_cat = n_adv * n_ver  # 6

    lines = []
    lines.append("\\begin{table*}[htbp]")
    lines.append("\\centering")
    lines.append("")

    for chunk_idx, cats in enumerate(cat_chunks):
        if chunk_idx > 0:
            lines.append("")
            lines.append("\\vspace{1em}")
            lines.append("")

        ncats = len(cats)
        ncols = 1 + cols_per_cat * ncats

        # Column spec: l | ccc|ccc | ccc|ccc | ...
        col_parts = ["!{\\vrule width 1.2pt}l"]
        for i in range(ncats):
            for j in range(n_adv):
                if i == ncats - 1 and j == n_adv - 1:
                    col_parts.append("|ccc!{\\vrule width 1.2pt}")
                else:
                    col_parts.append("|ccc")
        col_spec = "".join(col_parts)

        lines.append("\\begin{adjustbox}{max width=\\textwidth, center}")
        lines.append(f"\\begin{{tabular}}{{{col_spec}}}")
        lines.append("\\specialrule{1.2pt}{0pt}{0pt}")

        # Top header: category names spanning 6 columns each
        cat_header_parts = []
        for i, cat in enumerate(cats):
            escaped = cat.replace("_", "\\_")
            border = "c!{\\vrule width 1.2pt}" if i == ncats - 1 else "c|"
            cat_header_parts.append(
                f"\\multicolumn{{{cols_per_cat}}}{{{border}}}{{\\texttt{{{escaped}}}}}"
            )
        lines.append("& " + " \n& ".join(cat_header_parts) + " \\\\")

        # Second header: model names spanning 3 columns each
        cline_start = 2
        model_header_parts = []
        for i, cat in enumerate(cats):
            for j, (_, _, display) in enumerate(adv_models):
                is_last = i == ncats - 1 and j == n_adv - 1
                border = "c!{\\vrule width 1.2pt}" if is_last else "c|"
                model_header_parts.append(
                    f"\\multicolumn{{{n_ver}}}{{{border}}}{{{display}}}"
                )
        lines.append("& " + " \n& ".join(model_header_parts) + " \\\\")

        # Clines under model headers
        for i, cat in enumerate(cats):
            for j in range(n_adv):
                start = 2 + i * cols_per_cat + j * n_ver
                end = start + n_ver - 1
                lines.append(f"\\cline{{{start}-{end}}}")

        # Sub-header row: Base Mild Strong repeated
        sub = "\\textbf{Detector}"
        for _ in cats:
            for _ in adv_models:
                for vl in version_labels:
                    sub += f" & {vl}"
        sub += " \\\\"
        lines.append(sub)
        lines.append("\\specialrule{1.2pt}{0pt}{0pt}")

        # Data rows
        for det in detectors:
            det_display = DETECTOR_DISPLAY_NAMES.get(det, det)
            row = f"{det_display:<15}"
            for cat in cats:
                for _, gen_short, _ in adv_models:
                    for v in versions:
                        rate = rates.get((cat, det, gen_short, v))
                        row += f" & {_fmt_pct(rate)}"
            row += " \\\\"
            lines.append(row)

        lines.append("\\specialrule{1.2pt}{0pt}{0pt}")
        lines.append("\\end{tabular}")
        lines.append("\\end{adjustbox}")

    lines.append("")
    lines.append(
        "\\caption{Adversarial Detection Rates by Image Class (Base / Mild / Strong)}"
    )
    lines.append("\\label{tab:adversarial_detection_by_class}")
    lines.append("\\end{table*}")

    print("\n".join(lines))


def get_adversarial_paired_scores(group: int = None):
    """Get matched-pair detector scores for original vs adversarial images.

    Returns dict keyed by (gen_short, version) -> list of
    {category, detector, orig_score, adv_score} dicts.
    """
    conn = _connect()

    results = {}

    for source_model in MAIN_MODELS:
        gen_short = _MODEL_TO_GENERATOR.get(source_model)
        if not gen_short:
            continue

        # Source generators: edit and t2i variants
        source_gens = [gen_short, f"{gen_short}-t2i"]

        for version in _ADV_VERSIONS:
            pairs = []
            for src_gen in source_gens:
                adv_model = f"adversarial-{version}-{src_gen}"

                # Match pairs on same detector, plus merge truthscan variants
                query = """
                    SELECT i.category, d_orig.detector as orig_det,
                           d_adv.detector as adv_det,
                           d_orig.ai_score as orig_score,
                           d_adv.ai_score as adv_score
                    FROM generated_images g_orig
                    JOIN generated_images g_adv
                        ON g_orig.base_image_id = g_adv.base_image_id
                    JOIN images i ON g_orig.base_image_id = i.id
                    JOIN norm_det d_orig
                        ON g_orig.image_ref_id = d_orig.image_ref_id
                    JOIN norm_det d_adv
                        ON g_adv.image_ref_id = d_adv.image_ref_id
                    WHERE g_orig.generator = ?
                      AND g_adv.model = ?
                      AND (d_orig.detector = d_adv.detector
                           OR (d_orig.detector LIKE 'truthscan%'
                               AND d_adv.detector LIKE 'truthscan%'))
                """
                params = [src_gen, adv_model]
                if group is not None:
                    query += ' AND i."group" = ?'
                    params.append(group)

                for row in conn.execute(query, params).fetchall():
                    pairs.append(
                        {
                            "category": row["category"],
                            "detector": row["orig_det"],
                            "orig_score": row["orig_score"],
                            "adv_score": row["adv_score"],
                        }
                    )

            if pairs:
                results[(gen_short, version)] = pairs

    conn.close()
    return results


def compute_adversarial_paired_stats(group: int = None, ignore_detectors: list = None):
    """Compute matched-pairs statistics for adversarial effects on detector scores.

    Returns (rows, detectors, categories) where rows is a list of dicts with keys:
        model, version, category, detector, n, mean_orig, mean_adv,
        mean_reduction, cohens_d, p_value, beta.
    Returns (None, None, None) if no paired data is available.
    """
    from scipy import stats as scipy_stats
    from scipy.stats import nct

    paired_data = get_adversarial_paired_scores(group=group)

    if not paired_data:
        return None, None, None

    detectors = sorted({p["detector"] for pairs in paired_data.values() for p in pairs})
    if ignore_detectors:
        detectors = [d for d in detectors if d not in ignore_detectors]

    categories = sorted(
        {p["category"] for pairs in paired_data.values() for p in pairs}
    )

    alpha = 0.05

    # Collect results
    rows = []
    _ver_order = {v: i for i, v in enumerate(_ADV_VERSIONS)}
    for (gen_short, version), pairs in sorted(
        paired_data.items(), key=lambda kv: (kv[0][0], _ver_order.get(kv[0][1], 99))
    ):
        for category in categories:
            for det in detectors:
                # Filter to this (category, detector) subset
                orig = []
                adv = []
                for p in pairs:
                    if p["category"] == category and p["detector"] == det:
                        orig.append(p["orig_score"])
                        adv.append(p["adv_score"])

                if len(orig) < 2:
                    continue

                orig_arr = np.array(orig)
                adv_arr = np.array(adv)

                n = len(orig_arr)
                mean_orig = np.mean(orig_arr)
                mean_adv = np.mean(adv_arr)
                mean_diff = mean_orig - mean_adv

                cohens_d, p_value = _paired_effect_size(orig_arr, adv_arr)

                # Power and Type II error (beta) for one-sided test
                df = n - 1
                ncp = cohens_d * np.sqrt(n)  # non-centrality parameter
                t_crit = scipy_stats.t.ppf(1 - alpha, df)
                # Power = P(T > t_crit | H1 true)
                power = 1 - nct.cdf(t_crit, df, ncp)
                beta = max(0.0, 1 - power)

                rows.append(
                    {
                        "model": gen_short,
                        "version": version,
                        "category": category,
                        "detector": det,
                        "n": n,
                        "mean_orig": mean_orig,
                        "mean_adv": mean_adv,
                        "mean_reduction": mean_diff,
                        "cohens_d": cohens_d,
                        "p_value": p_value,
                        "beta": beta,
                    }
                )

    return rows, detectors, categories


def print_adversarial_paired_analysis(group: int = None, ignore_detectors: list = None):
    """Print matched-pairs analysis of adversarial effects on detector scores."""
    rows, detectors, categories = compute_adversarial_paired_stats(
        group=group,
        ignore_detectors=ignore_detectors,
    )

    if rows is None:
        print("\nNo adversarial paired data available.")
        return

    if not rows:
        print("\nNo matched pairs found.")
        return

    # Print table grouped by model
    col_w = {
        "cat": max(16, max(len(c) for c in categories)),
        "det": max(14, max(len(d) for d in detectors)),
        "ver": 6,
        "n": 5,
        "score": 8,
        "d": 8,
        "p": 10,
    }

    header = (
        f"{'Category':<{col_w['cat']}} "
        f"{'Detector':<{col_w['det']}} "
        f"{'Ver':>{col_w['ver']}} "
        f"{'N':>{col_w['n']}} "
        f"{'Before':>{col_w['score']}} "
        f"{'After':>{col_w['score']}} "
        f"{'Δ Mean':>{col_w['score']}} "
        f"{'Cohen d':>{col_w['d']}} "
        f"{'p-value':>{col_w['p']}} "
        f"{'β (TypeII)':>{col_w['p']}}"
    )
    separator = "-" * len(header)

    current_model = None
    for row in rows:
        if row["model"] != current_model:
            current_model = row["model"]
            display = MODEL_DISPLAY_NAMES.get(
                next(
                    (m for m, g in _MODEL_TO_GENERATOR.items() if g == current_model),
                    "",
                ),
                current_model,
            )
            print(f"\n{'=' * len(header)}")
            print(f"MODEL: {display}")
            print(f"{'=' * len(header)}")
            print(header)
            print(separator)

        p_str = (
            f"{row['p_value']:.2e}"
            if row["p_value"] < 0.001
            else f"{row['p_value']:.4f}"
        )
        b_str = f"{row['beta']:.2e}" if row["beta"] < 0.001 else f"{row['beta']:.4f}"

        print(
            f"{row['category']:<{col_w['cat']}} "
            f"{row['detector']:<{col_w['det']}} "
            f"{_VERSION_DISPLAY[row['version']]:>{col_w['ver']}} "
            f"{row['n']:>{col_w['n']}} "
            f"{row['mean_orig']:>{col_w['score']}.4f} "
            f"{row['mean_adv']:>{col_w['score']}.4f} "
            f"{row['mean_reduction']:>{col_w['score']}.4f} "
            f"{row['cohens_d']:>{col_w['d']}.3f} "
            f"{p_str:>{col_w['p']}} "
            f"{b_str:>{col_w['p']}}"
        )

    print("\n* Δ Mean = mean(before - after); positive = score decreased (evasion)")
    print("* Cohen's d: |d| > 0.8 = large effect, > 0.5 = medium, > 0.2 = small")
    print(
        "* p-value: Type I error probability (one-sided paired t-test, H1: score decreased, α = 0.05)"
    )
    print("* β: Type II error probability (1 - power)")


def _paired_effect_size(base_scores, adv_scores):
    """Compute paired Cohen's d and one-sided p-value (H1: base > adv).

    Returns (cohens_d, p_value).  Positive d means adversarial reduced detection.
    Requires equal-length arrays (paired observations).
    """
    base = np.asarray(base_scores, dtype=float)
    adv = np.asarray(adv_scores, dtype=float)
    diff = base - adv
    mean_diff = np.mean(diff)
    std_diff = np.std(diff, ddof=1)

    cohens_d = mean_diff / std_diff if std_diff > 0 else float("inf")
    _, p_value = ttest_rel(base, adv, alternative="greater")

    return cohens_d, p_value


def _compute_d_and_p(rows, gen_short, ver, categories, detectors):
    """Build Cohen's d and p-value matrices from pre-computed paired stats rows.

    Returns (d_matrix, p_matrix) of shape (len(categories), len(detectors)).
    Positive d means the base model has higher scores (adversarial reduced detection).
    """
    nc, nd = len(categories), len(detectors)
    d_mat = np.zeros((nc, nd))
    p_mat = np.ones((nc, nd))

    cat_to_idx = {c: i for i, c in enumerate(categories)}
    det_to_idx = {d: j for j, d in enumerate(detectors)}

    for r in rows:
        if r["model"] != gen_short or r["version"] != ver:
            continue
        i = cat_to_idx.get(r["category"])
        j = det_to_idx.get(r["detector"])
        if i is not None and j is not None:
            d_mat[i, j] = r["cohens_d"]
            p_mat[i, j] = r["p_value"]

    return d_mat, p_mat


_PLOT_RC = {
    "font.family": "sans-serif",
    "font.sans-serif": ["Helvetica Neue", "Helvetica", "Arial"],
}


def _style_heatmap_ax(ax):
    """Apply standard heatmap styling: remove spines, tick marks, add cell dividers."""
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.tick_params(axis="both", length=0)
    # Thin white lines between cells
    n_rows, n_cols = ax.images[0].get_array().shape
    ax.set_xticks([x - 0.5 for x in range(1, n_cols)], minor=True)
    ax.set_yticks([y - 0.5 for y in range(1, n_rows)], minor=True)
    ax.grid(which="minor", color="white", linewidth=1.5)
    ax.tick_params(which="minor", length=0)


def _annotate_heatmap(ax, matrix, n_rows, n_cols, fontsize=8, fmt="{:.0f}%"):
    """Annotate heatmap cells with values."""
    for i in range(n_rows):
        for j in range(n_cols):
            val = matrix[i, j]
            color = "white" if val > 55 else "black"
            ax.text(
                j,
                i,
                fmt.format(val),
                ha="center",
                va="center",
                fontsize=fontsize,
                color=color,
                fontweight="bold",
            )


def plot_detection_heatmap(
    output_path: str = "tmp/detection_heatmap.png",
    group: int = None,
    ignore_detectors: list = None,
):
    """Plot detection rate heatmaps: one row per model, columns for Base/Mild/Strong.

    Includes only models that have adversarial data. Each subplot shows detection
    rates per category (+ Overall row) and detector, using a blue colormap.
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.colors import LinearSegmentedColormap
    except ImportError:
        print("matplotlib not installed. Install with: pip install matplotlib")
        return

    _, gen_scores, detectors = get_raw_scores(group, ignore_detectors=ignore_detectors)
    categories = sorted(gen_scores.keys())
    cat_labels = [c.replace("_", " ").title() for c in categories]
    det_labels = [DETECTOR_DISPLAY_NAMES.get(d, d) for d in detectors]

    # Find models with adversarial data
    model_list = _find_adversarial_models(gen_scores, categories)

    if not model_list:
        print("No models with adversarial data found for detection heatmap.")
        return

    versions = [("base", "Base")] + [(v, _VERSION_DISPLAY[v]) for v in _ADV_VERSIONS]
    cmap = LinearSegmentedColormap.from_list(
        "blue", [(0.82, 0.88, 1.0), (0.05, 0.15, 0.50)]
    )

    n_models = len(model_list)
    with plt.rc_context(_PLOT_RC):
        fig, axes = plt.subplots(
            n_models,
            len(versions),
            figsize=(18, 5.5 * n_models),
            sharey="row",
            sharex=True,
            squeeze=False,
        )

        im = None
        for row_idx, (source_model, gen_short, display) in enumerate(model_list):
            for col_idx, (ver, ver_label) in enumerate(versions):
                # Determine which model key to use for this version
                if ver == "base":
                    model_key = source_model
                else:
                    model_key = f"adversarial-{ver}-{gen_short}"

                # Per-category rows + overall row
                per_cat = []
                for cat in categories:
                    row = []
                    for det in detectors:
                        scores = gen_scores[cat].get(model_key, {}).get(det, [])
                        rate = (
                            np.mean(np.array(scores) >= THRESHOLD) * 100
                            if scores
                            else 0.0
                        )
                        row.append(rate)
                    per_cat.append(row)

                overall = []
                for det in detectors:
                    pooled = []
                    for cat in categories:
                        pooled.extend(gen_scores[cat].get(model_key, {}).get(det, []))
                    overall.append(
                        np.mean(np.array(pooled) >= THRESHOLD) * 100 if pooled else 0.0
                    )

                matrix = np.array(per_cat + [overall])
                row_labels = cat_labels + ["Overall"]

                ax = axes[row_idx][col_idx]
                im = ax.imshow(matrix, cmap=cmap, vmin=0, vmax=100, aspect="auto")

                ax.set_xticks(range(len(detectors)))
                if row_idx == n_models - 1:
                    ax.set_xticklabels(det_labels, rotation=45, ha="right", fontsize=9)
                else:
                    ax.set_xticklabels([])

                if col_idx == 0:
                    ax.set_yticks(range(len(row_labels)))
                    ax.set_yticklabels(row_labels, fontsize=10)

                ax.set_title(
                    f"{display} \u2014 {ver_label}", fontsize=12, fontweight="bold"
                )
                _annotate_heatmap(ax, matrix, len(row_labels), len(detectors))
                _style_heatmap_ax(ax)
                # Thicker divider above the "Overall" row
                ax.axhline(y=len(categories) - 0.5, color="white", linewidth=3)

        fig.subplots_adjust(right=0.92)
        cbar_ax = fig.add_axes([0.94, 0.15, 0.015, 0.7])
        fig.colorbar(im, cax=cbar_ax, label="Detection Rate (%)")

        model_names = " and ".join(d for _, _, d in model_list)
        fig.suptitle(f"Detection Rates: {model_names}", fontsize=15, fontweight="bold")
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
    print(f"Detection heatmap saved to {output_path}")


def plot_effect_size_heatmap(
    output_path: str = "tmp/effect_size_heatmap.png",
    group: int = None,
    ignore_detectors: list = None,
):
    """Plot effect size (Cohen's d) heatmap with significance stars.

    Grayscale colormap for positive d (adversarial reduced detection), muted teal
    for negative d. Stars indicate p-value significance. One row per model that
    has adversarial data, columns for Mild/Strong.
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.colors import LinearSegmentedColormap, PowerNorm
    except ImportError:
        print("matplotlib not installed. Install with: pip install matplotlib")
        return

    rows, detectors, categories = compute_adversarial_paired_stats(
        group=group,
        ignore_detectors=ignore_detectors,
    )

    if rows is None or not rows:
        print("No models with adversarial data found for effect size heatmap.")
        return

    cat_labels = [c.replace("_", " ").title() for c in categories]
    det_labels = [DETECTOR_DISPLAY_NAMES.get(d, d) for d in detectors]

    # Find models with adversarial data from the stats rows
    gen_shorts_with_data = {r["model"] for r in rows}
    model_list = []
    for source_model in MAIN_MODELS:
        gen_short = _MODEL_TO_GENERATOR.get(source_model)
        if gen_short and gen_short in gen_shorts_with_data:
            display = MODEL_DISPLAY_NAMES.get(source_model, source_model)
            model_list.append((source_model, gen_short, display))

    if not model_list:
        print("No models with adversarial data found for effect size heatmap.")
        return

    cmap = LinearSegmentedColormap.from_list(
        "eff_gray", [(0.92, 0.92, 0.92), (0.45, 0.45, 0.45), (0.12, 0.12, 0.12)]
    )
    cmap = cmap.copy()
    neg_shade = (0.35, 0.55, 0.55)
    cmap.set_under(neg_shade)
    star_colors = ("#E8B030", "#AA7700")

    adv_versions = [("v7", "Mild"), ("v4", "Strong")]
    n_models = len(model_list)
    fig_h = 6 * n_models
    with plt.rc_context(_PLOT_RC):
        fig, axes = plt.subplots(
            n_models,
            len(adv_versions),
            figsize=(16, fig_h),
            sharey="row",
            squeeze=False,
        )

        im = None
        for row_idx, (source_model, gen_short, display) in enumerate(model_list):
            for col_idx, (ver, ver_label) in enumerate(adv_versions):
                ax = axes[row_idx][col_idx]
                d_mat, p_mat = _compute_d_and_p(
                    rows,
                    gen_short,
                    ver,
                    categories,
                    detectors,
                )
                vmax = max(d_mat.max(), 0.5)
                norm = PowerNorm(gamma=0.5, vmin=0, vmax=vmax)

                im = ax.imshow(d_mat, cmap=cmap, norm=norm, aspect="auto")
                ax.set_xticks(range(len(detectors)))
                if row_idx == n_models - 1:
                    ax.set_xticklabels(det_labels, rotation=45, ha="right", fontsize=9)
                else:
                    ax.set_xticklabels([])

                title = (
                    f"{display} \u2014 Base vs {ver_label}"
                    if n_models > 1
                    else f"Base vs {ver_label}"
                )
                ax.set_title(title, fontsize=13, fontweight="bold")

                if col_idx == 0:
                    ax.set_yticks(range(len(categories)))
                    ax.set_yticklabels(cat_labels, fontsize=10)

                for i in range(len(categories)):
                    for j in range(len(detectors)):
                        d = d_mat[i, j]
                        p = p_mat[i, j]
                        normed = norm(d) if d >= 0 else 0
                        txt_color = "white" if (d < 0 or normed > 0.55) else "black"

                        ax.text(
                            j,
                            i - 0.15,
                            f"{d:.2f}",
                            ha="center",
                            va="center",
                            fontsize=8,
                            color=txt_color,
                            fontweight="bold",
                        )

                        if p < 0.001:
                            stars = "***"
                        elif p < 0.01:
                            stars = "**"
                        elif p < 0.05:
                            stars = "*"
                        else:
                            stars = ""
                        if stars:
                            sc = (
                                star_colors[0]
                                if (d < 0 or normed > 0.35)
                                else star_colors[1]
                            )
                            ax.text(
                                j,
                                i + 0.2,
                                stars,
                                ha="center",
                                va="center",
                                fontsize=16,
                                color=sc,
                                fontweight="bold",
                            )

                _style_heatmap_ax(ax)

        fig.subplots_adjust(right=0.88)
        cbar_ax = fig.add_axes([0.90, 0.15, 0.015, 0.7])
        cbar = fig.colorbar(im, cax=cbar_ax, extend="min")
        cbar.set_label("Cohen's d")

        # Star legend — positioned just below the colorbar
        cbar_center_x = 0.90 + 0.015 / 2
        fig.text(
            cbar_center_x,
            0.10,
            "*** p < .001\n**  p < .01\n*   p < .05",
            fontsize=9,
            ha="center",
            va="center",
            color=star_colors[1],
            bbox=dict(
                boxstyle="round,pad=0.4",
                facecolor="white",
                edgecolor="#999999",
                linewidth=0.8,
            ),
        )

        model_names = " & ".join(d for _, _, d in model_list)
        fig.suptitle(
            f"Effect Size with Significance Stars ({model_names}) \u2014 Grayscale",
            fontsize=14,
            fontweight="bold",
        )
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
    print(f"Effect size heatmap saved to {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate detection report")
    parser.add_argument(
        "--adversarial", "-a", action="store_true", help="Show adversarial comparison"
    )
    parser.add_argument(
        "--output", "-o", default="tmp/roc_curves.png", help="ROC chart output path"
    )
    parser.add_argument(
        "--group",
        "-g",
        type=int,
        help="Only include images in this group (0, 1, 2, etc.)",
    )
    parser.add_argument(
        "--ignore-detectors",
        nargs="+",
        metavar="DETECTOR",
        help="Detectors to exclude from the report",
        default=["hive"],
    )
    parser.add_argument(
        "--adversarial-auc",
        action="store_true",
        help="Show AUC comparison: base vs adversarial mild/strong",
    )
    parser.add_argument(
        "--adversarial-paired",
        action="store_true",
        help="Matched-pairs analysis of adversarial effects on detector scores",
    )
    parser.add_argument(
        "--aggregate-roc",
        action="store_true",
        help="Plot ROC curves using aggregate (avg/min/max) detector scores",
    )
    parser.add_argument(
        "--latex", action="store_true", help="Output LaTeX tables for the paper"
    )
    parser.add_argument(
        "--detection-heatmap",
        action="store_true",
        help="Plot detection rate heatmap (Base/Mild/Strong per model)",
    )
    parser.add_argument(
        "--effect-size-heatmap",
        action="store_true",
        help="Plot effect size (Cohen's d) heatmap with significance stars",
    )
    args = parser.parse_args()

    if args.group is not None:
        print(f"Filtering to group {args.group}")
    if args.ignore_detectors:
        print(f"Ignoring detectors: {', '.join(args.ignore_detectors)}")

    ignore = args.ignore_detectors

    print_dataset_stats(group=args.group)
    print_report(group=args.group, ignore_detectors=ignore)

    if args.adversarial:
        print_adversarial_comparison(group=args.group, ignore_detectors=ignore)

    if args.adversarial_auc:
        print_adversarial_auc(group=args.group, ignore_detectors=ignore)

    if args.adversarial_paired:
        print_adversarial_paired_analysis(group=args.group, ignore_detectors=ignore)

    if args.aggregate_roc:
        plot_aggregate_roc(args.output, group=args.group, ignore_detectors=ignore)

    if args.latex:
        print_latex_by_class(group=args.group, ignore_detectors=ignore)
        print("\n")
        print_latex_aggregate(group=args.group, ignore_detectors=ignore)
        print("\n")
        print_latex_adversarial_by_class(group=args.group, ignore_detectors=ignore)

    if args.detection_heatmap:
        plot_detection_heatmap(
            "tmp/detection_heatmap.png", group=args.group, ignore_detectors=ignore
        )

    if args.effect_size_heatmap:
        plot_effect_size_heatmap(
            "tmp/effect_size_heatmap.png", group=args.group, ignore_detectors=ignore
        )