# flow_matching/src/debug_config.yml

# Global settings
out_dir: output/debug_run
seed: 42
overwrite: true
device: cpu # Use CPU for local debugging
batch_size: 2

# Stage 1: Mean Anchor Generation (MultiSubjectConvLinearEncoder)
stage1:
  epochs: 1
  lr: 1e-3
  weight_decay: 0.0
  
  model:
    embed_dim: 16
    encoder_kernel_size: 3
    decoder_kernel_size: 0
    hidden_model: null
    global_pool: avg
    encoder_causal: false
    encoder_positive: false
    encoder_blockwise: false
    pool_num_heads: 2
    with_shared_decoder: true
    with_subject_decoders: true

  transformer:
    num_heads: 2
    depth: 1
    mlp_ratio: 2.0
  conv1dnext:
    depth: 1
    kernel_size: 3
    causal: false

# Stage 2: Neural Vector Field (Flow Matching)
stage2:
  epochs: 1
  lr: 3e-4
  weight_decay: 0.0
  latent_dim: 32
  n_timesteps: 25

  # CFM Configuration
  cfm:
    solver: euler
    sigma_min: 1e-4

  # Decoder Architecture (Matcha-TTS style)
  decoder:
    # in_channels: 100 (Mock Voxel dim)
    # out_channels: 100
    channels: [32, 32]
    dropout: 0.0
    attention_head_dim: 16
    n_blocks: 1
    num_mid_blocks: 1
    num_heads: 2
    act_fn: snakebeta
    down_block_type: transformer
    mid_block_type: transformer
    up_block_type: transformer

# Dataset Configuration
subjects: [1, 2, 3, 5]

# Mock features list matches what the code expects to parse keys from
features:
  mock_feat_1:
    model: mock_model_1
    layers:
      layer1: layer1
  mock_feat_2:
    model: mock_model_2
    layers:
      layer2: layer2

include_features:
  - mock_feat_1/layer1
  - mock_feat_2/layer2

datasets:
  train:
    filter:
      seasons: []
      movies: []
    sample_length: 10
    num_samples: 4
    shuffle: True
    seed: 42

  val_debug:
    filter:
      seasons: []
      movies: []
    sample_length: null
    num_samples: null
    shuffle: false

val_set_name: val_debug
datasets_root: null