#!/bin/bash
# CASCADE HF Space — Asset Setup Script
# Run once from your terminal: bash setup_assets.sh
set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
NOVA="$SCRIPT_DIR/../NNPredictor/webapp_NMR/Predictor"
HF="$SCRIPT_DIR"

echo "=== CASCADE HF Space Asset Setup ==="
echo "Source: $NOVA"
echo "Target: $HF"
echo ""

# ── 1. Python inference code (with preprocessor path fix) ──────────────────
echo "[1/5] Copying inference code..."
sed "s|os.path.join('cascade', 'preprocessor.p')|os.path.join('NMR_Prediction', 'preprocessor.p')|g" \
    "$NOVA/NMR_Prediction/apply.py" > "$HF/NMR_Prediction/apply.py"
cp "$NOVA/NMR_Prediction/genConf.py" "$HF/NMR_Prediction/genConf.py"
echo "    apply.py   ✓  (preprocessor path fixed)"
echo "    genConf.py ✓"

# ── 2. Preprocessor pickle ─────────────────────────────────────────────────
echo "[2/5] Copying preprocessor..."
cp "$NOVA/NMR_Prediction/preprocessor.p" "$HF/NMR_Prediction/preprocessor.p"
echo "    preprocessor.p ✓  ($(du -sh "$HF/NMR_Prediction/preprocessor.p" | cut -f1))"

# ── 3. Model weights ───────────────────────────────────────────────────────
echo "[3/5] Copying model weights..."
mkdir -p "$HF/NMR_Prediction/schnet_edgeupdate"
cp "$NOVA/NMR_Prediction/schnet_edgeupdate/best_model.hdf5" \
   "$HF/NMR_Prediction/schnet_edgeupdate/best_model.hdf5"
echo "    best_model.hdf5 ✓  ($(du -sh "$HF/NMR_Prediction/schnet_edgeupdate/best_model.hdf5" | cut -f1))"

# ── 4. Vendored nfp package ────────────────────────────────────────────────
echo "[4/5] Copying nfp package..."
cp -r "$NOVA/nfp/." "$HF/nfp/"
echo "    nfp/ ✓  ($(find "$HF/nfp" -type f | wc -l | tr -d ' ') files)"

# ── 5. Static assets ───────────────────────────────────────────────────────
echo "[5/5] Copying static assets..."
cp -r "$NOVA/NMR_Prediction/static/NMR_Prediction/." "$HF/static/"
echo "    static/ ✓  ($(find "$HF/static" -type f | wc -l | tr -d ' ') files)"

echo ""
echo "=== Setup complete! ==="
echo "Total hf_space size: $(du -sh "$HF" | cut -f1)"
echo ""
echo "Next: docker build -t cascade-local . && docker run -p 7860:7860 cascade-local"
echo "Then: open http://localhost:7860/cascade_v1/predict/"