from collections import defaultdict
from transformers.modeling_utils import Conv1D
import torch.nn as nn

def discover_mlp_blocks(model, min_hidden_dim=512):
    """
    Auto-discovers MLP (fc, proj) pairs in any Hugging Face transformer.

    Returns a list of dicts:
      {
        "block_name":  str,
        "fc_name":     str,
        "proj_name":   str,
        "fc":          module,
        "proj":        module,
      }
    """
    # 1) Build a map: block_prefix → all Linear/Conv1D layers under it
    all_modules = dict(model.named_modules())
    block_map = defaultdict(list)
    for name, mod in all_modules.items():
        if isinstance(mod, (nn.Linear, Conv1D)):
            parts = name.split(".")
            # walk up the name to find the block prefix
            for i in range(len(parts)-1, 0, -1):
                prefix = ".".join(parts[:i])
                if prefix in all_modules:
                    block_map[prefix].append((name, mod))
                    break

    # 2) Now collect the actual MLP blocks
    mlp_blocks = []    # ← make sure this exists before the loop!
    for block_name, layers in block_map.items():
        # skip anything that isn't clearly an MLP/FFN
        if not any(tok in block_name.lower() for tok in ("mlp", "ffn")):
            continue

        # pick only non‑square Linear/Conv1D layers
        candidates = []
        for name, mod in layers:
            w0, w1 = mod.weight.shape
            # skip square layers (e.g., attention proj)
            if w0 == w1:
                continue
            # now filter by size
            if isinstance(mod, nn.Linear):
                if max(w0, w1) >= min_hidden_dim:
                    candidates.append((name, mod, (w0, w1)))
            else:  # Conv1D
                # Conv1D: weight.shape = (in, out)
                if w1 >= min_hidden_dim:
                    candidates.append((name, mod, (w0, w1)))

        # we need at least 2 to form an MLP block
        if len(candidates) >= 2:
            # pick the two biggest by param count
            top2 = sorted(candidates, key=lambda x: x[2][0]*x[2][1], reverse=True)[:2]
            fc_name,  fc_mod,  _ = top2[0]
            proj_name, proj_mod, _ = top2[1]
            mlp_blocks.append({
                "block_name": block_name,
                "fc_name":    fc_name,
                "proj_name":  proj_name,
                "fc":         fc_mod,
                "proj":       proj_mod,
            })

    return mlp_blocks
