# 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: 1e-3
  weight_decay: 0.0
  n_timesteps: 10

  # CFM and training regularization
  cfm:
    solver: euler
    kld_weight: 1.0
    kld_target_std: 1.0
    detach_ut: false
    time_dist_shift: 1.0

  # DiT-style velocity model
  velocity_net:
    hidden_dim: 64
    modality_dims: [1000]
    n_blocks: 2
    n_heads: 4
    dropout: 0.0
    modality_dropout: 0.0
    max_seq_len: 128
    temporal_attn_layers: 1

  # Source variational encoder
  source_ve:
    depth: 2
    num_heads: 4
    num_queries: 8
    dropout: 0.0
    use_variational: true
    init_logvar: 0.5
    fixed_std: null

  # CSFM transport + sampler settings
  transport:
    path_type: Linear
    prediction: velocity
    loss_weight: null
    time_dist_type: uniform
    time_dist_shift: 1.0

# 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