#!/bin/bash
# =============================================================================
# ARC-AI Physics Pretraining Pipeline
# =============================================================================
# Streams THE WELL physics data from HuggingFace → pretrains temporal encoder
# → fine-tunes on robot demonstrations with domain randomization
#
# Requirements:
#   pip install torch datasets huggingface_hub pyyaml numpy
#   huggingface-cli login  (for push_to_hub)
#
# Usage:
#   bash scripts/run_physics_pretrain.sh              # full pipeline
#   bash scripts/run_physics_pretrain.sh --pretrain   # pretrain only
#   bash scripts/run_physics_pretrain.sh --finetune   # finetune only (needs checkpoint)
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
CONFIG="$PROJECT_DIR/configs/physics_pretrain.yaml"

# Environment
export PYTHONPATH="$PROJECT_DIR/src:${PYTHONPATH:-}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

echo "============================================="
echo "ARC-AI Physics Pretraining Pipeline"
echo "============================================="
echo "Config: $CONFIG"
echo "GPU: $CUDA_VISIBLE_DEVICES"
echo "Python: $(python3 --version)"
echo ""

# Check dependencies
python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
python3 -c "import datasets; print(f'HF Datasets {datasets.__version__}')"

MODE="${1:-full}"

case "$MODE" in
  --pretrain)
    echo "Mode: Pretrain only"
    python3 -c "
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
from physics_pretraining import PhysicsPretrainer, PhysicsPretrainConfig
config = PhysicsPretrainConfig()
trainer = PhysicsPretrainer(config)
trainer.train()
"
    ;;
  --finetune)
    echo "Mode: Finetune only (requires pretrained checkpoint)"
    CHECKPOINT="${2:-checkpoints/physics_pretrain/final.pt}"
    echo "Loading: $CHECKPOINT"
    python3 -c "
import logging, torch
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
from physics_pretraining import (
    PhysicsTemporalEncoder, PhysicsPretrainConfig,
    PolicyFinetuner, FinetuneConfig
)
config = PhysicsPretrainConfig()
encoder = PhysicsTemporalEncoder(config)
ckpt = torch.load('$CHECKPOINT', map_location='cpu')
encoder.load_state_dict(ckpt['encoder_state'])
ft_config = FinetuneConfig(pretrained_path='$CHECKPOINT')
finetuner = PolicyFinetuner(encoder, ft_config)
print('Finetuner ready. Provide dataloader via API.')
"
    ;;
  *)
    echo "Mode: Full pipeline (pretrain + finetune)"
    python3 -c "
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
from physics_pretraining import run_physics_pretraining
policy = run_physics_pretraining('$CONFIG')
print('Pipeline complete.')
"
    ;;
esac

echo ""
echo "============================================="
echo "Done."
echo "============================================="