#!/usr/bin/env bash
set -euo pipefail

PROJECT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$PROJECT_DIR"

echo "=== Setting up HK Weather Prediction Market Pipeline ==="

# Use Python 3.10 for JAX compatibility
export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"

if ! pyenv versions | grep -q "3.10.17"; then
    echo "Python 3.10.17 not found in pyenv. Please install it first."
    exit 1
fi

if [ ! -d "venv" ]; then
    pyenv local 3.10.17
    python -m venv venv
fi

source venv/bin/activate

echo "=== Installing dependencies ==="

pip install --upgrade pip setuptools wheel

# JAX with CUDA 12 support
pip install "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Data & scientific
pip install numpy pandas xarray scipy netCDF4 h5netcdf cfgrib
pip install zarr fsspec gcsfs

# Weather data access
pip install ecmwf-api-client openmeteo-requests requests-cache retry-requests
pip install cdsapi

# Polymarket
pip install py-clob-client websocket-client

# Utilities
pip install httpx python-dotenv click rich tqdm schedule
pip install matplotlib cartopy

# WeatherNext itself
pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0

echo ""
echo "=== Verifying installations ==="
python -c "import jax; print(f'JAX version: {jax.__version__}'); print(f'Devices: {jax.devices()}')"
python -c "import weathernext; print('WeatherNext imported successfully')"

echo ""
echo "=== Setup complete ==="
echo "Next steps:"
echo "  1. source venv/bin/activate"
echo "  2. Copy .env.example to .env and fill in API keys"
echo "  3. Download model weights: python scripts/download_weights.py"
echo "  4. Run pipeline: python pipeline.py"