grid_search_v6.py

"""
grid_search_v6.py — Hatch v6 Ensemble Validation Grid Search
=============================================================

Validates the Multi-Scale Ensemble on the full test set and sweeps:
  1. Voting strategies: majority, weighted, unanimous, any
  2. Per-scale weight combinations for the "weighted" strategy
  3. Scale subsets: [20,30], [30,40], [20,40], [20,30,40]

Also measures individual scale performance for comparison.

Output: optimization_report_v6.csv
"""

import csv
import time
from itertools import product

import numpy as np
from sklearn.metrics import f1_score

from training_engine import train as _train
from inference_engine_v6 import HatchEnsembleV6, SCALE_PARAMS


# ─── Constants ────────────────────────────────────────────────────────────────

LOG_FILE   = "/tmp/grid_v6.log"
OUTPUT_CSV = "optimization_report_v6.csv"

# ─── Build test set ───────────────────────────────────────────────────────────

def build_test_set(n_per_class: int = 200) -> list:
    rng = np.random.default_rng(42)
    print(f"Loading test sequences (n={n_per_class} per class)...")
    _, _, meta = _train(window_size=30)
    all_test = meta.get("test_sequences", [])
    folded_valid    = [s for s in all_test if s["label"] == 1 and len(s["sequence"]) >= 40]
    disorder_valid  = [s for s in all_test if s["label"] == 0 and len(s["sequence"]) >= 40]
    idx  = rng.choice(len(folded_valid),   size=n_per_class, replace=False)
    idx2 = rng.choice(len(disorder_valid), size=n_per_class, replace=False)
    test_set = [folded_valid[i] for i in idx] + [disorder_valid[i] for i in idx2]
    rng.shuffle(test_set)
    print(f"  Test set: {len(test_set)} sequences ({n_per_class} folded, {n_per_class} disordered)")
    return test_set


# ─── Evaluation ───────────────────────────────────────────────────────────────

def evaluate(ensemble: HatchEnsembleV6, test_set: list) -> dict:
    y_true, y_pred = [], []
    for item in test_set:
        y_true.append(item["label"])
        y_pred.append(ensemble.classify(item["sequence"]))

    y_true = np.array(y_true)
    y_pred = np.array(y_pred)

    acc      = float(np.mean(y_true == y_pred))
    f1_fold  = float(f1_score(y_true, y_pred, pos_label=1, zero_division=0))
    f1_dis   = float(f1_score(y_true, y_pred, pos_label=0, zero_division=0))
    mean_f1  = (f1_fold + f1_dis) / 2.0

    return {
        "accuracy":      round(acc,     4),
        "f1_folded":     round(f1_fold, 4),
        "f1_disordered": round(f1_dis,  4),
        "mean_f1":       round(mean_f1, 4),
    }


# ─── Main ─────────────────────────────────────────────────────────────────────

def main():
    log = open(LOG_FILE, "w", buffering=1)

    def pr(msg):
        print(msg)
        log.write(msg + "\n")

    test_set = build_test_set(200)
    pr(f"  Test set ready: {len(test_set)} sequences")
    pr("")

    results = []
    best_mean_f1 = 0.0
    best_combo   = None
    combo_n      = 0
    t0           = time.time()

    # ── 1. Individual scale baselines ──────────────────────────────────────────
    pr("=== Individual Scale Baselines ===")
    for w in [20, 30, 40]:
        ens = HatchEnsembleV6(scales=[w], strategy="majority")
        m   = evaluate(ens, test_set)
        pr(f"  W={w:2d} alone:  mean_f1={m['mean_f1']:.4f}  "
           f"f1_fold={m['f1_folded']:.4f}  f1_dis={m['f1_disordered']:.4f}")
        row = {"combo": f"W{w}_solo", "scales": str(w), "strategy": "majority",
               "w20": 1.0, "w30": 1.0, "w40": 1.0, **m}
        results.append(row)
        if m["mean_f1"] > best_mean_f1:
            best_mean_f1 = m["mean_f1"]
            best_combo   = row.copy()
    pr("")

    # ── 2. Scale subsets ───────────────────────────────────────────────────────
    pr("=== Scale Subset Experiments ===")
    for scale_set in [[20, 30], [30, 40], [20, 40], [20, 30, 40]]:
        for strategy in ["majority", "unanimous", "any"]:
            combo_n += 1
            ens = HatchEnsembleV6(scales=scale_set, strategy=strategy)
            m   = evaluate(ens, test_set)
            marker = ""
            if m["mean_f1"] > best_mean_f1:
                best_mean_f1 = m["mean_f1"]
                best_combo   = {"combo": combo_n, "scales": str(scale_set),
                                "strategy": strategy, "w20": 1.0, "w30": 1.0, "w40": 1.0, **m}
                marker = "  *** NEW BEST ***"
            pr(f"  [{combo_n:3d}] scales={scale_set} strategy={strategy:<12}  "
               f"mean_f1={m['mean_f1']:.4f}  f1_fold={m['f1_folded']:.4f}  "
               f"f1_dis={m['f1_disordered']:.4f}{marker}")
            row = {"combo": combo_n, "scales": str(scale_set), "strategy": strategy,
                   "w20": 1.0, "w30": 1.0, "w40": 1.0, **m}
            results.append(row)
    pr("")

    # ── 3. Weighted strategy: sweep weight combinations ────────────────────────
    pr("=== Weighted Strategy: Weight Sweep ===")
    # Weight values to try for each scale
    weight_options = [0.5, 0.75, 1.0, 1.25, 1.5]
    # Only sweep [20, 30, 40] with weighted strategy
    for w20, w30, w40 in product(weight_options, repeat=3):
        combo_n += 1
        weights = {20: w20, 30: w30, 40: w40}
        ens = HatchEnsembleV6(scales=[20, 30, 40], strategy="weighted",
                              scale_weights=weights)
        m = evaluate(ens, test_set)
        marker = ""
        if m["mean_f1"] > best_mean_f1:
            best_mean_f1 = m["mean_f1"]
            best_combo   = {"combo": combo_n, "scales": "[20, 30, 40]",
                            "strategy": "weighted", "w20": w20, "w30": w30, "w40": w40, **m}
            marker = "  *** NEW BEST ***"
        pr(f"  [{combo_n:3d}] w=[{w20:.2f},{w30:.2f},{w40:.2f}]  "
           f"mean_f1={m['mean_f1']:.4f}  f1_fold={m['f1_folded']:.4f}  "
           f"f1_dis={m['f1_disordered']:.4f}{marker}")
        row = {"combo": combo_n, "scales": "[20, 30, 40]", "strategy": "weighted",
               "w20": w20, "w30": w30, "w40": w40, **m}
        results.append(row)

    elapsed = time.time() - t0
    pr("")
    pr("=" * 60)
    pr(f"GRID SEARCH COMPLETE ({elapsed:.1f}s)")
    pr("=" * 60)
    if best_combo:
        pr(f"Best: scales={best_combo['scales']} strategy={best_combo['strategy']} "
           f"w=[{best_combo['w20']},{best_combo['w30']},{best_combo['w40']}]")
        pr(f"  mean_f1={best_combo['mean_f1']:.4f}  "
           f"f1_fold={best_combo['f1_folded']:.4f}  "
           f"f1_dis={best_combo['f1_disordered']:.4f}")

    # Write CSV
    results.sort(key=lambda r: -r["mean_f1"])
    if results:
        with open(OUTPUT_CSV, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=list(results[0].keys()))
            writer.writeheader()
            writer.writerows(results)
        pr(f"\nResults saved to {OUTPUT_CSV}")
    log.close()


if __name__ == "__main__":
    main()