#!/usr/bin/env bash
# cleanup.sh — Purge venv, downloaded data, checkpoints, caches, and wandb logs
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

echo "=========================================="
echo "  HRM_optimised — Cleanup"
echo "=========================================="

# Deactivate venv if currently active
if [ -n "${VIRTUAL_ENV:-}" ]; then
    echo "  Deactivating current virtual environment..."
    deactivate 2>/dev/null || true
fi

# 1. Remove virtual environment
if [ -d "venv" ]; then
    echo "[1/6] Removing virtual environment (venv/)..."
    rm -rf venv
else
    echo "[1/6] No venv found, skipping."
fi

# 2. Remove downloaded/generated datasets
if [ -d "data" ]; then
    echo "[2/6] Removing generated datasets (data/)..."
    rm -rf data
else
    echo "[2/6] No data/ directory found, skipping."
fi

# 3. Remove raw downloaded data from HuggingFace
if [ -d "dataset/raw-data" ]; then
    echo "[3/6] Removing raw dataset downloads (dataset/raw-data/)..."
    rm -rf dataset/raw-data
else
    echo "[3/6] No dataset/raw-data/ found, skipping."
fi

# 4. Remove training checkpoints
if [ -d "checkpoints" ]; then
    echo "[4/6] Removing training checkpoints (checkpoints/)..."
    rm -rf checkpoints
else
    echo "[4/6] No checkpoints/ directory found, skipping."
fi

# 5. Remove wandb logs
if [ -d "wandb" ]; then
    echo "[5/6] Removing W&B logs (wandb/)..."
    rm -rf wandb
else
    echo "[5/6] No wandb/ directory found, skipping."
fi

# 6. Remove Python caches
echo "[6/6] Removing __pycache__ directories..."
find . -type d -name "__pycache__" -not -path "./.git/*" -exec rm -rf {} + 2>/dev/null || true
find . -type d -name ".hydra" -exec rm -rf {} + 2>/dev/null || true

echo ""
echo "=========================================="
echo "  ✅ Cleanup complete!"
echo "  Run ./startup.sh to set up a fresh environment."
echo "=========================================="