#!/usr/bin/env python3
"""
W&B sweep wrapper for SPAtop.
Called by `wandb agent` for each trial. Writes a SPANet options JSON and
runs spanet.train with the current W&B run context injected via env vars so
SPANet's Lightning WandbLogger streams metrics live to the sweep run.
"""
import json, os, subprocess, sys
import wandb

BASE_CONFIG = {
    "num_encoder_layers": 4,
    "num_branch_embedding_layers": 3,
    "num_branch_encoder_layers": 3,
    "num_jet_embedding_layers": 0,
    "num_jet_encoder_layers": 2,
    "num_detector_layers": 2,
    "num_regression_layers": 3,
    "num_classification_layers": 3,
    "split_symmetric_attention": True,
    "num_attention_heads": 8,
    "transformer_activation": "gelu",
    "linear_block_type": "GRU",
    "transformer_type": "Gated",
    "linear_activation": "gelu",
    "normalization": "LayerNorm",
    "masking": "Filling",
    "skip_connections": True,
    "initial_embedding_skip_connections": True,
    "event_info_file": "/data/spatop/event_files/v11/tt_hadronic_v7_full.yaml",
    "training_file": "/data/spatop/tpm70_wpm30_FB350_training/all_merged.h5",
    "normalize_features": True,
    "limit_to_num_jets": 0,
    "balance_jets": False,
    "partial_events": True,
    "balance_particles": False,
    "dataset_limit": 1.0,
    "train_validation_split": 0.95,
    "num_dataloader_workers": 8,
    "mask_sequence_vectors": True,
    "combine_pair_loss": "min",
    "optimizer": "AdamW",
    "focal_gamma": 0.0,
    "learning_rate_cycles": 1,
    "learning_rate_warmup_epochs": 1.0,
    "assignment_loss_scale": 1.0,
    "detection_loss_scale": 1.0,
    "kl_loss_scale": 0.0,
    "regression_loss_scale": 0.0,
    "classification_loss_scale": 0.0,
    "l2_penalty": 0.0002,
    "gradient_clip": 10.0,
    "epochs": 100,
    "num_gpu": 1,
    "verbose_output": True,
}

run = wandb.init(project="spatop-sweep")
config = dict(wandb.config)
run_id = run.id
run_project = run.project

hidden_dim = config.get("hidden_dim", 128)
full_config = {
    **BASE_CONFIG,
    **config,
    "hidden_dim": hidden_dim,
    "transformer_dim": hidden_dim,
    "initial_embedding_dim": hidden_dim,
    "position_embedding_dim": hidden_dim,
}

os.makedirs("/data/spatop/options_files/sweep", exist_ok=True)
os.makedirs("/data/spatop/logs/sweep", exist_ok=True)

config_path = f"/data/spatop/options_files/sweep/config_{run_id}.json"
with open(config_path, "w") as f:
    json.dump(full_config, f, indent=2)

run_name = f"sweep_{run_id}"
log_dir = "/data/spatop/logs/sweep"

print(f"=== Trial {run.name} | config: {json.dumps(config, indent=2)}")

# Finish the wrapper's hold on the run before launching the subprocess.
# SPANet's WandbLogger will resume it via WANDB_RUN_ID, so there is
# only one active wandb client at a time and metrics stream live.
wandb.finish()

proc = subprocess.run(
    ["python", "-m", "spanet.train", "-of", config_path, "-l", log_dir, "-n", run_name],
    env={
        **os.environ,
        "WANDB_PROJECT": run_project,
        "WANDB_RUN_ID": run_id,
        "WANDB_RESUME": "allow",
    },
)

if proc.returncode != 0:
    print(f"Training failed (exit code {proc.returncode})")
    sys.exit(1)