benchmark.py

"""
MedCode Intelligence Network (MCIN) — Bittensor Subnet
Benchmark management: gold label sets, rotation, and task sampling.

The benchmark is the single most important validator asset.
Keep it private. Rotate it regularly. Never publish it.
"""

import json
import random
import hashlib
import logging
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

# ── Pool allocation (fraction of scored requests per epoch) ───────────────
POOL_GOLD_FRACTION        = 0.30   # Private expert-labeled cases
POOL_ONLINE_FRACTION      = 0.50   # Rolling live-traffic cases (anonymized)
POOL_ADVERSARIAL_FRACTION = 0.20   # Paraphrase variants for consistency testing

# ── Rotation schedule ─────────────────────────────────────────────────────
GOLD_ROTATION_DAYS       = 30    # Rotate private gold set monthly
ONLINE_ROTATION_DAYS     = 7     # Roll online cases weekly

# ── Paraphrase variants per gold case ─────────────────────────────────────
PARAPHRASE_VARIANTS_N    = 3     # Number of adversarial paraphrases per case


@dataclass
class BenchmarkCase:
    """
    A single benchmark task with gold labels.

    note_text and gold_codes are the minimum required fields.
    All cases must be fully de-identified (no PHI).
    """
    case_id:      str
    note_text:    str
    specialty:    str
    gold_codes:   List[Dict]         # List of CodeEntry dicts
    gold_modifiers: List[str] = field(default_factory=list)
    pool:         str = "gold"       # "gold" | "online" | "adversarial"
    created_at:   str = ""
    paraphrase_of: Optional[str] = None   # case_id of original if adversarial

    def to_request_dict(self, nonce: str = "") -> Dict:
        """Convert to MCINRequest kwargs."""
        return {
            "job_id":       self._make_job_id(nonce),
            "note_text":    self.note_text,
            "specialty":    self.specialty,
            "code_families": list({e["system"] for e in self.gold_codes}),
            "nonce":        nonce,
        }

    def _make_job_id(self, nonce: str = "") -> str:
        h = hashlib.sha256(f"{self.case_id}{nonce}{datetime.utcnow().isoformat()}".encode()).hexdigest()
        return h[:16]


class BenchmarkManager:
    """
    Manages the three benchmark pools and task sampling.

    Usage:
        mgr = BenchmarkManager(data_dir="/path/to/benchmark")
        mgr.load()

        # Sample tasks for one epoch
        tasks = mgr.sample_epoch_tasks(n=60)

        # Get gold label for a case
        gold = mgr.get_gold(case_id)
    """

    def __init__(self, data_dir: str = "./benchmark"):
        self.data_dir   = Path(data_dir)
        self._gold:        List[BenchmarkCase] = []
        self._online:      List[BenchmarkCase] = []
        self._adversarial: List[BenchmarkCase] = []
        self._last_rotation: Dict[str, datetime] = {}

    # ── Loading ────────────────────────────────────────────────────────────

    def load(self):
        """Load all benchmark pools from disk."""
        self._gold        = self._load_pool("gold.jsonl")
        self._online      = self._load_pool("online.jsonl")
        self._adversarial = self._load_pool("adversarial.jsonl")
        logger.info(
            f"Benchmark loaded: {len(self._gold)} gold | "
            f"{len(self._online)} online | "
            f"{len(self._adversarial)} adversarial"
        )

    def _load_pool(self, filename: str) -> List[BenchmarkCase]:
        path = self.data_dir / filename
        if not path.exists():
            logger.warning(f"Benchmark pool not found: {path}. Starting empty.")
            return []
        cases = []
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                d = json.loads(line)
                cases.append(BenchmarkCase(**d))
        return cases

    # ── Sampling ───────────────────────────────────────────────────────────

    def sample_epoch_tasks(self, n: int = 60, nonce_prefix: str = "") -> List[Tuple[Dict, BenchmarkCase]]:
        """
        Sample n tasks across the three pools using configured fractions.

        Returns list of (request_dict, gold_case) tuples.
        """
        n_gold   = max(1, int(n * POOL_GOLD_FRACTION))
        n_online = max(1, int(n * POOL_ONLINE_FRACTION))
        n_adv    = n - n_gold - n_online

        sampled = []
        sampled += self._sample_pool(self._gold,        n_gold,   nonce_prefix)
        sampled += self._sample_pool(self._online,      n_online, nonce_prefix)
        sampled += self._sample_pool(self._adversarial, n_adv,    nonce_prefix)

        random.shuffle(sampled)
        return sampled

    def _sample_pool(
        self,
        pool: List[BenchmarkCase],
        n: int,
        nonce_prefix: str,
    ) -> List[Tuple[Dict, BenchmarkCase]]:
        if not pool:
            return []
        chosen = random.choices(pool, k=min(n, len(pool)))
        nonce  = hashlib.sha256(f"{nonce_prefix}{random.random()}".encode()).hexdigest()[:12]
        return [(case.to_request_dict(nonce=nonce), case) for case in chosen]

    # ── Burst probe task generation ────────────────────────────────────────

    def sample_burst_tasks(self, n: int, burst_size: int, nonce_prefix: str = "") -> List[List[Tuple[Dict, BenchmarkCase]]]:
        """
        Generate burst_size batches of n tasks each for concurrent capacity probing.

        Returns a list of bursts; each burst is a list of (request_dict, gold_case).
        The validator fires all tasks in a single burst simultaneously.
        """
        bursts = []
        for _ in range(burst_size):
            nonce = hashlib.sha256(f"{nonce_prefix}{random.random()}".encode()).hexdigest()[:12]
            batch = self._sample_pool(self._gold, n, nonce_prefix=nonce)
            # Mark as burst probes
            for req_dict, _ in batch:
                req_dict["is_burst_probe"] = True
                req_dict["deadline_ms"] = 5000   # extra time for burst
            bursts.append(batch)
        return bursts

    # ── Gold label lookup ──────────────────────────────────────────────────

    def get_gold(self, case_id: str) -> Optional[BenchmarkCase]:
        """Look up a gold case by case_id across all pools."""
        for pool in (self._gold, self._online, self._adversarial):
            for case in pool:
                if case.case_id == case_id:
                    return case
        return None

    # ── Rotation ───────────────────────────────────────────────────────────

    def maybe_rotate(self):
        """Rotate pools if their rotation schedule has elapsed."""
        now = datetime.utcnow()
        for pool_name, days in [("gold", GOLD_ROTATION_DAYS), ("online", ONLINE_ROTATION_DAYS)]:
            last = self._last_rotation.get(pool_name, datetime.min)
            if now - last >= timedelta(days=days):
                logger.info(f"Rotating benchmark pool: {pool_name}")
                self._load_pool(f"{pool_name}.jsonl")
                self._last_rotation[pool_name] = now

    # ── Consistency scoring ────────────────────────────────────────────────

    def compute_consistency_score(
        self,
        responses_by_variant: Dict[str, List[Dict]],
    ) -> float:
        """
        Measure how consistent a miner's predictions are across paraphrase variants
        of the same clinical scenario.

        responses_by_variant: {case_id: [codes_list_variant_1, codes_list_variant_2, ...]}

        Returns float in [0, 1]: 1.0 = perfectly consistent across all variants.
        """
        if not responses_by_variant:
            return 1.0

        consistency_scores = []
        for case_id, variants in responses_by_variant.items():
            if len(variants) < 2:
                continue
            # Jaccard similarity between all pairs of variants
            sims = []
            for i in range(len(variants)):
                for j in range(i + 1, len(variants)):
                    set_a = {e["code"] for e in variants[i]}
                    set_b = {e["code"] for e in variants[j]}
                    union = set_a | set_b
                    inter = set_a & set_b
                    sims.append(len(inter) / len(union) if union else 1.0)
            consistency_scores.append(sum(sims) / len(sims))

        return sum(consistency_scores) / len(consistency_scores) if consistency_scores else 1.0