#!/usr/bin/env bash
set -euo pipefail

# One-shot pipeline:
# 1) Train Combined
# 2) Generate Combined inference CSV
# 3) Train VA (initialized from Combined checkpoint)
# 4) Generate VA inference CSV
# 5) Evaluate all available inference CSV files

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$SCRIPT_DIR"

# Defaults can be overridden by environment variables before running this script.
export AFFECTNET_ROOT="${AFFECTNET_ROOT:-/workspace/data_affectnet/AffectNet}"
export AFFECTNET_EXTRACT_ROOT="${AFFECTNET_EXTRACT_ROOT:-${AFFECTNET_ROOT}/extracted}"
export AFFECTNET_ANNO_ROOT="${AFFECTNET_ANNO_ROOT:-$SCRIPT_DIR/affectnet_annotations}"
export NUM_WORKERS="${NUM_WORKERS:-0}"
export BATCHSIZE="${BATCHSIZE:-64}"
export TRAIN_ARCHIVE="${TRAIN_ARCHIVE:-${AFFECTNET_ROOT}/train_set.tar}"
export VAL_ARCHIVE="${VAL_ARCHIVE:-${AFFECTNET_ROOT}/val_set.zip}"
export EXTRACT_ROOT="${EXTRACT_ROOT:-${AFFECTNET_EXTRACT_ROOT}}"
export TRAIN_EXTRACT_DIR="${TRAIN_EXTRACT_DIR:-${EXTRACT_ROOT}/train_extracted}"
export VAL_EXTRACT_DIR="${VAL_EXTRACT_DIR:-${EXTRACT_ROOT}/val_extracted}"
export AUTO_EXTRACT="${AUTO_EXTRACT:-1}"

PY_BIN="${PY_BIN:-python3}"
SKIP_PIP_INSTALL="${SKIP_PIP_INSTALL:-0}"

echo "[INFO] SCRIPT_DIR=$SCRIPT_DIR"
echo "[INFO] AFFECTNET_ROOT=$AFFECTNET_ROOT"
echo "[INFO] AFFECTNET_EXTRACT_ROOT=$AFFECTNET_EXTRACT_ROOT"
echo "[INFO] TRAIN_ARCHIVE=$TRAIN_ARCHIVE"
echo "[INFO] VAL_ARCHIVE=$VAL_ARCHIVE"
echo "[INFO] EXTRACT_ROOT=$EXTRACT_ROOT"
echo "[INFO] TRAIN_EXTRACT_DIR=$TRAIN_EXTRACT_DIR"
echo "[INFO] VAL_EXTRACT_DIR=$VAL_EXTRACT_DIR"
echo "[INFO] AUTO_EXTRACT=$AUTO_EXTRACT"
echo "[INFO] AFFECTNET_ANNO_ROOT=$AFFECTNET_ANNO_ROOT"
echo "[INFO] NUM_WORKERS=$NUM_WORKERS"
echo "[INFO] BATCHSIZE=$BATCHSIZE"
echo "[INFO] PY_BIN=$PY_BIN"

has_dataset_layout() {
  local root="$1"
  if [[ ! -d "$root" ]]; then
    return 1
  fi
  if [[ -d "$root/images" && -d "$root/annotations" ]]; then
    return 0
  fi
  if find "$root" -type d -name images 2>/dev/null | grep -q .; then
    if find "$root" -type d -name annotations 2>/dev/null | grep -q .; then
      return 0
    fi
  fi
  return 1
}

extract_tar_if_needed() {
  local tar_path="$1"
  local out_dir="$2"
  local marker="$out_dir/.extracted_done"
  mkdir -p "$out_dir"
  if [[ -f "$marker" ]]; then
    echo "[SKIP] already extracted: $tar_path"
    return
  fi
  echo "[EXTRACT] $tar_path -> $out_dir"
  tar -xf "$tar_path" -C "$out_dir"
  touch "$marker"
  echo "[DONE] train extracted"
}

extract_zip_if_needed() {
  local zip_path="$1"
  local out_dir="$2"
  local marker="$out_dir/.extracted_done"
  mkdir -p "$out_dir"
  if [[ -f "$marker" ]]; then
    echo "[SKIP] already extracted: $zip_path"
    return
  fi
  echo "[EXTRACT] $zip_path -> $out_dir"
  unzip -q "$zip_path" -d "$out_dir"
  touch "$marker"
  echo "[DONE] val extracted"
}

if [[ "$SKIP_PIP_INSTALL" != "1" ]]; then
  echo "[STEP] Installing dependencies"
  "$PY_BIN" -m pip install -r requirements.txt
  echo "[STEP] Installing mmcv (required by APViT)"
  "$PY_BIN" -m pip install mmcv
else
  echo "[STEP] Skipping dependency install (SKIP_PIP_INSTALL=1)"
fi

# ── APViT repo setup ──────────────────────────────────────────────────────────
APVIT_PATH="$SCRIPT_DIR/APViT"
if [[ ! -d "$APVIT_PATH/.git" ]]; then
  echo "[STEP] Cloning APViT repo"
  git clone https://github.com/youqingxiaozhua/APViT.git --depth=1 "$APVIT_PATH"
else
  echo "[SKIP] APViT repo already present at $APVIT_PATH"
fi

WEIGHTS_DIR="$APVIT_PATH/weights"
mkdir -p "$WEIGHTS_DIR"

VIT_SMALL_W="$WEIGHTS_DIR/vit_small_p16_224-15ec54c9.pth"
if [[ ! -f "$VIT_SMALL_W" ]]; then
  echo "[STEP] Downloading ViT-Small pretrained weights"
  wget -q -O "$VIT_SMALL_W" \
    "https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth"
else
  echo "[SKIP] ViT-Small weights already present"
fi

IR50_W="$WEIGHTS_DIR/backbone_ir50_ms1m_epoch120.pth"
if [[ ! -f "$IR50_W" ]]; then
  echo "[WARNING] IR-50 pretrained weights not found at $IR50_W"
  echo "[WARNING] Download backbone_ir50_ms1m_epoch120.pth from face.evoLVe model zoo"
  echo "[WARNING] (https://github.com/ZhaoJ9014/face.evoLVe/#model-zoo) and place it at $IR50_W"
  echo "[WARNING] Training will proceed with random init for the IR-50 extractor."
fi
# ─────────────────────────────────────────────────────────────────────────────

if [[ ! -f "$AFFECTNET_ANNO_ROOT/train_set_annotation_without_lnd.csv" ]]; then
  echo "[ERROR] Missing train annotation CSV in $AFFECTNET_ANNO_ROOT"
  exit 1
fi
if [[ ! -f "$AFFECTNET_ANNO_ROOT/val_set_annotation_without_lnd.csv" ]]; then
  echo "[ERROR] Missing val annotation CSV in $AFFECTNET_ANNO_ROOT"
  exit 1
fi

train_ready=0
val_ready=0
if has_dataset_layout "$TRAIN_EXTRACT_DIR"; then
  train_ready=1
fi
if has_dataset_layout "$VAL_EXTRACT_DIR"; then
  val_ready=1
fi

if [[ "$train_ready" -eq 1 && "$val_ready" -eq 1 ]]; then
  echo "[SKIP] Found existing extracted train/val folders"
else
  if [[ "$AUTO_EXTRACT" != "1" ]]; then
    echo "[ERROR] Extracted folders not found and AUTO_EXTRACT is disabled"
    echo "[HINT] Set AUTO_EXTRACT=1 or point TRAIN_EXTRACT_DIR / VAL_EXTRACT_DIR to existing extracted data"
    exit 1
  fi

  if [[ ! -f "$TRAIN_ARCHIVE" ]]; then
    echo "[ERROR] Missing TRAIN_ARCHIVE: $TRAIN_ARCHIVE"
    exit 1
  fi
  if [[ ! -f "$VAL_ARCHIVE" ]]; then
    echo "[ERROR] Missing VAL_ARCHIVE: $VAL_ARCHIVE"
    exit 1
  fi

  extract_tar_if_needed "$TRAIN_ARCHIVE" "$TRAIN_EXTRACT_DIR"
  extract_zip_if_needed "$VAL_ARCHIVE" "$VAL_EXTRACT_DIR"
fi

echo "[STEP] Train Combined"
cd "$SCRIPT_DIR/models/AffectNet8_Maxvit_Combined"
"$PY_BIN" train.py

echo "[STEP] Generate Combined inference CSV"
"$PY_BIN" generate_csv.py

echo "[STEP] Train VA"
cd "$SCRIPT_DIR/models/AffectNet8_Maxvit_VA"
"$PY_BIN" train.py

echo "[STEP] Generate VA inference CSV"
"$PY_BIN" generate_csv.py

echo "[STEP] Evaluate"
cd "$SCRIPT_DIR/models"
"$PY_BIN" evaluation.py

echo "[DONE] Full pipeline completed successfully."