#!/bin/bash

# Default values
TJA_DIR="datas/tja"
MEASURE_CSV="datas/measure.csv"
DATASET_CSV="datas/dataset.csv"
THRESHOLD=0.1
TRAIN_COUNT=0 # 0 means all
MAX_EPOCHS=5000
MODEL_PATH="model/constant_predictor.joblib"

# Parse arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --tja_dir) TJA_DIR="$2"; shift ;;
        --measure_csv) MEASURE_CSV="$2"; shift ;;
        --threshold) THRESHOLD="$2"; shift ;;
        --train_count) TRAIN_COUNT="$2"; shift ;;
        --max_epochs) MAX_EPOCHS="$2"; shift ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

ITERATIONS=${ITERATIONS:-50}
BATCH_SIZE=${BATCH_SIZE:-200}

echo "--- Step 1: TJA Factorization ---"
bun run src/factorize.ts "$TJA_DIR" "$MEASURE_CSV" "$DATASET_CSV"

echo "--- Step 2: Incremental Training ---"
python3 model/train.py --dataset "$DATASET_CSV" --model_path "$MODEL_PATH" --batch_size "$BATCH_SIZE" --iterations "$ITERATIONS"

if [ $? -ne 0 ]; then
    echo "Training failed."
    exit 1
fi

echo "--- Step 3: Model Validation (Sample) ---"
# Pick a random TJA from the dir to see prediction
SAMPLE_TJA=$(ls $TJA_DIR/*.tja | head -n 1)
echo "Testing on $SAMPLE_TJA..."
bun run src/predict.ts "$MODEL_PATH" "$SAMPLE_TJA" "oni"

echo "--- Pipeline Complete ---"