import sys
sys.path.insert(0, '.')
from train_spatial_beats import (
    make_ov1_local_spatial_v3b_classwarmup_config,
    configure_stage1_trainable_parameters,
)
from spatial_beats import SpatialBEATs

cfg = make_ov1_local_spatial_v3b_classwarmup_config()
print("=== Config ===")
print(f"  freeze_trunk_in_stage1: {cfg.freeze_trunk_in_stage1}")
print(f"  unfreeze_top_n_layers: {cfg.unfreeze_top_n_layers}")
print(f"  unfreeze_full_trunk: {cfg.unfreeze_full_trunk}")
print(f"  freeze_local_spatial_in_classwarmup: {cfg.freeze_local_spatial_in_classwarmup}")
print(f"  ddp_find_unused_parameters: {cfg.ddp_find_unused_parameters}")
print(f"  loss.lambda_direction: {cfg.loss.lambda_direction}")
print(f"  loss.lambda_dist: {cfg.loss.lambda_dist}")
print(f"  loss.lambda_cls_aux: {cfg.loss.lambda_cls_aux}")
print(f"  readout_scheme: {cfg.model.readout_scheme}")
print(f"  class_finetuned_ckpt: {cfg.class_finetuned_ckpt}")
print(f"  supervision_mode: {cfg.loss.supervision_mode}")

model = SpatialBEATs(cfg.model)
configure_stage1_trainable_parameters(model, cfg)

# Count
trainable = []
frozen = []
for name, param in model.named_parameters():
    if param.requires_grad:
        trainable.append(name)
    else:
        frozen.append(name)

print(f"\n=== Trainable ({len(trainable)}) ===")
for n in trainable:
    print(f"  ✅ {n}")
print(f"\n=== Frozen ({len(frozen)}) ===")
for n in frozen:
    print(f"  ❄️  {n}")

# Summary
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"\nTrainable: {trainable_params:,} / {total_params:,} = {trainable_params/total_params:.1%}")