#!/bin/bash
# Setup Training Environment för AI Produktionsmodell
# Zoomer-submissions granskning

set -e

echo "=========================================="
echo "AI Training Environment Setup"
echo "=========================================="

# Konfiguration
PROJECT_DIR="/home/bernt/.openclaw/workspace/iom/ai_pipeline"
DATASET_DIR="/tmp/quixzoom_production_dataset"
MODELS_DIR="/tmp/quixzoom_production_models"
VENV_DIR="/tmp/quixzoom_training_venv"

# Skapa kataloger
echo ""
echo "1. Creating directories..."
mkdir -p "$DATASET_DIR"/{images,labels}/{train,val,test}
mkdir -p "$MODELS_DIR"
mkdir -p "$PROJECT_DIR"/logs

# Kontrollera Python
echo ""
echo "2. Checking Python..."
python3 --version

# Kontrollera tillgängliga bibliotek
echo ""
echo "3. Checking installed packages..."
python3 -c "
import sys
packages = ['torch', 'torchvision', 'ultralytics', 'transformers', 'pillow', 'opencv-python', 'numpy', 'pandas']
for pkg in packages:
    try:
        __import__(pkg.replace('-', '_'))
        print(f'  ✓ {pkg}')
    except ImportError:
        print(f'  ✗ {pkg} (missing)')
        sys.exit(1)
"

# Kontrollera GPU
echo ""
echo "4. Checking GPU availability..."
python3 -c "
import torch
if torch.cuda.is_available():
    print(f'  ✓ CUDA available: {torch.cuda.get_device_name(0)}')
    print(f'  ✓ GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
else:
    print('  ⚠ No GPU available - training will be slow')
    print('  💡 Consider using CPU with reduced batch size')
"

# Kontrollera diskutrymme
echo ""
echo "5. Checking disk space..."
DF_OUTPUT=$(df -h "$DATASET_DIR" | tail -1)
echo "  $DF_OUTPUT"

# Kontrollera minne
echo ""
echo "6. Checking memory..."
FREE_OUTPUT=$(free -h | grep Mem)
echo "  $FREE_OUTPUT"

# Ladda ner YOLOv8 modell
echo ""
echo "7. Downloading YOLOv8 model..."
if [ ! -f "$PROJECT_DIR/yolov8n.pt" ]; then
    python3 -c "from ultralytics import YOLO; YOLO('yolov8n.pt')"
    echo "  ✓ YOLOv8n downloaded"
else
    echo "  ✓ YOLOv8n already exists"
fi

# Ladda ner CLIP modell
echo ""
echo "8. Downloading CLIP model..."
python3 -c "
from transformers import CLIPModel, CLIPProcessor
try:
    CLIPModel.from_pretrained('openai/clip-vit-base-patch32')
    CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
    print('  ✓ CLIP downloaded')
except Exception as e:
    print(f'  ✗ CLIP download failed: {e}')
"

# Skapa data.yaml template
echo ""
echo "9. Creating data.yaml template..."
cat > "$DATASET_DIR/data.yaml" << 'EOF'
path: /tmp/quixzoom_production_dataset
train: images/train
val: images/val
test: images/test

nc: 20
names:
  - street_light
  - traffic_sign
  - bench
  - trash_can
  - sidewalk
  - road
  - building
  - bridge
  - tree
  - graffiti
  - pothole
  - crack
  - corrosion
  - broken_glass
  - missing_parts
  - water_damage
  - vegetation_overgrowth
  - illegal_dumping
  - broken_pavement
  - faded_markings
EOF

echo "  ✓ data.yaml created"

# Skapa träningslogg
echo ""
echo "10. Creating training log..."
cat > "$PROJECT_DIR/logs/training_setup.log" << EOF
Training Environment Setup
==========================
Date: $(date -u +"%Y-%m-%d %H:%M:%S UTC")
Host: $(hostname)
Python: $(python3 --version)
PyTorch: $(python3 -c "import torch; print(torch.__version__)")
CUDA: $(python3 -c "import torch; print('Available' if torch.cuda.is_available() else 'Not available')")

Directories:
  Dataset: $DATASET_DIR
  Models: $MODELS_DIR
  Logs: $PROJECT_DIR/logs

Status: READY
EOF

echo "  ✓ Training log created"

echo ""
echo "=========================================="
echo "Setup Complete!"
echo "=========================================="
echo ""
echo "Next steps:"
echo "  1. Place real images in: $DATASET_DIR/images/train/"
echo "  2. Place annotations in: $DATASET_DIR/labels/train/"
echo "  3. Run: python3 $PROJECT_DIR/scripts/train_production_model.py"
echo ""
echo "Dataset structure:"
echo "  $DATASET_DIR/"
echo "    images/"
echo "      train/     # 70% of data"
echo "      val/       # 15% of data"
echo "      test/      # 15% of data"
echo "    labels/"
echo "      train/     # YOLO format .txt files"
echo "      val/"
echo "      test/"
echo "    data.yaml    # Dataset configuration"
echo ""