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

echo "=========================================="
echo "  Quantum Syndrome Decoder - Setup & Train"
echo "=========================================="

# --- 1. Environment Setup ---
PYTHON="${PYTHON:-python3}"
VENV_DIR="${VENV_DIR:-.venv}"

if [ ! -d "$VENV_DIR" ]; then
    echo "[1/5] Creating virtual environment..."
    $PYTHON -m venv "$VENV_DIR"
fi

echo "[1/5] Activating virtual environment..."
source "$VENV_DIR/bin/activate"

echo "[2/5] Installing dependencies..."
pip install --upgrade pip setuptools wheel -q
pip install torch numpy stim pymatching -q

# Verify installations
echo "    torch:     $(python -c 'import torch; print(torch.__version__)')"
echo "    numpy:     $(python -c 'import numpy; print(numpy.__version__)')"
echo "    stim:      $(python -c 'import stim; print(stim.__version__)')"
echo "    pymatching:$(python -c 'import pymatching; print(pymatching.__version__)')"

# --- 2. Data Generation ---
echo "[3/5] Generating synthetic syndrome dataset..."
python generate_data.py

# --- 3. Training ---
echo "[4/5] Training supervised 3D CNN decoder..."
python train_decoder.py \
    --mode supervised \
    --distance 5 \
    --epochs 50 \
    --batch_size 128 \
    --lr 5e-4 \
    --channels 128 \
    --layers 4 \
    --data_dir data \
    --output_dir outputs

# --- 4. Evaluation ---
echo "[5/5] Evaluating trained model..."
python -c "
import torch
import numpy as np
import json
from decoder_model import QuantumSyndromeDecoder
from train_decoder import SyndromeDataset, DataLoader, evaluate_decoder

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'    Evaluation device: {device}')

model = QuantumSyndromeDecoder(distance=5, channels=128, num_layers=4)
checkpoint = torch.load('outputs/decoder_best.pt', map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
model = model.to(device)

syndromes = np.load('data/syndromes.npy')
errors = np.load('data/errors.npy')
n = len(syndromes)
test_dataset = SyndromeDataset(syndromes[int(0.9*n):], errors[int(0.9*n):])
test_loader = DataLoader(test_dataset, batch_size=64)

results = evaluate_decoder(model, test_loader, device)

# Save detailed results
with open('outputs/final_results.json', 'w') as f:
    json.dump(results, f, indent=2)

print(f'    Precision: {results[\"precision\"]:.4f}')
print(f'    Recall:    {results[\"recall\"]:.4f}')
print(f'    F1 Score:  {2*results[\"precision\"]*results[\"recall\"]/(results[\"precision\"]+results[\"recall\"]+1e-8):.4f}')
"

echo "=========================================="
echo "  Training Complete!"
echo "  Model: outputs/decoder_best.pt"
echo "  Results: outputs/final_results.json"
echo "=========================================="