@echo off
REM ============================================================================
REM Atomizer Neural Network Training Script
REM ============================================================================
REM Trains a parametric neural network on collected FEA data
REM
REM Usage:
REM   train_neural.bat                          (uses default: bracket study)
REM   train_neural.bat study_name               (specify study name)
REM   train_neural.bat study_name 200           (specify epochs)
REM ============================================================================

echo.
echo ============================================================================
echo              ATOMIZER NEURAL NETWORK TRAINING
echo ============================================================================
echo.

REM Set defaults
set STUDY_NAME=bracket_stiffness_optimization_atomizerfield
set EPOCHS=100

REM Parse arguments
if not "%~1"=="" set STUDY_NAME=%~1
if not "%~2"=="" set EPOCHS=%~2

set TRAIN_DIR=atomizer_field_training_data\%STUDY_NAME%
set OUTPUT_DIR=atomizer-field\runs\%STUDY_NAME%

echo Study:     %STUDY_NAME%
echo Epochs:    %EPOCHS%
echo Train Dir: %TRAIN_DIR%
echo Output:    %OUTPUT_DIR%
echo.

REM Check training data exists
if not exist "%TRAIN_DIR%" (
    echo ERROR: Training data not found at %TRAIN_DIR%
    echo.
    echo Available training data directories:
    dir atomizer_field_training_data /b 2>nul
    echo.
    pause
    exit /b 1
)

REM Count trials
echo Counting training trials...
set /a count=0
for /d %%d in ("%TRAIN_DIR%\trial_*") do set /a count+=1
echo Found %count% trials
echo.

if %count% LSS 20 (
    echo WARNING: Less than 20 trials. Training may not converge well.
    echo Consider running more FEA simulations first.
    echo.
)

REM Check PyTorch
python -c "import torch; print(f'PyTorch {torch.__version__}')" 2>nul
if errorlevel 1 (
    echo ERROR: PyTorch not installed
    echo Run install.bat first
    pause
    exit /b 1
)

REM Check for GPU
echo.
echo Checking for GPU...
python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')"
echo.

REM Start training
echo ============================================================================
echo Starting training...
echo ============================================================================
echo.
echo Press Ctrl+C to stop training early (checkpoint will be saved)
echo.

cd atomizer-field
python train_parametric.py ^
    --train_dir "..\%TRAIN_DIR%" ^
    --epochs %EPOCHS% ^
    --output_dir "runs\%STUDY_NAME%" ^
    --batch_size 16 ^
    --learning_rate 0.001 ^
    --hidden_channels 128 ^
    --num_layers 4

if errorlevel 1 (
    echo.
    echo Training failed! Check error messages above.
    cd ..
    pause
    exit /b 1
)

cd ..

echo.
echo ============================================================================
echo                    TRAINING COMPLETE!
echo ============================================================================
echo.
echo Model saved to: %OUTPUT_DIR%\checkpoint_best.pt
echo.
echo To use the trained model in optimization:
echo   cd studies\%STUDY_NAME%
echo   python run_optimization.py --run --trials 100 --enable-nn --resume
echo.

pause