install.sh

#!/bin/bash
# ============================================================================
# ERA5 Parallel Downloader - Installation Script
# ============================================================================
#
# This script sets up the ERA5 downloader in your system
#
# Usage:
#   bash install.sh
#
# ============================================================================

set -e  # Exit on any error

echo "============================================================================"
echo "ERA5 Parallel Downloader - Installation"
echo "============================================================================"
echo ""

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# ----------------------------------------------------------------------------
# Step 1: Check Python version
# ----------------------------------------------------------------------------

echo -e "${BLUE}[1/6]${NC} Checking Python version..."

if ! command -v python3 &> /dev/null; then
    echo -e "${RED}✗${NC} Python 3 not found. Please install Python 3.8 or higher."
    exit 1
fi

PYTHON_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
echo -e "${GREEN}✓${NC} Found Python ${PYTHON_VERSION}"

# Check if version is 3.8 or higher
MAJOR=$(echo $PYTHON_VERSION | cut -d. -f1)
MINOR=$(echo $PYTHON_VERSION | cut -d. -f2)

if [ "$MAJOR" -lt 3 ] || ([ "$MAJOR" -eq 3 ] && [ "$MINOR" -lt 8 ]); then
    echo -e "${RED}✗${NC} Python 3.8+ required (found ${PYTHON_VERSION})"
    exit 1
fi

# ----------------------------------------------------------------------------
# Step 2: Create directory structure
# ----------------------------------------------------------------------------

echo -e "${BLUE}[2/6]${NC} Setting up directory structure..."

TARGET_DIR="/DATA/datasets/era5"

# Check if directory exists
if [ -d "$TARGET_DIR" ]; then
    echo -e "${YELLOW}⚠${NC} Directory already exists: $TARGET_DIR"
    read -p "   Continue and use existing directory? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "   Installation cancelled."
        exit 1
    fi
else
    echo "   Creating: $TARGET_DIR"
    mkdir -p "$TARGET_DIR"
fi

# Create subdirectories
mkdir -p "$TARGET_DIR/pressure_levels"
mkdir -p "$TARGET_DIR/single_levels"
mkdir -p "$TARGET_DIR/logs"
mkdir -p "$TARGET_DIR/temp"

echo -e "${GREEN}✓${NC} Directory structure created"

# ----------------------------------------------------------------------------
# Step 3: Install Python dependencies
# ----------------------------------------------------------------------------

echo -e "${BLUE}[3/6]${NC} Installing Python dependencies..."

# Check if pip is available
if ! command -v pip3 &> /dev/null; then
    echo -e "${RED}✗${NC} pip3 not found. Please install pip."
    exit 1
fi

echo "   Installing required packages..."

# Install core dependencies
pip3 install --user cdsapi pyyaml

# Install optional dependencies if available
echo "   Installing optional packages (for enhanced features)..."
pip3 install --user tqdm colorlog 2>/dev/null || echo "   (Some optional packages skipped)"

echo -e "${GREEN}✓${NC} Dependencies installed"

# ----------------------------------------------------------------------------
# Step 4: Copy files to target directory
# ----------------------------------------------------------------------------

echo -e "${BLUE}[4/6]${NC} Installing ERA5 downloader files..."

# Get script directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

# Copy main scripts
cp "$SCRIPT_DIR/era5_downloader.py" "$TARGET_DIR/"
cp "$SCRIPT_DIR/validate_config.py" "$TARGET_DIR/"

# Copy configuration files
cp "$SCRIPT_DIR/config_wrf.yaml" "$TARGET_DIR/"
cp "$SCRIPT_DIR/config_surface_only.yaml" "$TARGET_DIR/"
cp "$SCRIPT_DIR/config_year_2025.yaml" "$TARGET_DIR/"

# Copy documentation
cp "$SCRIPT_DIR/README.txt" "$TARGET_DIR/"
cp "$SCRIPT_DIR/requirements.txt" "$TARGET_DIR/"

# Copy SLURM script if on HPC
if command -v sbatch &> /dev/null; then
    cp "$SCRIPT_DIR/slurm_download_era5.sh" "$TARGET_DIR/"
    chmod +x "$TARGET_DIR/slurm_download_era5.sh"
    echo -e "${GREEN}✓${NC} SLURM job script installed"
fi

# Make scripts executable
chmod +x "$TARGET_DIR/era5_downloader.py"
chmod +x "$TARGET_DIR/validate_config.py"

echo -e "${GREEN}✓${NC} Files installed to $TARGET_DIR"

# ----------------------------------------------------------------------------
# Step 5: Check CDS API setup
# ----------------------------------------------------------------------------

echo -e "${BLUE}[5/6]${NC} Checking CDS API configuration..."

CDSAPIRC="$HOME/.cdsapirc"

if [ -f "$CDSAPIRC" ]; then
    echo -e "${GREEN}✓${NC} Found CDS API configuration: $CDSAPIRC"
    
    # Test connection
    echo "   Testing CDS API connection..."
    if python3 -c "import cdsapi; cdsapi.Client()" 2>/dev/null; then
        echo -e "${GREEN}✓${NC} CDS API connection successful"
    else
        echo -e "${YELLOW}⚠${NC} CDS API connection failed - check your credentials"
    fi
else
    echo -e "${YELLOW}⚠${NC} CDS API not configured"
    echo ""
    echo "   To set up CDS API:"
    echo "   1. Register at: https://cds.climate.copernicus.eu/"
    echo "   2. Get your API key from: https://cds.climate.copernicus.eu/user"
    echo "   3. Create $CDSAPIRC with:"
    echo ""
    echo "      url: https://cds.climate.copernicus.eu/api"
    echo "      key: YOUR_UID:YOUR_API_KEY"
    echo ""
    echo "   4. Set permissions: chmod 600 $CDSAPIRC"
    echo ""
fi

# ----------------------------------------------------------------------------
# Step 6: Validate installation
# ----------------------------------------------------------------------------

echo -e "${BLUE}[6/6]${NC} Validating installation..."

cd "$TARGET_DIR"

# Test if scripts are executable
if [ -x "era5_downloader.py" ] && [ -x "validate_config.py" ]; then
    echo -e "${GREEN}✓${NC} Scripts are executable"
else
    echo -e "${RED}✗${NC} Scripts are not executable"
    exit 1
fi

# Test if configurations are readable
if [ -r "config_wrf.yaml" ]; then
    echo -e "${GREEN}✓${NC} Configuration files are readable"
else
    echo -e "${RED}✗${NC} Configuration files not found"
    exit 1
fi

# Test dry run
echo "   Testing dry run..."
if python3 era5_downloader.py -c config_wrf.yaml --dry-run > /dev/null 2>&1; then
    echo -e "${GREEN}✓${NC} Dry run successful"
else
    echo -e "${YELLOW}⚠${NC} Dry run test failed (may be due to CDS API setup)"
fi

# ----------------------------------------------------------------------------
# Installation complete
# ----------------------------------------------------------------------------

echo ""
echo "============================================================================"
echo -e "${GREEN}Installation Complete!${NC}"
echo "============================================================================"
echo ""
echo "Installation directory: $TARGET_DIR"
echo ""
echo "📋 Next Steps:"
echo ""
echo "1. Review and edit configuration:"
echo "   cd $TARGET_DIR"
echo "   nano config_wrf.yaml"
echo ""
echo "2. Validate your configuration:"
echo "   python3 validate_config.py config_wrf.yaml"
echo ""
echo "3. Start downloading:"
echo "   python3 era5_downloader.py -c config_wrf.yaml"
echo ""
echo "4. For full year 2025:"
echo "   python3 era5_downloader.py -c config_year_2025.yaml"
echo ""

# Print SLURM instructions if applicable
if command -v sbatch &> /dev/null; then
    echo "5. For HPC/SLURM systems:"
    echo "   sbatch slurm_download_era5.sh config_wrf.yaml"
    echo ""
fi

echo "📚 Documentation: $TARGET_DIR/README.txt"
echo ""
echo "============================================================================"