@echo off
setlocal enabledelayedexpansion

echo ================================================================
echo DANCE-ST Neural Model Training for All CMAPSS Datasets
echo ================================================================
echo.

REM Create necessary directories
if not exist DANCEST_model\models\saved mkdir DANCEST_model\models\saved
if not exist DANCEST_model\models\tuning mkdir DANCEST_model\models\tuning
if not exist DANCEST_model\results mkdir DANCEST_model\results
if not exist DANCEST_model\logs mkdir DANCEST_model\logs

REM Check if a data directory was specified
set DATA_DIR=data
if not "%~1"=="" set DATA_DIR=%~1

REM Set timestamp for log files
set "TIMESTAMP=%date:~-4,4%%date:~-10,2%%date:~-7,2%_%time:~0,2%%time:~3,2%%time:~6,2%"
set "TIMESTAMP=%TIMESTAMP: =0%"

echo Training will be logged to: DANCEST_model\logs\cmapss_training_%TIMESTAMP%.log
echo Looking for CMAPSS data in: %DATA_DIR%

REM Step 1: Train neural models for all datasets
echo.
echo Step 1: Training neural models for all CMAPSS datasets...
echo This may take some time. See log file for detailed progress.

python DANCEST_model\train_cmapss_neural.py --datasets FD001 FD002 FD003 FD004 --verbose --data_dir %DATA_DIR% > DANCEST_model\logs\cmapss_training_%TIMESTAMP%.log 2>&1
if errorlevel 1 (
  echo Error in training neural models. See DANCEST_model\logs\cmapss_training_%TIMESTAMP%.log for details.
  exit /b 1
)

REM Check for trained models
echo.
echo ================================================================
echo                   Training Complete
echo ================================================================
echo.
echo Checking for trained models:

for %%d in (FD001 FD002 FD003 FD004) do (
  set "FOUND_MODEL="
  for /f "delims=" %%f in ('dir /b "DANCEST_model\models\saved\cmapss_%%d_model_*.keras" 2^>nul') do (
    set "FOUND_MODEL=%%f"
  )
  
  if defined FOUND_MODEL (
    echo - Dataset %%d: Model trained successfully - !FOUND_MODEL!
  ) else (
    echo - Dataset %%d: No model found
  )
)

echo.
echo Training complete! You can now run the models with:
echo DANCEST_model\run_cmapss_agents.bat [dataset] [cycle]
echo.
echo Example: DANCEST_model\run_cmapss_agents.bat FD002 120
echo.

endlocal 