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

# Pre-install Python packages for the NBCU interview environment.
# Run this BEFORE the interview — torch alone is ~2.5GB.

if [ ! -d ".venv" ]; then
    echo "No .venv found. Enter the devshell first: nix develop"
    exit 1
fi

source .venv/bin/activate

echo "Installing core ML stack (trying ROCm torch first, CPU fallback)..."
echo "  Attempting ROCm 6.4 wheels..."
if uv pip install \
    torch \
    torchvision \
    --index-url https://download.pytorch.org/whl/rocm6.4; then
    echo "  ROCm wheels installed. Checking GPU visibility..."
    if python -c "import torch; print(f'  GPU available: {torch.cuda.is_available()}'); assert torch.cuda.is_available()"; then
        echo "  ROCm torch with GPU acceleration ready."
    else
        echo ""
        echo "  ROCm torch installed but GPU not detected."
        echo "  To debug GPU detection later:"
        echo "    - Check /dev/kfd exists and is accessible"
        echo "    - Check user is in 'video' and 'render' groups"
        echo "    - Run: HSA_OVERRIDE_GFX_VERSION=10.3.0 python -c \"import torch; print(torch.cuda.is_available())\""
        echo ""
    fi
else
    echo "  ROCm wheel install failed. Falling back to CPU torch."
    uv pip install torch torchvision
fi

uv pip install \
    numpy \
    scipy \
    scikit-learn \
    matplotlib \
    Pillow

echo ""
echo "Installing geospatial stack..."
uv pip install \
    rasterio \
    laspy[lazrs] \
    opencv-python-headless \
    shapely \
    geopandas \
    fiona

echo ""
echo "Installing extras (seaborn for confusion matrix plots, etc.)..."
uv pip install \
    seaborn \
    h5py

echo ""
echo "Verifying critical imports..."
python -c "
import torch
print(f'  torch {torch.__version__}')
if torch.cuda.is_available():
    print(f'    GPU: {torch.cuda.get_device_name(0)}')
    print(f'    ROCm/CUDA: {torch.version.hip or torch.version.cuda}')
else:
    print(f'    CPU only')

import torchvision
print(f'  torchvision {torchvision.__version__}')

import numpy
print(f'  numpy {numpy.__version__}')

import scipy
print(f'  scipy {scipy.__version__}')

import sklearn
print(f'  sklearn {sklearn.__version__}')

import rasterio
print(f'  rasterio {rasterio.__version__}')

import laspy
print(f'  laspy {laspy.__version__}')

import cv2
print(f'  opencv {cv2.__version__}')

import shapely
print(f'  shapely {shapely.__version__}')

print()
print('All imports OK.')
"

echo ""
echo "Smoke-testing torch inference..."
python -c "
import torch
from torchvision.models.segmentation import deeplabv3_resnet50, DeepLabV3_ResNet50_Weights

weights = DeepLabV3_ResNet50_Weights.DEFAULT
model = deeplabv3_resnet50(weights=weights)
model.eval()

x = torch.randint(0, 256, (3, 256, 256), dtype=torch.uint8)
x = weights.transforms()(x).unsqueeze(0)

with torch.no_grad():
    out = model(x)['out']

print(f'  DeepLabV3 inference OK — output shape: {out.shape}')
print(f'  Predicted classes: {torch.unique(out.argmax(dim=1)).tolist()}')
"

echo ""
echo "Smoke-testing rasterio + laspy..."
python -c "
import rasterio
import numpy as np
from rasterio.transform import from_bounds

# Write and read a tiny GeoTIFF
data = np.random.randint(0, 255, (3, 64, 64), dtype=np.uint8)
transform = from_bounds(-73.58, 45.49, -73.55, 45.52, 64, 64)

with rasterio.open(
    '/tmp/_test_rio.tif', 'w', driver='GTiff',
    height=64, width=64, count=3, dtype='uint8',
    crs='EPSG:4326', transform=transform,
) as dst:
    dst.write(data)

with rasterio.open('/tmp/_test_rio.tif') as src:
    print(f'  rasterio read/write OK — CRS: {src.crs}, shape: {src.shape}')

import laspy
header = laspy.LasHeader(point_format=0, version='1.2')
header.offsets = [0, 0, 0]
header.scales = [0.001, 0.001, 0.001]
las = laspy.LasData(header)
las.x = np.random.uniform(0, 100, 1000)
las.y = np.random.uniform(0, 100, 1000)
las.z = np.random.uniform(0, 50, 1000)
las.write('/tmp/_test_las.las')
las2 = laspy.read('/tmp/_test_las.las')
print(f'  laspy read/write OK — {las2.header.point_count} points')

import os
os.remove('/tmp/_test_rio.tif')
os.remove('/tmp/_test_las.las')
"

echo ""
echo "========================================="
echo "          Environment ready."
echo "========================================="
echo ""