#!/bin/bash
set -e
exec > >(tee -a ~/pytorch-build.log) 2>&1

echo "=========================================="
echo " PyTorch 2.5.1 Build for BC-250 (gfx1010)"
echo " ROCm 7.2.0 | HIP 7.2 | $(date)"
echo "=========================================="

BUILD_DIR="$HOME/pytorch-build"
ROCM_ARCH="gfx1010"

##############################################
# Phase 1: Fix incomplete clone
##############################################
echo ""
echo "[1/5] Fixing submodules..."
cd "$BUILD_DIR"

# Reset any partial checkouts
git submodule sync --recursive
git submodule update --init --recursive --force --jobs 4 2>&1 || {
    echo "WARN: Some submodules failed, retrying one by one..."
    git submodule foreach --recursive 'git checkout . 2>/dev/null; true'
    git submodule update --init --recursive --force 2>&1 || true
}

echo "Clone/submodule size: $(du -sh "$BUILD_DIR" | cut -f1)"

##############################################
# Phase 2: Install build dependencies
##############################################
echo ""
echo "[2/5] Checking build dependencies..."
for pkg in cmake ninja gcc python3; do
    if ! command -v $pkg &>/dev/null; then
        echo "ERROR: $pkg not found!"
        exit 1
    fi
    echo "  $pkg: $(command -v $pkg)"
done

# Ensure python build deps
pip3 install --user cmake ninja pyyaml typing-extensions setuptools wheel 2>&1 | tail -3

##############################################
# Phase 3: Hipify (CUDA -> HIP conversion)
##############################################
echo ""
echo "[3/5] Running hipify (CUDA -> HIP)..."
cd "$BUILD_DIR"

if [ ! -f "aten/src/ATen/hip" ] || [ ! -d "aten/src/ATen/hip" ]; then
    python3 tools/amd_build/build_amd.py 2>&1 | tail -20
    echo "Hipify complete."
else
    echo "Hipify already done, skipping."
fi

##############################################
# Phase 4: Configure environment
##############################################
echo ""
echo "[4/5] Configuring build environment..."

# ROCm paths
export ROCM_PATH=/opt/rocm
export HIP_PATH=/opt/rocm
export ROCM_HOME=/opt/rocm
export HCC_HOME=/opt/rocm/hcc
export HIP_PLATFORM=amd

# Target architecture
export PYTORCH_ROCM_ARCH="$ROCM_ARCH"
export AMDGPU_TARGETS="$ROCM_ARCH"

# Build configuration
export USE_ROCM=1
export USE_CUDA=0
export USE_NCCL=0
export USE_DISTRIBUTED=0
export USE_MKLDNN=0
export USE_FBGEMM=0
export USE_KINETO=0
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export USE_NNPACK=0
export USE_XNNPACK=0
export BUILD_TEST=0
export BUILD_CAFFE2=0
export USE_FLASH_ATTENTION=0
export USE_MEM_EFF_ATTENTION=0

# BC-250 has limited RAM — reduce parallel jobs
export MAX_JOBS=3

# Fix CMake 4.x compatibility
export CMAKE_POLICY_VERSION_MINIMUM=3.5

# Compiler
export CC=gcc
export CXX=g++
export CMAKE_C_COMPILER=gcc
export CMAKE_CXX_COMPILER=g++

# HIP compiler
export HIP_CLANG_PATH=/opt/rocm/llvm/bin
export HIPCC_COMPILE_FLAGS_APPEND="--offload-arch=$ROCM_ARCH"

# BC-250 specific: need managed memory
export HSA_ENABLE_SDMA=0

echo "  PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
echo "  MAX_JOBS=$MAX_JOBS"
echo "  ROCM_PATH=$ROCM_PATH"
echo "  Python: $(python3 --version)"

##############################################
# Phase 5: Build
##############################################
echo ""
echo "[5/5] Building PyTorch wheel..."
echo "  Start time: $(date)"

cd "$BUILD_DIR"

# Clean any previous build artifacts
python3 setup.py clean 2>/dev/null || true

# Build wheel
python3 setup.py bdist_wheel 2>&1

echo ""
echo "=========================================="
echo " BUILD COMPLETE: $(date)"
echo "=========================================="
ls -lh dist/*.whl 2>/dev/null || echo "ERROR: No wheel file found!"
echo ""
echo "Install with:"
echo "  pip install dist/torch-*.whl"