#!/bin/bash
set -e

echo "==================================="
echo "PVMD - Parakeet Romanian Fine-tuning"
echo "==================================="

# Activare mediu
source /opt/conda/etc/profile.d/conda.sh
conda activate pvmd

# Variabile
export HF_HOME=/mnt/parakeet-training/cache/huggingface
export HF_DATASETS_CACHE=/mnt/parakeet-training/cache/datasets
export PYSPARK_PYTHON=/opt/conda/envs/pvmd/bin/python
export PYSPARK_DRIVER_PYTHON=/opt/conda/envs/pvmd/bin/python

DATASET_NAME="datadriven-company/TTS-Romanian"
MODEL_PATH="/mnt/parakeet-training/models/parakeet-romanian/SpeD-ParakeetRo_110M_TDT-CTC.nemo"
AUDIO_DIR="/mnt/parakeet-training/datasets/audio_wav"
MANIFEST_DIR="/mnt/parakeet-training/datasets/manifests"
OUTPUT_DIR="/mnt/parakeet-training/outputs/nemo-finetune"

mkdir -p "$AUDIO_DIR" "$MANIFEST_DIR" "$OUTPUT_DIR"

echo "[1/4] Extragere audio din dataset..."
python src/extract_audio_for_nemo.py \
  --dataset_name "$DATASET_NAME" \
  --output_audio_dir "$AUDIO_DIR" \
  --output_manifest "$MANIFEST_DIR/full_manifest.json" \
  --target_sr 16000 \
  --cache_dir /mnt/parakeet-training/cache/datasets

echo "[2/4] Împărțire train/validation..."
python -c "
import json, random
random.seed(42)
with open('$MANIFEST_DIR/full_manifest.json') as f:
    lines = f.readlines()
random.shuffle(lines)
split = int(len(lines) * 0.95)
train = lines[:split]
val = lines[split:]
with open('$MANIFEST_DIR/train_manifest.json', 'w') as f:
    f.writelines(train)
with open('$MANIFEST_DIR/val_manifest.json', 'w') as f:
    f.writelines(val)
print(f'Train: {len(train)}, Val: {len(val)}')
"

echo "[3/4] Preprocesare PySpark (opțional - filtre avansate)..."
# python src/preprocess_nemo_manifest.py ...

echo "[4/4] Antrenament NeMo..."
python src/train_nemo.py \
  --model_path "$MODEL_PATH" \
  --train_manifest "$MANIFEST_DIR/train_manifest.json" \
  --val_manifest "$MANIFEST_DIR/val_manifest.json" \
  --output_dir "$OUTPUT_DIR" \
  --max_epochs 10 \
  --batch_size 16 \
  --lr 2e-3 \
  --warmup_steps 5000 \
  --precision bf16-mixed \
  --accumulate_grad_batches 4

echo "==================================="
echo "Antrenament finalizat!"
echo "Output: $OUTPUT_DIR"
echo "==================================="