#!/bin/bash

# Parallel hyperparameter tuning across 8 GPUs
# This script distributes experiments evenly across all available GPUs

clear

# Activate conda environment
source ~/miniconda3/etc/profile.d/conda.sh
conda activate /home/ec2-user/aev

# Configuration
DATASET_TYPE="pickapic"  # "coco" or "pickapic"
MODEL_VARIANT="lpo"      # "origin", "spo", "diffusion_dpo", or "lpo"
MAX_SAMPLES=500           # Number of samples for tuning
NUM_STEPS=50             # Fixed inference steps
SEARCH_TYPE="grid"       # "grid" or "random"
OUTPUT_DIR="RESULTS_TURNING/run_2"
NUM_GPUS=8               # Number of GPUs to use

echo "=============================================="
echo "  PARALLEL HYPERPARAMETER TUNING"
echo "=============================================="
echo ""
echo "Configuration:"
echo "  Dataset: $DATASET_TYPE"
echo "  Model: $MODEL_VARIANT"
echo "  Samples: $MAX_SAMPLES"
echo "  Inference Steps: $NUM_STEPS"
echo "  Search Type: $SEARCH_TYPE"
echo "  GPUs: $NUM_GPUS"
echo "  Output: $OUTPUT_DIR"
echo ""

# First, calculate total number of experiments
echo "Calculating total experiments..."
TOTAL_CONFIGS=$(python -c "
from tune_hyperparams import HyperparameterTuner
import sys
tuner = HyperparameterTuner()
configs = tuner.define_search_space()
sys.stderr.write(f'Generated {len(configs)} configurations\n')
print(len(configs))
" 2>&1 | tail -1)

echo "Total configurations: $TOTAL_CONFIGS"
echo ""

# Calculate experiments per GPU
CONFIGS_PER_GPU=$((TOTAL_CONFIGS / NUM_GPUS))
REMAINDER=$((TOTAL_CONFIGS % NUM_GPUS))

echo "Distributing work:"
echo "  Base configs per GPU: $CONFIGS_PER_GPU"
echo "  Extra configs for first GPUs: $REMAINDER"
echo ""

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Array to store background process IDs
PIDS=()

# Launch parallel processes on each GPU
for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do
    # Calculate start and end indices for this GPU
    START_IDX=$((GPU_ID * CONFIGS_PER_GPU))
    
    # Give extra configs to first GPUs
    if [ $GPU_ID -lt $REMAINDER ]; then
        START_IDX=$((START_IDX + GPU_ID))
        END_IDX=$((START_IDX + CONFIGS_PER_GPU + 1))
    else
        START_IDX=$((START_IDX + REMAINDER))
        END_IDX=$((START_IDX + CONFIGS_PER_GPU))
    fi
    
    # Create GPU-specific output directory
    GPU_OUTPUT_DIR="${OUTPUT_DIR}/gpu_${GPU_ID}"
    mkdir -p "$GPU_OUTPUT_DIR"
    
    echo "GPU $GPU_ID: configs $START_IDX to $END_IDX"
    
    # Launch tuning process in background
    nohup python tune_hyperparams.py \
        --output_dir "$GPU_OUTPUT_DIR" \
        --max_samples $MAX_SAMPLES \
        --num_steps $NUM_STEPS \
        --dataset_type "$DATASET_TYPE" \
        --model_variant "$MODEL_VARIANT" \
        --cuda $GPU_ID \
        --search_type "$SEARCH_TYPE" \
        --start_idx $START_IDX \
        --end_idx $END_IDX \
        --metrics clip aesthetic pickscore hpsv2 imagereward \
        > "${GPU_OUTPUT_DIR}/tuning.log" 2>&1 &
    
    # Store PID
    PIDS+=($!)
    
    echo "  Launched with PID: ${PIDS[$GPU_ID]}"
    
    # Small delay to avoid race conditions
    sleep 2
done

echo ""
echo "=============================================="
echo "  ALL PROCESSES LAUNCHED"
echo "=============================================="
echo ""
echo "Background processes running:"
for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do
    echo "  GPU $GPU_ID: PID ${PIDS[$GPU_ID]} -> ${OUTPUT_DIR}/gpu_${GPU_ID}/tuning.log"
done
echo ""
echo "To monitor progress:"
echo "  tail -f ${OUTPUT_DIR}/gpu_0/tuning.log"
echo "  tail -f ${OUTPUT_DIR}/gpu_1/tuning.log"
echo "  ... etc"
echo ""
echo "To check all GPU processes:"
echo "  ps aux | grep tune_hyperparams.py"
echo ""
echo "To monitor GPU usage:"
echo "  watch -n 1 nvidia-smi"
echo ""
echo "To kill all processes:"
echo "  kill ${PIDS[@]}"
echo ""
echo "Waiting for all processes to complete..."
echo "(Press Ctrl+C to stop waiting, processes will continue in background)"
echo ""

# Wait for all background processes
for PID in "${PIDS[@]}"; do
    wait $PID
done

echo ""
echo "=============================================="
echo "  ALL TUNING PROCESSES COMPLETE"
echo "=============================================="
echo ""

# Merge results from all GPUs
echo "Merging results from all GPUs..."

# Activate conda environment for Python script
source ~/miniconda3/etc/profile.d/conda.sh
conda activate /home/ec2-user/aev

python - <<'EOF'
import json
from pathlib import Path
import sys

output_dir = Path("RESULTS_TURNING")
all_results = []
baseline_result = None

# Collect results from each GPU
for gpu_id in range(8):
    gpu_dir = output_dir / f"gpu_{gpu_id}"
    results_file = gpu_dir / "tuning_results.json"
    
    if results_file.exists():
        with open(results_file, 'r') as f:
            data = json.load(f)
            
        # Get baseline (should be same from all)
        if baseline_result is None and "baseline" in data:
            baseline_result = data["baseline"]
        
        # Collect experiments
        if "experiments" in data:
            all_results.extend(data["experiments"])
        
        print(f"GPU {gpu_id}: {len(data.get('experiments', []))} results")

# Merge all results
merged_data = {
    "baseline": baseline_result,
    "experiments": all_results,
    "num_gpus": 8,
    "total_experiments": len(all_results)
}

# Save merged results
merged_file = output_dir / "merged_results.json"
with open(merged_file, 'w') as f:
    json.dump(merged_data, f, indent=2)

print(f"\nMerged {len(all_results)} total results")
print(f"Saved to: {merged_file}")

# Find best configuration
successful = [r for r in all_results if "metrics" in r]
if successful:
    # Compute aggregate scores
    def compute_score(metrics):
        weights = {
            "reward": 1.0, "clip": 0.8, "aesthetic": 0.8,
            "pickscore": 1.0, "hpsv2": 1.0, "imagereward": 1.0,
            "fid": -0.5
        }
        score = sum(weights.get(k, 0) * v for k, v in metrics.items())
        return score / sum(abs(w) for w in weights.values())
    
    for r in successful:
        r["aggregate_score"] = compute_score(r["metrics"])
    
    successful.sort(key=lambda x: x["aggregate_score"], reverse=True)
    
    best = successful[0]
    best_file = output_dir / "best_config.json"
    with open(best_file, 'w') as f:
        json.dump({
            "config": best["config"],
            "metrics": best["metrics"],
            "aggregate_score": best["aggregate_score"],
            "improvements": best.get("improvements", {})
        }, f, indent=2)
    
    print(f"\n{'='*60}")
    print("BEST CONFIGURATION:")
    print(f"{'='*60}")
    print(json.dumps(best["config"], indent=2))
    print(f"\nAggregate Score: {best['aggregate_score']:.4f}")
    print(f"Saved to: {best_file}")
else:
    print("\nNo successful experiments found!")
    sys.exit(1)
EOF

if [ $? -eq 0 ]; then
    echo ""
    echo "=============================================="
    echo "  TUNING COMPLETE!"
    echo "=============================================="
    echo ""
    echo "Results:"
    echo "  Merged results: ${OUTPUT_DIR}/merged_results.json"
    echo "  Best config: ${OUTPUT_DIR}/best_config.json"
    echo ""
    echo "View best configuration:"
    echo "  cat ${OUTPUT_DIR}/best_config.json"
    echo ""
else
    echo ""
    echo "ERROR: Failed to merge results"
    exit 1
fi