image-classifier/validate_classifier.py

"""
validate_classifier.py — Test the design code classifier against examiner ground truth
========================================================================================
Pulls a sample of image marks from your Supabase database that already have
USPTO examiner-assigned design codes, runs each through the classifier, and
compares the classifier's output against the examiner's codes.

The output tells you concretely whether the classifier is good enough to
ship — or whether we need to iterate before exposing it to attorney-tier
customers.

WHAT IT MEASURES
----------------
For each test image, four agreement levels:
  - Section-exact:  XX.YY.ZZ matches exactly (strictest)
  - Division-level: XX.YY matches (same code family)
  - Category-level: XX matches (same broad category)
  - Any overlap:    at least one code in common

Plus retrieval-style metrics:
  - Precision:  of codes the classifier returned, what fraction were correct?
  - Recall:     of examiner codes, what fraction did the classifier find?
  - F1:         harmonic mean

WHAT TO LOOK FOR
----------------
For an MVP shipping to entrepreneurs at $49/mo:
  - Category-level agreement >= 80% is reasonable
  - Recall >= 60% means we catch most relevant matches

For attorney-tier customers at $250+/mo:
  - Section-exact agreement >= 70%
  - Recall >= 80% (missing codes is a liability risk)

If the numbers come in lower, look at the per-image CSV — usually the
classifier is right but using a slightly different code than the examiner,
or vice versa. That's data we can use to tune the prompt.

REQUIREMENTS
------------
Same as design_code_classifier.py, plus:
    pip install httpx supabase

USAGE
-----
    # Quick test — 20 random samples
    python validate_classifier.py --samples 20

    # Full validation — 100 samples, slower but more reliable
    python validate_classifier.py --samples 100

    # Resume after a crash
    python validate_classifier.py --resume

OUTPUT
------
    validation_results.csv  — per-image: serial, examiner codes, classifier codes, scores
    validation_summary.json — aggregate metrics
"""

import os
import sys
import csv
import json
import asyncio
import logging
import random
from pathlib import Path
from datetime import datetime, timezone
from typing import List, Dict, Any

import httpx
from dotenv import load_dotenv
from supabase import create_client, Client

# Local import — must be in same directory
from design_code_classifier import classify_image


# ============================================================================
# CONFIG
# ============================================================================

env_path = Path(__file__).parent / ".env"
load_dotenv(dotenv_path=env_path)

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

CONFIDENCE_THRESHOLD = 0.7   # only count classifier codes at or above this confidence

OUTPUT_CSV = Path(__file__).parent / "validation_results.csv"
OUTPUT_JSON = Path(__file__).parent / "validation_summary.json"
CHECKPOINT_PATH = Path(__file__).parent / ".validation_checkpoint.json"

# Politeness — don't hammer your DB or Claude API
DELAY_BETWEEN_SAMPLES_S = 1.0

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("validate")


# ============================================================================
# CHECKPOINT (resume after crash)
# ============================================================================

def load_checkpoint() -> Dict[str, Any]:
    if CHECKPOINT_PATH.exists():
        try:
            return json.loads(CHECKPOINT_PATH.read_text())
        except Exception as e:
            logger.warning(f"Checkpoint unreadable, starting fresh: {e}")
    return {"completed_serials": [], "results": []}


def save_checkpoint(state: Dict[str, Any]):
    state["updated_at"] = datetime.now(timezone.utc).isoformat()
    CHECKPOINT_PATH.write_text(json.dumps(state, indent=2))


# ============================================================================
# SAMPLE FETCHING
# ============================================================================

async def fetch_sample(supabase: Client, n: int, exclude_serials: set) -> List[Dict]:
    """Pull N random image marks that have both an image and examiner codes.

    We over-fetch and randomize client-side because Postgres ORDER BY RANDOM()
    on millions of rows is brutally slow.
    """
    # Fetch a pool of candidates with codes + image (skip NOT_FOUND)
    pool_size = min(n * 20, 5000)   # 20x oversample, capped
    logger.info(f"📡 Fetching {pool_size} candidate image marks from Supabase...")

    response = (
        supabase.table("trademarks_images")
        .select("serial_number,image_url,design_search_codes,mark_text")
        .neq("image_url", "NOT_FOUND")
        .not_.is_("image_url", "null")
        .not_.is_("design_search_codes", "null")
        .neq("design_search_codes", "")
        .limit(pool_size)
        .execute()
    )

    pool = [
        r for r in response.data
        if r["serial_number"] not in exclude_serials
        and r.get("image_url")
        and r.get("design_search_codes")
    ]
    logger.info(f"  Pool size after filtering: {len(pool)}")

    # Random sample
    random.shuffle(pool)
    return pool[:n]


# ============================================================================
# IMAGE FETCHING
# ============================================================================

async def fetch_image_bytes(url: str, http_client: httpx.AsyncClient) -> bytes:
    """Download a Supabase Storage image."""
    resp = await http_client.get(url, timeout=20.0)
    resp.raise_for_status()
    return resp.content


# ============================================================================
# COMPARISON METRICS
# ============================================================================

def normalize_codes(codes_str: str) -> set:
    """Parse a comma-separated codes string into a set of XX.YY.ZZ codes.

    USPTO bulk XML stores codes as 6-digit strings without separators
    (e.g., "260121"). The classifier returns dotted format (e.g., "26.01.21").
    This function accepts either format and normalizes everything to dotted
    XX.YY.ZZ so set comparison works correctly.
    """
    if not codes_str:
        return set()
    normalized = set()
    for c in codes_str.split(","):
        c = c.strip()
        if not c:
            continue
        # Already-dotted format passes through unchanged
        if "." in c:
            normalized.add(c)
            continue
        # 6-digit USPTO format → insert dots: "260121" → "26.01.21"
        if len(c) == 6 and c.isdigit():
            normalized.add(f"{c[0:2]}.{c[2:4]}.{c[4:6]}")
            continue
        # Anything else: keep as-is (will simply not match, which is correct)
        normalized.add(c)
    return normalized


def compare_codes(examiner: set, classifier: set) -> Dict[str, Any]:
    """Compute multi-level agreement between two code sets."""
    examiner_divisions = {".".join(c.split(".")[:2]) for c in examiner}
    examiner_categories = {c.split(".")[0] for c in examiner}
    classifier_divisions = {".".join(c.split(".")[:2]) for c in classifier}
    classifier_categories = {c.split(".")[0] for c in classifier}

    section_overlap = examiner & classifier
    division_overlap = examiner_divisions & classifier_divisions
    category_overlap = examiner_categories & classifier_categories

    return {
        "examiner_count": len(examiner),
        "classifier_count": len(classifier),
        "section_exact_matches": len(section_overlap),
        "division_matches": len(division_overlap),
        "category_matches": len(category_overlap),
        "any_section_overlap": bool(section_overlap),
        "any_division_overlap": bool(division_overlap),
        "any_category_overlap": bool(category_overlap),
        # Retrieval metrics (treat examiner codes as ground truth)
        "precision": len(section_overlap) / max(len(classifier), 1),
        "recall": len(section_overlap) / max(len(examiner), 1),
    }


def f1_score(precision: float, recall: float) -> float:
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


# ============================================================================
# PER-SAMPLE RUN
# ============================================================================

async def validate_one(
    record: Dict,
    http_client: httpx.AsyncClient,
) -> Dict[str, Any]:
    """Run the classifier on one image and compare to examiner codes."""
    serial = record["serial_number"]
    image_url = record["image_url"]
    examiner_codes_raw = record["design_search_codes"]
    examiner_codes = normalize_codes(examiner_codes_raw)

    try:
        image_bytes = await fetch_image_bytes(image_url, http_client)
    except Exception as e:
        return {
            "serial_number": serial,
            "status": "image_fetch_failed",
            "error": str(e),
        }

    try:
        result = await classify_image(image_bytes)
    except Exception as e:
        return {
            "serial_number": serial,
            "status": "classification_failed",
            "error": str(e),
        }

    classifier_codes = set(result.high_confidence_codes(threshold=CONFIDENCE_THRESHOLD))
    metrics = compare_codes(examiner_codes, classifier_codes)

    return {
        "serial_number": serial,
        "mark_text": record.get("mark_text", ""),
        "image_url": image_url,
        "examiner_codes": sorted(examiner_codes),
        "classifier_codes": sorted(classifier_codes),
        "image_description": result.image_description,
        "status": "ok",
        **metrics,
    }


# ============================================================================
# MAIN
# ============================================================================

async def run(n_samples: int, resume: bool):
    if not all([SUPABASE_URL, SUPABASE_KEY]):
        logger.error("❌ SUPABASE_URL / SUPABASE_KEY not set in .env")
        sys.exit(1)

    state = load_checkpoint() if resume else {"completed_serials": [], "results": []}
    completed = set(state["completed_serials"])
    results: List[Dict] = state["results"]

    if resume and completed:
        logger.info(f"📋 Resuming — {len(completed)} samples already done")
        n_samples = max(0, n_samples - len(completed))
        if n_samples == 0:
            logger.info("✅ Sample target already met from checkpoint")

    supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

    if n_samples > 0:
        samples = await fetch_sample(supabase, n_samples, exclude_serials=completed)
        logger.info(f"🎯 Will validate {len(samples)} new samples")
    else:
        samples = []

    async with httpx.AsyncClient() as http_client:
        for idx, record in enumerate(samples, 1):
            serial = record["serial_number"]
            logger.info(f"\n[{idx}/{len(samples)}] Validating {serial} ({record.get('mark_text', '')[:40]})")

            try:
                result = await validate_one(record, http_client)
                results.append(result)
                completed.add(serial)

                if result["status"] == "ok":
                    logger.info(
                        f"  Examiner: {result['examiner_codes']}\n"
                        f"  Classifier: {result['classifier_codes']}\n"
                        f"  P={result['precision']:.2f} R={result['recall']:.2f} "
                        f"category_match={result['any_category_overlap']}"
                    )
                else:
                    logger.warning(f"  ⚠️ {result['status']}: {result.get('error', '')[:200]}")
            except Exception as e:
                logger.error(f"  ❌ Unexpected error: {e}")
                results.append({
                    "serial_number": serial,
                    "status": "unexpected_error",
                    "error": str(e),
                })
                completed.add(serial)

            # Checkpoint after each sample
            state["completed_serials"] = list(completed)
            state["results"] = results
            save_checkpoint(state)

            await asyncio.sleep(DELAY_BETWEEN_SAMPLES_S)

    # ── Aggregate ──
    ok_results = [r for r in results if r.get("status") == "ok"]
    if not ok_results:
        logger.error("❌ No successful validations to aggregate")
        return

    avg_precision = sum(r["precision"] for r in ok_results) / len(ok_results)
    avg_recall = sum(r["recall"] for r in ok_results) / len(ok_results)
    avg_f1 = f1_score(avg_precision, avg_recall)

    section_match_rate = sum(1 for r in ok_results if r["any_section_overlap"]) / len(ok_results)
    division_match_rate = sum(1 for r in ok_results if r["any_division_overlap"]) / len(ok_results)
    category_match_rate = sum(1 for r in ok_results if r["any_category_overlap"]) / len(ok_results)

    summary = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "total_samples": len(results),
        "successful_samples": len(ok_results),
        "failed_samples": len(results) - len(ok_results),
        "confidence_threshold": CONFIDENCE_THRESHOLD,
        "metrics": {
            "section_exact_match_rate": section_match_rate,
            "division_match_rate": division_match_rate,
            "category_match_rate": category_match_rate,
            "avg_precision": avg_precision,
            "avg_recall": avg_recall,
            "avg_f1": avg_f1,
        },
    }

    OUTPUT_JSON.write_text(json.dumps(summary, indent=2))

    # CSV — per-sample, easy to sort/filter in a spreadsheet
    with OUTPUT_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "serial_number", "mark_text", "image_url",
            "examiner_codes", "classifier_codes",
            "section_overlap", "division_overlap", "category_overlap",
            "precision", "recall", "image_description", "status",
        ])
        for r in results:
            if r.get("status") != "ok":
                writer.writerow([r.get("serial_number"), "", "", "", "", "", "", "", "", "", "", r.get("status")])
                continue
            writer.writerow([
                r["serial_number"],
                r["mark_text"],
                r["image_url"],
                ",".join(r["examiner_codes"]),
                ",".join(r["classifier_codes"]),
                r["any_section_overlap"],
                r["any_division_overlap"],
                r["any_category_overlap"],
                f"{r['precision']:.3f}",
                f"{r['recall']:.3f}",
                r["image_description"],
                "ok",
            ])

    # ── Print summary ──
    logger.info("\n" + "=" * 70)
    logger.info("📊 VALIDATION RESULTS")
    logger.info("=" * 70)
    logger.info(f"  Samples:                    {len(ok_results)} successful, {len(results) - len(ok_results)} failed")
    logger.info(f"  Confidence threshold:       {CONFIDENCE_THRESHOLD}")
    logger.info("")
    logger.info(f"  Section-exact match rate:   {section_match_rate:.1%}  (any XX.YY.ZZ in common)")
    logger.info(f"  Division match rate:        {division_match_rate:.1%}  (any XX.YY in common)")
    logger.info(f"  Category match rate:        {category_match_rate:.1%}  (any XX in common)")
    logger.info("")
    logger.info(f"  Avg precision:              {avg_precision:.1%}  (classifier codes that were right)")
    logger.info(f"  Avg recall:                 {avg_recall:.1%}  (examiner codes the classifier found)")
    logger.info(f"  Avg F1:                     {avg_f1:.3f}")
    logger.info("")
    logger.info(f"  Per-sample CSV: {OUTPUT_CSV}")
    logger.info(f"  Summary JSON:   {OUTPUT_JSON}")
    logger.info("=" * 70)


# ============================================================================
# CLI
# ============================================================================

async def main():
    import argparse
    parser = argparse.ArgumentParser(description="Validate the design code classifier")
    parser.add_argument(
        "--samples", type=int, default=20,
        help="How many random image marks to test against (default: 20)"
    )
    parser.add_argument(
        "--resume", action="store_true",
        help="Skip samples already in .validation_checkpoint.json"
    )
    parser.add_argument(
        "--reset-checkpoint", action="store_true",
        help="Delete checkpoint and start fresh"
    )
    args = parser.parse_args()

    if args.reset_checkpoint and CHECKPOINT_PATH.exists():
        CHECKPOINT_PATH.unlink()
        logger.info("🗑️  Checkpoint cleared")

    await run(n_samples=args.samples, resume=args.resume)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("\n⚠️ Interrupted — checkpoint saved, safe to re-run with --resume")
        sys.exit(0)