#!/bin/bash
set -e

echo "============================================"
echo "  HTSR-VL Data Setup"
echo "============================================"

TRAIN_IMGS=$(find Data/training_series -name '*.png' 2>/dev/null | wc -l)
BENCH_IMGS=$(find MultiAgentTS/Data/benchmark_tb -name '*.png' 2>/dev/null | wc -l)
echo "  Training images found: $TRAIN_IMGS"
echo "  Benchmark images found: $BENCH_IMGS"

if [ "$TRAIN_IMGS" -lt 3000 ]; then
    echo "  WARNING: Expected ~3060 training images but found $TRAIN_IMGS."
    echo "  Make sure you cloned the repo with Git LFS enabled:"
    echo "    git lfs install && git clone https://huggingface.co/ANTICH/htsr-grpo-train"
    echo "  Or run: git lfs pull"
    exit 1
fi

# 1. Download base model (Qwen3-VL-8B-Instruct)
echo "[1/2] Downloading Qwen3-VL-8B-Instruct model..."
if [ -d "models/qwen3vl-8b-base" ] && [ "$(ls models/qwen3vl-8b-base/*.safetensors 2>/dev/null | wc -l)" -ge 4 ]; then
    echo "  Model already exists, skipping."
else
    mkdir -p models/qwen3vl-8b-base
    huggingface-cli download Qwen/Qwen3-VL-8B-Instruct \
        --local-dir models/qwen3vl-8b-base \
        --local-dir-use-symlinks False
    echo "  Model downloaded."
fi

# 2. Patch preprocessor_config.json for 512x512 resolution
echo "[2/2] Patching preprocessor_config.json (max_pixels=262144 for ~512x512)..."
python3 -c "
import json, os
cfg_path = 'models/qwen3vl-8b-base/preprocessor_config.json'
with open(cfg_path) as f:
    cfg = json.load(f)
cfg['size']['longest_edge'] = 262144
cfg['size']['shortest_edge'] = 16384
with open(cfg_path, 'w') as f:
    json.dump(cfg, f, indent=2)
print(f'  Patched: longest_edge={cfg[\"size\"][\"longest_edge\"]}, shortest_edge={cfg[\"size\"][\"shortest_edge\"]}')
"

echo ""
echo "============================================"
echo "  Setup complete!"
echo "  Model:    models/qwen3vl-8b-base/"
echo "  Train:    Data/training_series/ ($TRAIN_IMGS files)"
echo "  Bench:    MultiAgentTS/Data/benchmark_tb/ ($BENCH_IMGS files)"
echo ""
echo "  To start training:   uv run bash run_grpo.sh"
echo "  To run benchmark:    uv run python run_benchmark_eval.py --model-path models/qwen3vl-8b-base"
echo "============================================"