#!/usr/bin/env python3
"""Run model-family ablation with balanced config selection."""
import os, sys, yaml, json, time
from datetime import datetime
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from scripts.run_synthetic import build_model_family_grid, run_single_config
from src.utils import save_jsonl, ensure_dir

with open('config/synthetic_grid.yaml') as f:
    cfg = yaml.safe_load(f)

configs = build_model_family_grid(cfg)

# Select balanced: 8 per family
pg = [c for c in configs if c['model_family']=='poisson_gamma'][:8]
gg = [c for c in configs if c['model_family']=='gaussian_gaussian'][:8]
gm = [c for c in configs if c['model_family']=='gaussian_gamma_map'][:8]
selected = pg + gg + gm

print(f"Running {len(selected)} balanced configs: {len(pg)} PG, {len(gg)} GG, {len(gm)} GM")

output_dir = ensure_dir('results/raw')
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = os.path.join(output_dir, f'model_family_balanced_{timestamp}.jsonl')

all_records = []
for idx, config in enumerate(selected):
    print(f"\n>>> Config {idx+1}/{len(selected)} ({config['model_family']})")
    try:
        records = run_single_config(config, mode='model_family')
        all_records.extend(records)
        save_jsonl(records, output_file)
        print(f"  Saved {len(records)} records (total: {len(all_records)})")
    except Exception as e:
        print(f"  ERROR: {e}")
        import traceback
        traceback.print_exc()

print(f"\nCompleted. Total: {len(all_records)} records in {output_file}")