#!/bin/bash
# Companion to scripts/repr_experiment.sh.
#
# Timeline:
#   Phase I   (parallel with main pipeline after S2b):
#               extract ISCXTor2016 pcaps into unified artifacts.
#               CPU-bound, coexists with GPU training of E0/E1/E2.
#   Phase II  (after main pipeline DONE):
#               for each trained model (E0, E1, E2), run detect + per_class
#               against ISCXTor2016 (benign=nontor, attack=tor), emitting
#               `iscxtor_eval/` subdir per model.
#   Phase III : unified summary across both transfer targets.
#
# Log layout (reused from main repr_experiment):
#   $MAIN_DIR/companion.log              — this script's orchestration log
#   $MAIN_DIR/companion_iscxtor_extract.log
#   $MAIN_DIR/<tag>/iscxtor_eval/        — per-model ISCXTor2016 results

set -uo pipefail

ROOT=/home/chy/mambafortrafficmodeling
cd "$ROOT"

MAIN_DIR="runs/repr_experiment_20260423_092147"
MAIN_LOG="$MAIN_DIR/orch.log"
COMP_LOG="$MAIN_DIR/companion.log"
mkdir -p "$MAIN_DIR"
exec > >(tee -a "$COMP_LOG") 2>&1

N_VAL=20000
N_ATK=20000        # ISCXTor2016 has fewer attack flows than CICDDoS2019
SPLIT_SEED=42

echo "========================================================================"
echo "= $(date): iscxtor_companion START                                      ="
echo "= main dir: $MAIN_DIR                                                   ="
echo "========================================================================"

wait_for_pattern() {
    local pattern=$1 log=$2 desc=$3
    echo ">>> $(date): waiting for '$desc' (pattern='$pattern' in $log)"
    local waited=0
    while ! grep -q "$pattern" "$log" 2>/dev/null; do
        sleep 60
        waited=$((waited + 60))
        if (( waited % 600 == 0 )); then
            echo "    [heartbeat $(date +%H:%M:%S)] waited ${waited}s for $desc"
        fi
    done
    echo "<<< $(date): '$desc' detected after ${waited}s wait"
}

run_stage() {
    local name=$1; shift
    local log="$MAIN_DIR/${name}.log"
    echo ""
    echo ">>> $(date): [$name] START"
    local t0=$(date +%s)
    if ! "$@" > "$log" 2>&1; then
        local t1=$(date +%s)
        echo "!!! $(date): [$name] FAILED after $((t1-t0))s — see $log"
        tail -30 "$log"
        exit 1
    fi
    local t1=$(date +%s)
    echo "<<< $(date): [$name] OK in $((t1-t0))s  (log: $log)"
    tail -6 "$log" | sed 's/^/    | /'
}

# =====================================================================
# Phase I — ISCXTor2016 extraction (after S2b, parallel with main training)
# =====================================================================

wait_for_pattern "s2b_extract_cicddos2019_01-12.*OK" "$MAIN_LOG" \
    "S2b CICDDoS2019 01-12 extraction to complete"

run_stage "companion_iscxtor_extract" \
    nice -n 10 ionice -c 3 uv run python scripts/extract_iscxtor2016.py \
        --skip-decompress --jobs 6

# =====================================================================
# Phase II — wait for main pipeline DONE, then detect + per_class
# =====================================================================

wait_for_pattern "repr_experiment DONE" "$MAIN_LOG" \
    "main repr_experiment to finish (S7 summary)"

detect_and_per_class_iscxtor() {
    local tag=$1
    local src="$MAIN_DIR/$tag"
    local dst="$MAIN_DIR/$tag/iscxtor_eval"

    if [ ! -f "$src/model.pt" ]; then
        echo "!!! $(date): [$tag] model.pt not found at $src — skipping"
        return 1
    fi

    mkdir -p "$dst"
    # Symlink the trained model into the eval subdir — detect.py reads model.pt
    # from --save-dir. This keeps the original $tag/ directory pristine
    # (CICDDoS2019 artifacts stay where they were).
    ln -sf "../model.pt" "$dst/model.pt"

    run_stage "${tag}_detect_iscxtor" \
        uv run python -m detect \
            --save-dir "$dst" \
            --packets-npz  datasets/iscxtor2016/processed/packets.npz \
            --flows-parquet datasets/iscxtor2016/processed/flows.parquet \
            --benign-label nontor \
            --per-class-column activity \
            --n-val "$N_VAL" --n-atk "$N_ATK" --seed "$SPLIT_SEED"

    run_stage "${tag}_per_class_iscxtor" \
        uv run python -m eval.per_class --save-dir "$dst"
}

detect_and_per_class_iscxtor "e0_baseline"
detect_and_per_class_iscxtor "e1_relv2"
detect_and_per_class_iscxtor "e2_relv2_ctx"

# =====================================================================
# Phase III — unified summary across both transfer targets
# =====================================================================

run_stage "companion_summary" \
    uv run python scripts/summarize_repr_exp.py --root "$MAIN_DIR" --with-iscxtor

echo ""
echo "========================================================================"
echo "= $(date): iscxtor_companion DONE                                       ="
echo "= results: $MAIN_DIR/{summary.txt, summary.json}                         ="
echo "========================================================================"
