"""Explore available timm backbones and inspect model properties.

This example demonstrates:
- Listing available backbones by architecture family
- Filtering by pretrained availability
- Inspecting model properties (parameters, features)
- Comparing different model configurations

Usage:
    python examples/backbone_discovery.py
"""

import autotimm
from autotimm import BackboneConfig


def main():
    print("=" * 70)
    print("AutoTimm Backbone Discovery")
    print("=" * 70)

    # ========================================================================
    # List all available backbones
    # ========================================================================
    all_models = autotimm.list_backbones()
    print(f"\nTotal available backbones: {len(all_models)}")

    # ========================================================================
    # Filter by architecture family
    # ========================================================================
    print("\n" + "-" * 70)
    print("Architecture Families")
    print("-" * 70)

    families = {
        "ResNet": "resnet*",
        "EfficientNet": "efficientnet*",
        "Vision Transformer": "vit_*",
        "ConvNeXt": "convnext*",
        "MobileNet": "*mobilenet*",
        "RegNet": "regnet*",
        "DenseNet": "densenet*",
        "Swin Transformer": "swin_*",
        "MaxViT": "maxvit*",
        "CoAtNet": "coatnet*",
    }

    for name, pattern in families.items():
        models = autotimm.list_backbones(pattern)
        pretrained = autotimm.list_backbones(pattern, pretrained_only=True)
        print(f"{name:.<30} {len(models):>5} total, {len(pretrained):>5} pretrained")

    # ========================================================================
    # Show sample models for each family
    # ========================================================================
    print("\n" + "-" * 70)
    print("Sample Models by Family")
    print("-" * 70)

    for name, pattern in list(families.items())[:4]:  # First 4 families
        models = autotimm.list_backbones(pattern, pretrained_only=True)
        print(f"\n{name}:")
        for m in models[:5]:
            print(f"  - {m}")
        if len(models) > 5:
            print(f"  ... and {len(models) - 5} more")

    # ========================================================================
    # Inspect specific backbones
    # ========================================================================
    print("\n" + "-" * 70)
    print("Model Inspection")
    print("-" * 70)

    models_to_inspect = [
        "resnet18",
        "resnet50",
        "efficientnet_b0",
        "efficientnet_b3",
        "vit_base_patch16_224",
        "convnext_tiny",
        "swin_tiny_patch4_window7_224",
    ]

    print(f"\n{'Model':<35} {'Params (M)':>12} {'Features':>10}")
    print("-" * 60)

    for model_name in models_to_inspect:
        try:
            backbone = autotimm.create_backbone(model_name)
            params = autotimm.count_parameters(backbone, trainable_only=False)
            print(f"{model_name:<35} {params / 1e6:>12.2f} {backbone.num_features:>10}")
        except Exception as e:
            print(f"{model_name:<35} Error: {e}")

    # ========================================================================
    # Advanced backbone configuration
    # ========================================================================
    print("\n" + "-" * 70)
    print("Advanced Backbone Configuration")
    print("-" * 70)

    # Standard backbone
    print("\n1. Standard backbone (pretrained):")
    backbone = autotimm.create_backbone("resnet50")
    print(f"   Output features: {backbone.num_features}")
    print(f"   Total params: {autotimm.count_parameters(backbone):,}")

    # With BackboneConfig for more control
    print("\n2. BackboneConfig with drop path:")
    config = BackboneConfig(
        model_name="vit_base_patch16_224",
        pretrained=True,
        drop_path_rate=0.1,  # Stochastic depth
    )
    backbone = autotimm.create_backbone(config)
    print(f"   Model: {config.model_name}")
    print(f"   Drop path rate: {config.drop_path_rate}")
    print(f"   Output features: {backbone.num_features}")
    print(f"   Total params: {autotimm.count_parameters(backbone):,}")

    # With dropout
    print("\n3. BackboneConfig with dropout:")
    config = BackboneConfig(
        model_name="efficientnet_b0",
        pretrained=True,
        drop_rate=0.2,  # Classifier dropout
    )
    backbone = autotimm.create_backbone(config)
    print(f"   Model: {config.model_name}")
    print(f"   Drop rate: {config.drop_rate}")
    print(f"   Output features: {backbone.num_features}")

    # Without pretrained weights
    print("\n4. Training from scratch (no pretrained):")
    config = BackboneConfig(
        model_name="resnet18",
        pretrained=False,
    )
    backbone = autotimm.create_backbone(config)
    print(f"   Model: {config.model_name}")
    print(f"   Pretrained: {config.pretrained}")
    print(f"   Output features: {backbone.num_features}")

    # ========================================================================
    # Search for specific model variants
    # ========================================================================
    print("\n" + "-" * 70)
    print("Searching for Model Variants")
    print("-" * 70)

    # Find all ImageNet-21k pretrained models
    print("\nImageNet-21k pretrained ViTs:")
    vit_21k = autotimm.list_backbones("vit_*in21k*", pretrained_only=True)
    for m in vit_21k[:5]:
        print(f"  - {m}")

    # Find CLIP models
    print("\nCLIP models:")
    clip_models = autotimm.list_backbones("*clip*", pretrained_only=True)
    for m in clip_models[:5]:
        print(f"  - {m}")

    # Find models by size
    print("\nSmall models (good for edge deployment):")
    small_models = [
        "mobilenetv3_small_100",
        "efficientnet_lite0",
        "resnet18",
        "convnext_tiny",
    ]
    for model_name in small_models:
        try:
            backbone = autotimm.create_backbone(model_name)
            params = autotimm.count_parameters(backbone)
            print(f"  - {model_name}: {params / 1e6:.1f}M params")
        except Exception:
            pass

    print("\n" + "=" * 70)
    print("Done! Use autotimm.list_backbones() to explore more models.")
    print("=" * 70)


if __name__ == "__main__":
    main()
