#!/usr/bin/env bash
set -uo pipefail
cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
export TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 MOLT5_PATH=models/molt5-base PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
mkdir -p logs models/abl models/pretrain reports

PT=models/pretrain_delivery.pt
NT=${NT:-15}; EPT=${EPT:-30}
MAX_RETRY=${MAX_RETRY:-5} 
NEED_MIB=${NEED_MIB:-8000}  

GPUS=${GPUS:-0}  
IFS=',' read -r -a GPU_ARR <<< "$GPUS"
PRETRAIN_GPU=${PRETRAIN_GPU:-${GPU_ARR[0]}}

ts(){ awk '{ print strftime("%F %T"), $0; fflush() }'; }

retry(){
  local log=$1; shift; local n=1 start
  while :; do
    start=$(date +%s)
    echo "[$(date '+%F %T')] >>> START (try $n/$((MAX_RETRY+1))) : $*" | tee -a "$log"
    if "$@" 2>&1 | ts | tee -a "$log"; then
      echo "[$(date '+%F %T')] <<< DONE ($(( ($(date +%s)-start)/60 )) min)" | tee -a "$log"; return 0; fi
    (( n > MAX_RETRY )) && { echo "[$(date '+%F %T')] !!! FAILED after $n tries" | tee -a "$log"; return 1; }
    echo "[$(date '+%F %T')] retry in 30s..." | tee -a "$log"; sleep 30; ((n++))
  done
}

NR=${NR:-3}
COMMON="--device cuda --n-trials $NT --epochs-per-trial $EPT --n-repeats $NR"
wait_free(){
  while :; do
    local free
    free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits \
           | sed -n "$((${CUDA_VISIBLE_DEVICES:-0}+1))p")
    [ "${free:-0}" -ge "$NEED_MIB" ] && return 0
    echo "[$(date '+%F %T')] GPU${CUDA_VISIBLE_DEVICES:-0} 仅剩 ${free}MiB(<${NEED_MIB}),等待空闲...(60s)" >&2
    sleep 60
  done
}

pretrain_arch(){ local name=$1; shift
  local pt="models/pretrain/${name}/pretrain_delivery.pt"
  if [ ! -f "$pt" ]; then
    pt_py(){ wait_free; python -m lnp_ml.modeling.pretrain main --device cuda --output-dir "models/pretrain/${name}" "$@"; }
    retry "logs/pretrain_${name}.log" pt_py "$@" >&2
  fi
  echo "$pt"
}

run_variant(){ local l=$1; shift
  if ls models/abl/"$l"/*/summary.json >/dev/null 2>&1; then
    echo "[$(date '+%F %T')] SKIP $l (已有完成结果)"; return 0
  fi
  local pt; pt=$(pretrain_arch "$l" "$@")
  run_py(){ wait_free; python -m lnp_ml.modeling.nested_cv_optuna $COMMON --init-from-pretrain "$pt" "$@" --output-dir "models/abl/$l"; }
  retry "logs/nested_${l}.log" run_py "$@"
}
echo "[$(date '+%F %T')] GPUS=$GPUS  PRETRAIN_GPU=$PRETRAIN_GPU"

# 0) 预训练已并入 run_variant

# 1) 消融
if [ "${#GPU_ARR[@]}" -ge 2 ]; then
  ( export CUDA_VISIBLE_DEVICES=${GPU_ARR[0]}
    run_variant baseline
    run_variant both  --use-moe --use-llm --llm-model-path models/molt5-base
  ) & G0=$!
  ( export CUDA_VISIBLE_DEVICES=${GPU_ARR[1]}
    run_variant moe   --use-moe
    run_variant llm   --use-llm --llm-model-path models/molt5-base
  ) & G1=$!
  wait $G0 $G1
else
  export CUDA_VISIBLE_DEVICES=${GPU_ARR[0]}
  run_variant baseline
  run_variant moe   --use-moe
  run_variant llm   --use-llm --llm-model-path models/molt5-base
  run_variant both  --use-moe --use-llm --llm-model-path models/molt5-base
fi

# 2) 汇总
latest(){ ls -dt "models/abl/$1"/*/summary.json 2>/dev/null | head -1 | xargs -r dirname; }
sum_args=()
for pair in "baseline:baseline" "+moe:moe" "+llm:llm" "+both:both"; do
  label=${pair%%:*}; name=${pair##*:}; d=$(latest "$name")
  [ -n "$d" ] && sum_args+=(--run "${label}=$d")
done
python scripts/summarize_ablation.py "${sum_args[@]}" \
  --out reports/ablation_summary.csv 2>&1 | ts | tee logs/summary.log