#!/bin/bash

# =================================================================
# Quick Start Script for FedNewton Experiments
# Usage: ./run.sh
# =================================================================

# ===========================
# 1. Model Selection (Uncomment to activate)
# ===========================

MODEL_NAME="resnet18" # Default: Lightweight, good for debugging
# MODEL_NAME="tv_resnet50"
# MODEL_NAME="vit_base_patch16_clip_224"      # CLIP (OpenAI)
# MODEL_NAME="vit_base_patch14_dinov2.lvd142m" # DINOv2 (Meta AI)


# ===========================
# 2. Dataset Selection (Uncomment to activate)
# ===========================
DATASET="cifar10"      # Default: Standard benchmark
# DATASET="fashionmnist" # Simple benchmark
# DATASET="mnist"        # Basic sanity check


# ===========================
# 3. Hyperparameters Configuration
# ===========================
METHOD="FedNewton"      # Options: FedAvg, FedProx, FedNewton, FedSophia, etc.
TRAIN_MODE="head"       # 'full': Full fine-tuning, 'head': Train head only
ROUNDS=10                # Number of communication rounds
BATCH_SIZE=64
NEWTON_LR=0.1                 # Learning rate
NEWTON_DAMPING=0.0001       # Lambda factor


# ===========================
# 4. Directory Setup & Execution
# ===========================

# Generate unique directory based on timestamp and config
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
RUN_ROOT="runs/run_${TIMESTAMP}_${MODEL_NAME}_${TRAIN_MODE}"
CURRENT_LOG_DIR="${RUN_ROOT}/${DATASET}/logs"
CURRENT_FIG_DIR="${RUN_ROOT}/${DATASET}/figs"

echo "----------------------------------------------------------------"
echo "Starting Experiment..."
echo "Model: ${MODEL_NAME} | Dataset: ${DATASET} | Method: ${METHOD}"
echo "Logs will be saved to: ${CURRENT_LOG_DIR}"
echo "----------------------------------------------------------------"

# Ensure directories exist (optional, python handles this usually, but good practice)
mkdir -p "${CURRENT_LOG_DIR}"
mkdir -p "${CURRENT_FIG_DIR}"

# Run the training script
python main.py \
  --dataset "${DATASET}" \
  --training_mode "${TRAIN_MODE}" \
  --model_name "${MODEL_NAME}" \
  --method "${METHOD}" \
  --rounds "${ROUNDS}" \
  --batch_size "${BATCH_SIZE}" \
  --newton_lr $NEWTON_LR \
  --damping $NEWTON_DAMPING \
  --log_dir "${CURRENT_LOG_DIR}"

# Check if training succeeded before running analysis
if [ $? -eq 0 ]; then
    echo "Training finished. Starting analysis..."
    python "./analysis.py" --log_dir "${CURRENT_LOG_DIR}" --fig_dir "${CURRENT_FIG_DIR}"
    echo "Done! Check results in ${RUN_ROOT}"
else
    echo "Experiment failed. Analysis skipped."
    exit 1
fi