# Lunar Crater Detection (LRO WAC, 5-band vis) — 100% training data
# Backbone: NASA-IBM Lunar Foundation Model at patch size 8,  LoRA PEFT variant
#
# Requires: pip install terratorch[peft]
#
# RUN:
#   PYTHONPATH=.:$PYTHONPATH \
#   terratorch fit --config terratorch_integration/configs/wac_craters/full_data/ni_lfm_ps8_lora.yaml

seed_everything: 46
custom_modules_path: terratorch_integration

trainer:
  accelerator: gpu
  devices: 1
  precision: "32"
  max_epochs: 100
  check_val_every_n_epoch: 1
  log_every_n_steps: 10
  deterministic: true
  benchmark: false
  enable_checkpointing: true
  logger:
    class_path: TensorBoardLogger
    init_args:
      save_dir: tb_logs/wac-crater
      name: ni_lfm_ps8_lora_100pct
  callbacks:
    - class_path: LearningRateMonitor
      init_args:
        logging_interval: epoch
    - class_path: terratorch_integration.DeterministicAugmentation
    - class_path: ModelCheckpoint
      init_args:
        dirpath: tb_logs/wac-crater/checkpoints/ni_lfm_ps8_lora_100pct
        monitor: val_map
        mode: max
        filename: best-{epoch:02d}-{val_map:.3f}
        save_last: true

data:
  class_path: terratorch_integration.data_adapter.LunarWACCraterDataModule
  init_args:
    data_dir: data/wac_craters_dataset
    annotations_dir: data/wac_craters_dataset
    train_fraction: 1.0
    batch_size: 4
    num_workers: 4
    image_size: 512
    min_diameter_px: 10.0
    drop_no_crater_images: true
    wac_norm_means: [0.018106, 0.025262, 0.027113, 0.028616, 0.030757]   # train-split stats (800 tiles)
    wac_norm_stds: [0.010051, 0.013646, 0.014491, 0.015267, 0.016555]
    train_transforms:
      class_path: terratorch_integration.data_utils.D4DetectionTransform
      init_args:
        image_keys: [image, vis]
        p: 1.0

model:
  class_path: terratorch_integration.LunarObjectDetectionTask
  init_args:
    model_factory: ObjectDetectionModelFactory
    model_args:
      framework: faster-rcnn
      backbone: ni_lfm_v1_base
      backbone_checkpoint_path: backbone/checkpoint.pt
      backbone_cfg: backbone/config.yaml

      num_classes: 2
      framework_min_size: 512
      framework_max_size: 512

      backbone_modalities:
        - vis

      backbone_patch_size: 8
      backbone_remove_register_tokens: false
      backbone_merge_method: concat

      backbone_drop_path_rate_encoder: 0.1

      in_channels: 5
      necks:
        - name: SelectIndices
          indices: [2, 5, 8, 11]
        - name: ReshapeTokensToImage
          remove_cls_token: false
          h: 64
        - name: MultilayerSimpleFeaturePyramid
          out_channels: 256
        - name: FeaturePyramidNetworkNeck

    freeze_backbone: false
    freeze_decoder: false

    class_names:
      - Background
      - Crater

    score_threshold: 0.05
    iou_threshold: 0.5
    eval_box_detections_per_img: 500
    metric_kwargs:
      max_detection_thresholds: [100, 300, 500]
      class_metrics: true

    anchor_sizes: [[8], [16], [32], [64]]
    anchor_aspect_ratios: [0.5, 1.0, 2.0]

    peft_config:
      method: LORA
      peft_config_kwargs:
        r: 16
        lora_alpha: 32
        lora_dropout: 0.05
        bias: none
        target_modules:
          - qkv
          - mlp.fc1
          - mlp.fc2
          - mlp.fc3
        exclude_modules: '.*decoder\..*'

    backbone_lr: 1.0e-4
    head_lr: 1.0e-4
    layer_decay: 1.0
    weight_decay: 0.0
    head_weight_decay: 1.0e-3
    betas: [0.9, 0.999]
    warmup_steps: 500
    eta_min: 1.0e-6