"""Jubair's hyperparameter grid (email of 2026-10-17) on BRCA, using the Workbench's train() unchanged.

Usage: grid.py <workspace.pkl> <out.jsonl> <shard> <n_shards> [seed]
Resumable: configs already in <out.jsonl> are skipped.

Grid as sent, mapped onto Workbench Hyperparameters:
  layer_one -> projection_dim; activation_function_one/two/three -> biological/projection/fusion activation;
  l2_reg -> weight_decay; epoch -> epochs (early stopping, patience 12, still applies);
  SGD gets momentum 0.03 as instructed. num_ex_neuron is not supported by the Workbench (fixed at 0).
"""
import itertools
import json
import sys
from dataclasses import asdict, replace
from pathlib import Path

import numpy as np
import pickle
import torch
from sklearn.metrics import recall_score

import biolmnet.training as training

WORKSPACE, OUT, SHARD, SHARDS = Path(sys.argv[1]), Path(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4])
SEED = int(sys.argv[5]) if len(sys.argv) > 5 else 42
torch.set_num_threads(2)

ACT = ["relu", "sigmoid", "tanh"]
GRID = list(itertools.product([64, 128], ACT, ACT, ACT, [16, 32], ["SGD", "Adam"]))

workspace = pickle.loads(WORKSPACE.read_bytes())
names = workspace.label_names
done = set()
if OUT.exists():
    done = {json.dumps(json.loads(line)["key"]) for line in OUT.read_text().splitlines() if line.strip()}

for index, (proj, a1, a2, a3, batch, opt) in enumerate(GRID):
    if index % SHARDS != SHARD:
        continue
    key = {"projection_dim": proj, "act1": a1, "act2": a2, "act3": a3, "batch_size": batch, "optimizer": opt, "seed": SEED}
    if json.dumps(key) in done:
        continue
    hp = replace(
        training.Hyperparameters(),
        epochs=200, batch_size=batch, learning_rate=0.01, weight_decay=0.01, dropout=0.3,
        projection_dim=proj, optimizer=opt, sgd_momentum=0.03 if opt == "SGD" else 0.0,
        biological_activation=a1, projection_activation=a2, fusion_activation=a3, seed=SEED,
    )
    try:
        result = training.train(workspace, hp)
    except Exception as exc:  # record and keep going
        row = {"key": key, "error": repr(exc)}
    else:
        m = result.bundle.metrics
        vp = result.validation_predictions
        counts = vp["predicted"].value_counts().reindex(names, fill_value=0).astype(int).tolist()
        history = result.bundle.history
        best_epoch = int(min(history, key=lambda h: h["validation_loss"])["epoch"])
        row = {
            "key": key, "hyperparameters": asdict(hp),
            "accuracy": m["accuracy"], "f1_macro": m["f1_macro"], "f1_weighted": m["f1_weighted"],
            "epochs_completed": m["epochs_completed"], "best_epoch": best_epoch,
            "pred_counts": counts, "classes_predicted": int((np.array(counts) > 0).sum()),
            "recall": recall_score(vp["observed"], vp["predicted"], labels=names, average=None, zero_division=0).tolist(),
            "final_train_loss": history[-1]["training_loss"], "best_val_loss": m["validation_loss"],
        }
    with OUT.open("a") as handle:
        handle.write(json.dumps(row) + "\n")
    print(f"[{SHARD}] {index} {key} -> " + (row.get("error") or f"f1m={row['f1_macro']:.3f} classes={row['classes_predicted']}"), flush=True)
print(f"SHARD {SHARD} DONE", flush=True)