src/wisp/search/controller.py

# Paper-scoped implementation; see SOURCE_PROVENANCE.json.
from __future__ import annotations
import hashlib
import json
import os
import pickle
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
from wisp.core.data import HARData, concatenate_har
from wisp.core.metrics import classification_report_dict
from wisp.cpu.algorithms import build_cpu_estimator, estimate_pickle_size_mb
from wisp.search.cascade import evaluate_candidate_cascade
from wisp.search.dataset_fingerprint import DatasetFingerprint, compute_dataset_fingerprint
from wisp.search.grammar import CandidateSpec, initial_candidates, mutate_candidate, normalize_spec
from wisp.search.operator_program import TemporalProgram
from wisp.search.posterior import MotifPosterior, island_key, sample_parent
from wisp.search.profiles import SearchProfile, get_profile
from wisp.utils.jsonl import append_jsonl, read_jsonl

def _positive_int_env(name: str, default: int) -> int:
    try:
        return max(int(os.environ.get(name, default)), 1)
    except (TypeError, ValueError):
        return max(int(default), 1)

def _stable_id(*, seed: int, generation: int, index: int, spec: CandidateSpec) -> str:
    payload = json.dumps(spec.to_dict(), sort_keys=True, separators=(',', ':'), default=str)
    digest = hashlib.sha1(f'{seed}:{generation}:{index}:{payload}'.encode()).hexdigest()[:12]
    return f'g{generation:03d}c{index:04d}_{digest}'

def _spec_key(spec: CandidateSpec) -> str:
    return json.dumps(normalize_spec(spec).to_dict(), sort_keys=True, separators=(',', ':'), default=str)

@dataclass
class WISPSearchController:
    train: HARData
    valid: HARData
    output: str | Path
    profile: str | SearchProfile = 'wisp_evolution'
    test: HARData | None = None
    population: int = 24
    generations: int = 20
    seed: int = 42
    cost_weight: float = 0.01
    robustness_weight: float = 0.1
    candidate_workers: int | None = None
    n_jobs_per_candidate: int | None = None

    def __post_init__(self) -> None:
        if isinstance(self.profile, str):
            self.profile = get_profile(self.profile)
        if self.train.n_classes != self.valid.n_classes:
            raise ValueError('train and validation must share the same global label vocabulary')
        if self.test is not None and self.test.n_classes != self.train.n_classes:
            raise ValueError('test must share the same global label vocabulary')

    def _worker_count(self) -> int:
        if self.candidate_workers is not None:
            return max(int(self.candidate_workers), 1)
        return _positive_int_env('WISP_CANDIDATE_WORKERS', 1)

    def _per_candidate_jobs(self) -> int:
        if self.n_jobs_per_candidate is not None:
            return max(int(self.n_jobs_per_candidate), 1)
        explicit = os.environ.get('WISP_N_JOBS_PER_CANDIDATE')
        if explicit is not None:
            return _positive_int_env('WISP_N_JOBS_PER_CANDIDATE', 1)
        total = _positive_int_env('WISP_N_JOBS', 1)
        return max(total // self._worker_count(), 1)

    def _runtime_spec(self, spec: CandidateSpec) -> CandidateSpec:
        spec = normalize_spec(spec)
        if spec.route != 'cpu':
            raise ValueError('WISPSearchController is CPU-only; GPU search is not released')
        params = dict(spec.params)
        params.setdefault('n_jobs', self._per_candidate_jobs())
        return CandidateSpec(name=spec.name, params=params, operators=list(spec.operators), route='cpu', parent_id=spec.parent_id, mutation=spec.mutation)

    def _constrain_spec(self, spec: CandidateSpec) -> CandidateSpec | None:
        spec = self._runtime_spec(spec)
        allowed = self.profile.allowed_cpu_families
        if allowed is not None and spec.name not in set(allowed):
            return None
        if self.profile.force_no_hmm:
            params = dict(spec.params)
            params['use_hmm'] = False
            operators = [op for op in spec.operators if op != 'hmm']
            if 'argmax' not in operators:
                operators.append('argmax')
            spec = CandidateSpec(name=spec.name, params=params, operators=operators, route='cpu', parent_id=spec.parent_id, mutation=spec.mutation)
        if spec.mutation in set(self.profile.disabled_mutations):
            return None
        return spec

    def _evaluate_one(self, spec: CandidateSpec) -> tuple[str, dict[str, Any]]:
        try:
            cascade = evaluate_candidate_cascade(spec, self.train, self.valid, cost_weight=self.cost_weight, robustness_weight=self.robustness_weight if self.profile.use_stress else 0.0, max_stress_transforms=4 if self.profile.use_stress else 0)
            result = cascade.to_dict()
            predictive = float(result.get('cascade_score', result.get('score', -1000000000.0)))
            result['predictive_utility'] = predictive
            return ('ok' if cascade.accepted else 'rejected', result)
        except Exception as exc:
            return ('failed', {'predictive_utility': -1000000000.0, 'breeding_score': -1000000000.0, 'error': repr(exc), 'candidate': spec.to_dict()})

    def _evaluate_many(self, specs: list[CandidateSpec]) -> list[tuple[str, dict[str, Any]]]:
        if not specs:
            return []
        workers = min(self._worker_count(), len(specs))
        if workers <= 1:
            return [self._evaluate_one(spec) for spec in specs]
        with ThreadPoolExecutor(max_workers=workers, thread_name_prefix='wisp') as pool:
            return list(pool.map(self._evaluate_one, specs))

    def _initial_specs(self, posterior: MotifPosterior, fingerprint: DatasetFingerprint | None, family_transfer_bias: dict[str, float], rng: np.random.Generator, count: int) -> list[CandidateSpec]:
        specs = posterior.rank_initial_candidates(fingerprint=fingerprint, seed=self.seed, route='cpu', include_v2=self.profile.include_v2_families, motif_mode=self.profile.motif_mode, family_transfer_bias=family_transfer_bias)
        specs = [constrained for spec in specs if (constrained := self._constrain_spec(spec)) is not None]
        if not specs:
            raise RuntimeError(f'Profile {self.profile.name} removed every initial CPU family')
        base = list(specs)
        i = 0
        while len(specs) < count:
            parent = base[i % len(base)]
            clone = CandidateSpec.from_dict(parent.to_dict())
            clone.params = dict(clone.params)
            clone.params['seed'] = int(clone.params.get('seed', self.seed)) + 1009 * (i + 1)
            clone.mutation = 'initial_seed_variant'
            specs.append(clone)
            i += 1
        return [spec for spec in specs[:count]]

    def _random_spec(self, rng: np.random.Generator, fingerprint: DatasetFingerprint | None) -> CandidateSpec:
        bases = initial_candidates(seed=int(rng.integers(0, 2 ** 31 - 1)), include_v2=self.profile.include_v2_families, route='cpu')
        spec = CandidateSpec.from_dict(bases[int(rng.integers(0, len(bases)))].to_dict())
        depth = int(rng.integers(1, 5))
        for _ in range(depth):
            spec = mutate_candidate(spec, rng=rng, mutation_bias=None, motif_weights=None if fingerprint is None else fingerprint.motif_weights())
        spec.parent_id = None
        spec.mutation = f'random_depth_{depth}'
        constrained = self._constrain_spec(spec)
        if constrained is None:
            return self._random_spec(rng, fingerprint)
        return constrained

    def _record(self, *, db_path: Path, candidate_id: str, generation: int, spec: CandidateSpec, status: str, result: dict[str, Any], fingerprint: DatasetFingerprint | None, parent_id: str | None) -> dict[str, Any]:
        row = {'id': candidate_id, 'generation': int(generation), 'parent_id': parent_id, 'island': island_key(spec), 'status': status, 'profile': self.profile.to_dict(), 'dataset_fingerprint': None if fingerprint is None else fingerprint.to_dict(), 'motif_weights': None if fingerprint is None else fingerprint.motif_weights(), 'operator_certificate': TemporalProgram.from_candidate(spec).certificate(), 'result': result, 'spec': spec.to_dict()}
        append_jsonl(db_path, row)
        return row

    def _final_refit(self, best: dict[str, Any], out_dir: Path) -> dict[str, Any] | None:
        if self.test is None:
            return None
        spec = CandidateSpec.from_dict(best['spec'])
        model = build_cpu_estimator(spec.to_dict())
        refit = concatenate_har(self.train, self.valid, source='train+validation')
        t0 = time.perf_counter()
        model.fit(refit.X, refit.y, groups=refit.subject, time_index=refit.time_index, n_classes=refit.n_classes)
        fit_s = time.perf_counter() - t0
        t1 = time.perf_counter()
        pred = model.predict(self.test.X, groups=self.test.subject, time_index=self.test.time_index)
        pred_s = time.perf_counter() - t1
        metrics = classification_report_dict(self.test.y, pred, fit_seconds=fit_s, predict_seconds=pred_s, model_size_mb=estimate_pickle_size_mb(model), label_names=self.test.label_names)
        np.savez_compressed(out_dir / 'locked_test_predictions.npz', y_true=self.test.y, y_pred=pred, subject=self.test.subject, time_index=self.test.time_index)
        with (out_dir / 'selected_model.pkl').open('wb') as f:
            pickle.dump(model, f)
        (out_dir / 'locked_test_metrics.json').write_text(json.dumps(metrics, indent=2, ensure_ascii=False, default=str), encoding='utf-8')
        return metrics

    def run(self) -> dict[str, Any]:
        out_dir = Path(self.output)
        out_dir.mkdir(parents=True, exist_ok=True)
        db_path = out_dir / 'program_database.jsonl'
        summary_path = out_dir / 'summary.json'
        config = {'profile': self.profile.to_dict(), 'population': int(self.population), 'generations': int(self.generations), 'seed': int(self.seed), 'cost_weight': float(self.cost_weight), 'robustness_weight': float(self.robustness_weight), 'candidate_workers': self._worker_count(), 'n_jobs_per_candidate': self._per_candidate_jobs()}
        if summary_path.exists():
            old = json.loads(summary_path.read_text(encoding='utf-8'))
            if old.get('status') == 'completed' and old.get('search_config') == config:
                return old
        if db_path.exists() and db_path.stat().st_size:
            raise RuntimeError(f'Refusing to append to non-empty search database {db_path}')
        rng = np.random.default_rng(self.seed)
        fingerprint = compute_dataset_fingerprint(self.train, seed=self.seed) if self.profile.use_fingerprint else None
        family_transfer, mutation_transfer = ({}, {})
        posterior = MotifPosterior()
        archive: list[CandidateSpec] = []
        population: list[tuple[str, CandidateSpec, dict[str, Any]]] = []
        rows: list[dict[str, Any]] = []
        seen: set[str] = set()
        budget = int(self.profile.max_evaluations)
        initial_count = min(9, budget)
        seeds = self._initial_specs(posterior, fingerprint, family_transfer, rng, initial_count)
        seed_results = self._evaluate_many(seeds)
        for i, (spec, (status, result)) in enumerate(zip(seeds, seed_results)):
            novelty = 0.0
            result['novelty_bonus'] = novelty
            result['breeding_score'] = float(result.get('predictive_utility', -1000000000.0)) + novelty
            cid = _stable_id(seed=self.seed, generation=0, index=i, spec=spec)
            row = self._record(db_path=db_path, candidate_id=cid, generation=0, spec=spec, status=status, result=result, fingerprint=fingerprint, parent_id=None)
            rows.append(row)
            seen.add(_spec_key(spec))
            archive.append(spec)
            if status == 'ok':
                population.append((cid, spec, result))
                if self.profile.use_posterior:
                    posterior.update(spec, result, score_key='predictive_utility')
        if not population:
            raise RuntimeError('No accepted initial candidate')
        evaluated = len(seeds)
        generation = 0
        while evaluated < budget and generation < int(self.generations):
            generation += 1
            remaining = budget - evaluated
            proposal_count = min(max(self.population - max(2, self.population // 3), 1), remaining)
            proposals: list[tuple[str, str | None, CandidateSpec]] = []
            attempts = 0
            mutation_bias = posterior.mutation_bias(fingerprint, motif_mode=self.profile.motif_mode, transfer_bias=mutation_transfer) if self.profile.use_posterior else mutation_transfer if self.profile.transfer_prior else {}
            for disabled in self.profile.disabled_mutations:
                mutation_bias.pop(disabled, None)
            while len(proposals) < proposal_count:
                attempts += 1
                if attempts > proposal_count * 100:
                    raise RuntimeError('Unable to generate enough unique candidate specifications')
                if self.profile.search_strategy == 'random':
                    parent_id = None
                    child = self._random_spec(rng, fingerprint)
                else:
                    parent_id, parent, parent_result = sample_parent(population, rng, use_islands=self.profile.use_islands, score_key='breeding_score')
                    child = None
                    if child is None:
                        child = mutate_candidate(parent, rng=rng, parent_id=parent_id, mutation_bias=mutation_bias, motif_weights=None if fingerprint is None else fingerprint.motif_weights())
                    child = self._constrain_spec(child)
                    if child is None:
                        continue
                key = _spec_key(child)
                if key in seen:
                    continue
                seen.add(key)
                index = evaluated + len(proposals)
                proposals.append((_stable_id(seed=self.seed, generation=generation, index=index, spec=child), parent_id, child))
            evaluated_results = self._evaluate_many([spec for _, _, spec in proposals])
            new_items: list[tuple[str, CandidateSpec, dict[str, Any]]] = []
            for (cid, parent_id, spec), (status, result) in zip(proposals, evaluated_results):
                novelty = posterior.novelty_bonus(spec, archive) if self.profile.use_novelty and status == 'ok' else 0.0
                result['novelty_bonus'] = float(novelty)
                result['breeding_score'] = float(result.get('predictive_utility', -1000000000.0)) + float(novelty)
                row = self._record(db_path=db_path, candidate_id=cid, generation=generation, spec=spec, status=status, result=result, fingerprint=fingerprint, parent_id=parent_id)
                rows.append(row)
                archive.append(spec)
                if status == 'ok':
                    new_items.append((cid, spec, result))
                    if self.profile.use_posterior:
                        posterior.update(spec, result, score_key='predictive_utility')
            evaluated += len(proposals)
            if self.profile.search_strategy == 'evolution':
                combined = population + new_items
                combined.sort(key=lambda item: float(item[2].get('breeding_score', -1000000000.0)), reverse=True)
                population = combined[:int(self.population)]
            else:
                population = (population + new_items)[-int(self.population):]
            if not population:
                raise RuntimeError(f'No accepted candidate after generation {generation}')
        ok_rows = [row for row in rows if row['status'] == 'ok']
        if not ok_rows:
            raise RuntimeError('Search completed without an accepted candidate')
        final_key = self.profile.final_selection
        best = max(ok_rows, key=lambda row: float(row['result'].get(final_key, -1000000000.0)))
        best_breeding = max(ok_rows, key=lambda row: float(row['result'].get('breeding_score', -1000000000.0)))
        best_predictive = max(ok_rows, key=lambda row: float(row['result'].get('predictive_utility', -1000000000.0)))
        (out_dir / 'selected_candidate.json').write_text(json.dumps(best, indent=2, ensure_ascii=False, default=str), encoding='utf-8')
        test_metrics = self._final_refit(best, out_dir)
        summary = {'status': 'completed', 'version': 'wisp_search_controller', 'profile': self.profile.name, 'n_rows': len(rows), 'n_ok': len(ok_rows), 'n_evaluated': evaluated, 'best': best, 'best_predictive': best_predictive, 'best_breeding': best_breeding, 'locked_test_metrics': test_metrics, 'search_config': config, 'dataset_fingerprint': None if fingerprint is None else {**fingerprint.to_dict(), 'motifs': fingerprint.motifs(), 'motif_weights': fingerprint.motif_weights()}, 'posterior_summary': posterior.summary(), 'cross_task_family_bias': family_transfer, 'cross_task_mutation_bias': mutation_transfer}
        summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False, default=str), encoding='utf-8')
        return summary
__all__ = ['WISPSearchController']