# Curriculum Learning Configuration - MEDIUM MODEL
# Balanced configuration with good performance/memory trade-off

# Model configuration - BALANCED CAPACITY
model:
  target_dim: 18  # Will be overridden based on dataset
  is_unconditional: 0
  timeemb: 256    # Increased from 128 → 256 (2x larger)
  featureemb: 64  # Increased from 16 → 64 (4x larger)
  target_strategy: "random"
  use_aux_loss: False
  aux_weight_order1: 0.4
  aux_weight_order2: 0.4
  aux_loss_normalize: False  # CRITICAL FIX: Disable dangerous normalization
  aux_loss_max_value: 2.0

# Training configuration
training:
  seed: 1
  batch_size: 256  # Reduced for stability with larger model
  validation_split: 0.05
  num_workers: 4
  mask_generation_count: 100
  gradient_clip: 0.5   # Tighter clipping to prevent gradient explosion
  weight_decay: 1e-4  # CRITICAL FIX: Proper weight decay for Adam
  ratio: 0.7
  checkpoint_save_step: 10
  use_mixed_precision: False  # CRITICAL FIX: Disable mixed precision for stability

# Dataset configuration
dataset:
  window_length: 100
  split: 10
  mask_ratio: 0.5
  scale_factor: 1

# Curriculum Learning Parameters
curriculum:
  phase1_epochs: 30
  phase2_epochs: 30

  # Phase 1: Easy - Reduced learning rates for gradient stability
  mask_ratio_phase1_start: 0.1
  mask_ratio_phase1_end: 0.3
  noise_ratio_phase1_start: 0.0001
  noise_ratio_phase1_end: 0.1
  lr_phase1_start: 2e-4  # CRITICAL FIX: Conservative, stable LR
  lr_phase1_end: 2e-4    # CRITICAL FIX: Fixed LR prevents momentum disruption

  # Phase 2: Medium - Reduced learning rates for gradient stability
  mask_ratio_phase2_start: 0.3
  mask_ratio_phase2_end: 0.6
  noise_ratio_phase2_start: 0.1
  noise_ratio_phase2_end: 0.3
  lr_phase2_start: 2e-4    # Keep consistent with phase 1
  lr_phase2_end: 2e-4      # Keep consistent with phase 1

  # Phase 3: Hard - Reduced learning rates for gradient stability
  mask_ratio_phase3_start: 0.6
  mask_ratio_phase3_end: 0.8
  noise_ratio_phase3_start: 0.3
  noise_ratio_phase3_end: 0.5
  lr_phase3_start: 2e-4    # Keep consistent with phase 1
  lr_phase3_end: 2e-4      # Keep consistent with phase 1

# Diffusion configuration - BALANCED CAPACITY
diffusion:
  layers: 6              # Increased from 4 → 6 (1.5x more layers)
  channels: 128          # Increased from 64 → 128 (2x larger)
  nheads: 8              # Keep at 8 for divisibility (128/8 = 16)
  diffusion_embedding_dim: 256  # Increased from 128 → 256 (2x larger)
  beta_start: 0.0001
  beta_end: 0.01
  num_steps: 500
  schedule: "linear"

# Expected parameter count: ~2.4M parameters (good balance)
# Fixes applied for gradient stability:
# - Reduced learning rates by 50% across all phases
# - Tighter gradient clipping (1.0 → 0.5)
# - Reduced batch size (384 → 256) for stability