from __future__ import annotations

# Manual-first bundling specs.
# mode:
#   - "file"        -> return the file verbatim
#   - "class"       -> extract only the canonical class from 'path', helpers inlined
# Extras:
#   - "helpers": list[(repo, path)] extra files to inline (3rd-party imports kept, internals stripped)
#   - "keep_functions": function names to extract from the SAME file
#   - "extra_classes":  class names to extract from the SAME file
#   - "strict_class_extract": if True, never fallback to whole file when class not found
INTERNAL_PREFIXES = (
    "adapters", "src.adapters",
)
MANUAL_BUNDLES: dict[str, dict] = {
    # Self-contained full files
    "ConformerBlock": {
        "repo": "lucidrains/conformer",
        "path": "conformer/conformer.py",
        "mode": "file",
        "helpers": [],
    },
    "RepConv": {
        "repo": "ultralytics/ultralytics",
        "path": "ultralytics/nn/modules/conv.py",
        "mode": "file",
        "helpers": [],
    },
    "CBAMBlock": {
        "repo": "xmu-xiaoma666/External-Attention-pytorch",
        "path": "model/attention/CBAM.py",
        "mode": "file",
        "helpers": [],
    },
    "ConvMixerBlock": {
        "repo": "SeuTao/FaceBagNet",
        "path": "model/ConvMixer.py",
        "mode": "file"
    },
    "AttentionAugmentedConv": {
        "repo": "pquochuy/attention_augmented_conv",
        "path": "attn_augconv.py",
        "mode": "file",
        "helpers": [],
    },
    "ConvNeXtBlock": {
        "repo": "hmorimitsu/ptlflow",
        "path": "ptlflow/models/memfof/layer.py",
        "mode": "file",
        "helpers": [],
    },
    "CrossAttention": {
        "repo": "rohankar7/Hope",
        "path": "ddpm.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "CSPBlock": {
        "repo": "Shohruh72/YOLOv11",
        "path": "nets/nn.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["Conv", "Residual"],
        "strict_class_extract": True,
    },
    "CSPELAN": {
        "repo": "johnvonueman/my_rtdetr",
        "path": "ultralytics/nn/Addmodules/RepNCSPELAN4.py",
        "mode": "file",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
    },
    "DeformableConv2d": {
        "repo": "RoboticImaging/RectConv",
        "path": "conv/DeformableConv2d.py",
        "mode": "file",
    },
    "DropBlock": {
        "repo": "SebastianPopescu/U-NET-for-LocalBrainAge-prediction",
        "path": "dropblock.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["_bernoulli"],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "DropPath": {
        "repo": "feng-lab/pytorch-image-models",
        "path": "timm/models/layers/drop.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["drop_path"],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "ESEBlock": {
        "repo": "youngwanLEE/centermask2",
        "path": "centermask/modeling/backbone/vovnet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["Hsigmoid"],
        "strict_class_extract": True,
    },
    "FeedForward": {
        "repo": "lucidrains/DALLE-pytorch",
        "path": "dalle_pytorch/transformer.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["GEGLU"],
        "strict_class_extract": True,
    },
    "Focus": {
        "repo": "ultralytics/yolov5",
        "path": "models/common.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["autopad"],
        "extra_classes": ["Conv"],
        "strict_class_extract": True,
    },
    "FusedMBConv": {
        "repo": "JunSeongLee1102/2023-Kaggle-Stanford-Ribonanza-RNA-Folding",
        "path": "main/modules.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "GhostBottleneck": {
        "repo": "huawei-noah/Efficient-AI-Backbones",
        "path": "ghostnet_pytorch/ghostnet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["hard_sigmoid", "_make_divisible"],
        "extra_classes": ["SqueezeExcite", "GhostModule"],
        "strict_class_extract": True,
    },
    "GhostModule": {
        "repo": "huawei-noah/Efficient-AI-Backbones",
        "path": "ghostnet_pytorch/ghostnet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "GRUCellBlock": {
        "repo": "pytorch/pytorch",
        "path": "torch/nn/modules/rnn.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "HardSwishConv": {
        "repo": "ultralytics/yolov5",
        "path": "utils/activations.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "InvertedResidual": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/shufflenetv2.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["channel_shuffle"],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "LSTMCellBlock" : {
        "repo": "pytorch/pytorch",
        "path": "torch/nn/modules/rnn.py",
        "mode": "file",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
    },
    "MaxVitBlock": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/maxvit.py",
        "mode": "file",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
    },
    "MBConv": {
        "repo": "d-li14/efficientnetv2.pytorch",
        "path": "effnetv2.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["_make_divisible"],
        "extra_classes": ["SiLU", "SELayer"],
        "strict_class_extract": True,
    },
    "MetaFormerBlock": {
        "repo": "huggingface/pytorch-image-models",
        "path": "timm/models/metaformer.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["Pooling", "LayerNorm2d", "DropPath", "Scale", "Mlp", "StarReLU"],
        "strict_class_extract": True,
    },
    "MixerBlock": {
        "repo": "google-research/scenic",
        "path": "scenic/projects/baselines/mixer.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "MobileViTBlock": {
        "repo": "JiwiChong/Crack-Segmentation",
        "path": "mobileVitblock.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["conv_nxn_bn", "conv_1x1_bn"],
        "extra_classes": ["Transformer", "PreNorm", "Attention", "FeedForward"],
        "strict_class_extract": True,
    },
    "MoEExpert": {
        "repo": "GeneZC/MiniMoE",
        "path": "modules/modeling_moebert_gate.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "MultiHeadAttention": {
        "repo": "alief46/amazon",
        "path": "LLM.txt",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "NonLocalBlock": {
        "repo": "dome272/MaskGIT-pytorch",
        "path": "helper.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["GroupNorm"],
        "strict_class_extract": True,
    },
    "OctaveConv": {
        "repo": "d-li14/octconv.pytorch",
        "path": "octconv.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "PatchEmbed": {
        "repo": "microsoft/Swin-Transformer",
        "path": "models/swin_transformer_v2.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "QFormerBlock": {
        "repo": "svical/Q-Formers-training",
        "path": "qformer.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["QFormerLayer"],
        "strict_class_extract": True,
    },
    "RefineBlock": {
        "repo": "Xt-Chen/SARPN",
        "path": "models/modules.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "RelPosAttention": {
        "repo": "ziplab/LITv2",
        "path": "classification/models/attentions.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "RepConv": {
        "repo": "ultralytics/ultralytics",
        "path": "ultralytics/nn/modules/conv.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["autopad"],
        "extra_classes": ["Conv"],
        "strict_class_extract": True,
    },
    "LayerScaleBlock": {
        "repo": "Laboratoire-de-Chemoinformatique/VQGAE",
        "path": "VQGAE/layers.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["AttentionTalkingHead", "FFN"],
        "strict_class_extract": True,
    },
    "RepLKBlock": {
        "repo": "DingXiaoH/RepLKNet-pytorch",
        "path": "replknet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["conv_bn_relu", "conv_bn"],
        "extra_classes": ["ReparamLargeKernelConv"],
        "strict_class_extract": True,
    },
    "ReZeroBlock": {
        "repo": "majumderb/rezero",
        "path": "rezero/transformer/rztx.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "RMSNorm": {
        "repo": "lucidrains/x-transformers",
        "path": "x_transformers/x_transformers.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "RotaryEmbedding": {
        "repo": "BlinkDL/SmallInitEmb",
        "path": "model.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "SEBottleneck": {
        "repo": "NagatoYuki0943/classification",    
        "path": "Heavy/pytorch/resnest_timm_s.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["SplitAttn","SplitAttn", "RadixSoftmax"],
        "strict_class_extract": True,
    },
    "InceptionBlock": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/inception.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["BasicConv2d"],
        "strict_class_extract": True,
    },
    # Adapter – take Adapter base class + external helper files; drop internal imports.
    "Adapter": {
        "repo": "adapter-hub/adapters",
        "path": "src/adapters/methods/adapter_layer_base.py",
        "mode": "class",
        "helpers": [
            ("adapter-hub/adapters", "src/adapters/composition.py"),
            ("adapter-hub/adapters", "src/adapters/context.py"),
        ],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "DenseBlock": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/densenet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["_DenseLayer"],
        "strict_class_extract": True,
    },
    "DownBlock": {
        "repo": "milesial/Pytorch-UNet",
        "path": "unet/unet_parts.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],   
        "extra_classes": ["DoubleConv"],
        "strict_class_extract": True,
    },

    # AntiAliasedResBlock -> keep class 'BlurPool' + same-file function 'get_pad_layer'
    "AntiAliasedResBlock": {
        "repo": "adobe/antialiased-cnns",
        "path": "antialiased_cnns/blurpool.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["get_pad_layer"],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "Bottleneck": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/resnet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["conv1x1", "conv3x3"],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "ShuffleUnit": {
        "repo": "MendelXu/ANN",
        "path": "models/cls/nets/shufflenetv2.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["ConvBnRelu", "ConvBn", "ShuffleUnit"],
        "strict_class_extract": True,
    },
    "SparseConvBlock": {
        "repo": "yxgeee/DepthComplete",
        "path": "models/SparseConvNet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["SparseConv"],
        "strict_class_extract": True,
    },
    "SplitAttnConv": {
        "repo": "MenghaoGuo/Awesome-Vision-Attentions",
        "path": "code/branch_attentions/resnest_module.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["make_divisible"],
        "extra_classes": [ "RadixSoftmax"],
        "strict_class_extract": True,
    },
    "SPPBlock": {
        "repo": "JiYuanFeng/MCTrans",
        "path": "mctrans/models/centers/cenet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "strict_class_extract": True,
    },
    "DacBlock": {
        "repo": "JiYuanFeng/MCTrans",
        "path": "mctrans/models/centers/cenet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "strict_class_extract": True,
    },
    "SPPFBlock": {
        "repo": "fly2023me/DB-FusionNet",
        "path": "main.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "SEBlock": {
        "repo": "fly2023me/DB-FusionNet",
        "path": "main.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "SqueezeExcitation": {
        "repo": "WZMIAOMIAO/deep-learning-for-image-processing",
        "path": "pytorch_classification/Test10_regnet/model.py",    
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "SRUppBlock": {
        "repo": "antspy/quantized_distillation",
        "path": "onmt/modules/SRU.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["check_sru_requirement"],
        "extra_classes": ["SRU_Compute"],
        "strict_class_extract": True,
    },
    "SwiGLU": {
        "repo": "lucidrains/PaLM-pytorch",
        "path": "palm_pytorch/palm_lite.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "strict_class_extract": True,
    },
    "SwinTransformerBlock": {
        "repo": "rishigami/Swin-Transformer-TF",
        "path": "swintransformer/model.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": ["window_partition", "window_reverse"],
        "extra_classes": ["WindowAttention", "Mlp", "DropPath"],
        "strict_class_extract": True,
    },
    "TemporalConvBlock": {
        "repo": "rhymes-ai/Allegro",
        "path": "allegro/models/vae/modules.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "TransformerEncoderLayer": {
        "repo": "JosefAlbers/Phi-3-Vision-MLX",
        "path": "gte.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    "UpBlock": {
        "repo": "Luoxd1996/nnunet_mini",
        "path": "unet.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["ConvBlock"],
        "strict_class_extract": True,
    },
    "WindowAttention": {
        "repo": "microsoft/Swin-Transformer",
        "path": "models/swin_transformer_v2.py",
        "mode": "class",
        "helpers": [],
        "strict_class_extract": True,
    },
    # ASPP -> keep only ASPPConv + ASPPPooling + ASPP classes (no extra helpers)
    "ASPP": {
        "repo": "pytorch/vision",
        "path": "torchvision/models/segmentation/deeplabv3.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["ASPPConv", "ASPPPooling"],
        "strict_class_extract": True,
    },
    "ConvBnSilu": {
        "repo": "152334H/DL-Art-School",
        "path": "codes/models/arch_util.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": ["SiLU"],
        "strict_class_extract": True,
    },

    # CausalConv1D (canonical class name is CausalConv1d)
    # Keep only the class with 3rd-party imports.
    "CausalConv1D": {
        "repo": "jakepoz/rnnt",
        "path": "rnnt/causalconv.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
    # Also allow canonical key to hit manual bundle:
    "CausalConv1d": {
        "repo": "jakepoz/rnnt",
        "path": "rnnt/causalconv.py",
        "mode": "class",
        "helpers": [],
        "keep_functions": [],
        "extra_classes": [],
        "strict_class_extract": True,
    },
}