#!/usr/bin/env bash
# Run a real-LOB backtest sweep via Argo Workflows.
#
# Reads the same `fxt-backtest sweep` grid YAML as the single-machine
# binary and fans each cell out to its own GPU pod. Mirrors the
# argo-train.sh + train-multi-seed-template.yaml conventions:
# `# __SWEEP_CELLS__` marker in the template is replaced with one
# WorkflowTask stanza per cell before submission.
#
# Usage:
#   ./scripts/argo-lob-sweep.sh --grid config/ml/sweep_decision_stride_example.yaml
#   ./scripts/argo-lob-sweep.sh --grid <path> --sha abc1234
#   ./scripts/argo-lob-sweep.sh --grid <path> --gpu-pool ci-training-h100
#   ./scripts/argo-lob-sweep.sh --grid <path> --watch
#   ./scripts/argo-lob-sweep.sh --grid <path> --dry-run > /tmp/wf.yaml
#
# Requires: argo CLI, kubectl, python3 (with PyYAML — universally
# available in foxhunt CI images).

set -euo pipefail

GRID=""
SHA="HEAD"
BRANCH="main"
GPU_POOL=""
SWEEP_TAG=""
WATCH=false
DRY_RUN=false
TAG=""

usage() {
  cat <<EOF
Usage: $(basename "$0") --grid <yaml> [OPTIONS]

Options:
  --grid <yaml>        REQUIRED. Sweep grid YAML (same format as
                       \`fxt-backtest sweep --grid\`).
  --sha <commit>       Git commit SHA (default: HEAD).
  --branch <branch>    Git branch (default: main).
  --gpu-pool <pool>    GPU node pool (default: ci-training-l40s; opt in to
                       ci-training-h100 for sm_90 / 80 GB).
  --sweep-tag <name>   Subdirectory name under <sweep-root>/ (default:
                       <basename of grid>-<short-sha>).
  --watch              Follow workflow logs after submission.
  --tag <t>            Label workflow with foxhunt-tag=<t>.
  --dry-run            Print rendered workflow YAML to stdout, do not submit.
  -h, --help           Show this help.
EOF
  exit 0
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --grid)       GRID="$2"; shift 2 ;;
    --sha)        SHA="$2"; shift 2 ;;
    --branch)     BRANCH="$2"; shift 2 ;;
    --gpu-pool)   GPU_POOL="$2"; shift 2 ;;
    --sweep-tag)  SWEEP_TAG="$2"; shift 2 ;;
    --watch)      WATCH=true; shift ;;
    --tag)        TAG="$2"; shift 2 ;;
    --dry-run)    DRY_RUN=true; shift ;;
    -h|--help)    usage ;;
    *) echo "unknown arg: $1" >&2; usage ;;
  esac
done

if [[ -z "$GRID" ]]; then
  echo "ERROR: --grid is required" >&2
  usage
fi
if [[ ! -f "$GRID" ]]; then
  echo "ERROR: grid file not found: $GRID" >&2
  exit 1
fi

# Default GPU pool follows feedback_default_to_l40s_pool.md (2026-05-09).
if [[ -z "$GPU_POOL" ]]; then
  GPU_POOL="ci-training-l40s"
fi
if [[ "$GPU_POOL" == "ci-training-h100" ]]; then
  CUDA_COMPUTE_CAP="90"
else
  CUDA_COMPUTE_CAP="89"
fi

# Short SHA for sweep-tag default.
if [[ "$SHA" == "HEAD" ]]; then
  RESOLVED_SHA=$(git rev-parse HEAD 2>/dev/null || echo "00000000")
else
  RESOLVED_SHA="$SHA"
fi
SHORT_SHA="${RESOLVED_SHA:0:9}"
if [[ -z "$SWEEP_TAG" ]]; then
  SWEEP_TAG="$(basename "$GRID" .yaml)-$SHORT_SHA"
fi

# Single Python invocation parses the grid + emits a JSON object with
# parsed base + cell records + the rendered Argo WorkflowTask stanzas
# for the __SWEEP_CELLS__ marker. Keeps the shell side parser-free.
TPL="$(dirname "$0")/../infra/k8s/argo/lob-backtest-sweep-template.yaml"
if [[ ! -f "$TPL" ]]; then
  echo "ERROR: template not found: $TPL" >&2
  exit 1
fi

RENDER_PY=$(python3 - "$GRID" <<'PYEOF'
import json, os, sys, yaml

with open(sys.argv[1], "r") as f:
    grid = yaml.safe_load(f)

base = grid.get("base") or {}
cells = grid.get("cells") or []
if not cells:
    print("ERROR: grid has no cells", file=sys.stderr)
    sys.exit(1)

def base_or(key, default):
    return base.get(key, default)

out = {
    "n_cells": len(cells),
    "base": {
        "data":            base_or("data", ""),
        "predecoded_dir":  base.get("predecoded_dir") or base.get("data", ""),
        "n_parallel":      base_or("n_parallel", 1),
        "latency_ns":      base_or("latency_ns", 100_000_000),
        "target_annual_vol_units": base_or("target_annual_vol_units", 50.0),
        "annualisation_factor":    base_or("annualisation_factor", 825.0),
        "max_lots":        base_or("max_lots", 5),
        "max_events":      base_or("max_events", 0),
        "seed":            base_or("seed", 0xC0FFEE),
        "checkpoint":      base.get("checkpoint") or "",
    },
}

# Render per-cell WorkflowTask stanzas + collect dependency names.
tasks = []
deps  = []
for c in cells:
    name = c.get("name")
    if not name:
        print("ERROR: cell missing 'name'", file=sys.stderr)
        sys.exit(1)
    stride       = c.get("decision_stride", base_or("decision_stride", 4))
    latency      = c.get("latency_ns", out["base"]["latency_ns"])
    target_vol   = c.get("target_annual_vol_units", out["base"]["target_annual_vol_units"])
    ann_factor   = c.get("annualisation_factor", out["base"]["annualisation_factor"])
    max_lots     = c.get("max_lots", out["base"]["max_lots"])
    max_events   = c.get("max_events", out["base"]["max_events"])
    seed         = c.get("seed", out["base"]["seed"])

    tasks.append(f'''
          - name: run-cell-{name}
            template: run-cell
            dependencies: [ensure-binary]
            arguments:
              parameters:
                - name: sha
                  value: "{{{{tasks.ensure-binary.outputs.parameters.sha}}}}"
                - name: cell-name
                  value: "{name}"
                - name: decision-stride
                  value: "{stride}"
                - name: latency-ns
                  value: "{latency}"
                - name: target-annual-vol-units
                  value: "{target_vol}"
                - name: annualisation-factor
                  value: "{ann_factor}"
                - name: max-lots
                  value: "{max_lots}"
                - name: max-events
                  value: "{max_events}"
                - name: seed
                  value: "{seed}"''')
    deps.append(f"run-cell-{name}")

out["cell_tasks"]   = "".join(tasks)
out["aggregate_deps"] = ", ".join(["ensure-binary", *deps])
out["short_sha"] = os.environ.get("SHORT_SHA", "")
out["sweep_tag"] = os.environ.get("SWEEP_TAG", "")
print(json.dumps(out))
PYEOF
)

export SHORT_SHA SWEEP_TAG

# Extract render outputs via python -c (no jq dependency).
N_CELLS=$(echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["n_cells"])')
BASE_DATA=$(       echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["data"])')
BASE_PREDECODED=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["predecoded_dir"])')
BASE_N_PARALLEL=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["n_parallel"])')
BASE_LATENCY_NS=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["latency_ns"])')
BASE_TARGET_VOL=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["target_annual_vol_units"])')
BASE_ANN_FACTOR=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["annualisation_factor"])')
BASE_MAX_LOTS=$(   echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["max_lots"])')
BASE_MAX_EVENTS=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["max_events"])')
BASE_CHECKPOINT=$( echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["base"]["checkpoint"])')
CELL_TASKS=$(      echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["cell_tasks"])')
AGG_DEPS=$(        echo "$RENDER_PY" | python3 -c 'import json,sys;print(json.load(sys.stdin)["aggregate_deps"])')

# Render template: replace # __SWEEP_CELLS__ marker AND inject all
# per-cell tasks into the aggregate's `dependencies: [...]` list.
RENDERED=$(awk -v cells="$CELL_TASKS" -v deps="$AGG_DEPS" '
  # Anchor: only match the actual marker (whitespace-indented at line start),
  # not the docstring reference at the top of the template that mentions
  # `# __SWEEP_CELLS__` inside a sentence.
  /^[[:space:]]+# __SWEEP_CELLS__/ { print cells; next }
  /dependencies: \[ensure-binary\]    # plus all run-cell-/ {
    sub(/\[ensure-binary\]/, "[" deps "]")
  }
  { print }
' "$TPL")

# Build the submission command.
SUBMIT_ARGS=(
  --from "workflowtemplate/lob-backtest-sweep"
  -p "commit-sha=$SHA"
  -p "git-branch=$BRANCH"
  -p "cuda-compute-cap=$CUDA_COMPUTE_CAP"
  -p "gpu-pool=$GPU_POOL"
  -p "data-root=$BASE_DATA"
  -p "predecoded-dir=$BASE_PREDECODED"
  -p "n-parallel=$BASE_N_PARALLEL"
  -p "latency-ns=$BASE_LATENCY_NS"
  -p "target-annual-vol-units=$BASE_TARGET_VOL"
  -p "annualisation-factor=$BASE_ANN_FACTOR"
  -p "max-lots=$BASE_MAX_LOTS"
  -p "max-events=$BASE_MAX_EVENTS"
  -p "sweep-tag=$SWEEP_TAG"
)
if [[ -n "$BASE_CHECKPOINT" ]]; then
  SUBMIT_ARGS+=(-p "checkpoint=$BASE_CHECKPOINT")
fi
if [[ -n "$TAG" ]]; then
  SUBMIT_ARGS+=(-l "foxhunt-tag=$TAG")
fi
if [[ "$WATCH" == "true" ]]; then
  SUBMIT_ARGS+=(--watch)
fi

if [[ "$DRY_RUN" == "true" ]]; then
  echo "# === Rendered workflow (cells: $N_CELLS, sweep-tag: $SWEEP_TAG) ==="
  echo "$RENDERED"
  echo ""
  echo "# === Would submit with: ==="
  echo "# argo submit ${SUBMIT_ARGS[*]}"
  exit 0
fi

# Apply the rendered template (covers any in-flight template edits) then submit.
echo "$RENDERED" | kubectl apply -f - -n foxhunt
argo submit -n foxhunt "${SUBMIT_ARGS[@]}"