# src/slimformers/mlp_discovery.py
import torch
from torch import nn
from transformers.modeling_utils import Conv1D
from collections import OrderedDict, defaultdict

# GPT‑2 FFN finder
def discover_gpt2_ffns(model, sample_inputs=None):
    blocks = []
    for i, block in enumerate(model.transformer.h):
        blocks.append({
            "type":      "ffn",
            "fc_name":   f"transformer.h.{i}.mlp.c_fc",
            "proj_name": f"transformer.h.{i}.mlp.c_proj",
            "fc":        block.mlp.c_fc,
            "proj":      block.mlp.c_proj,
        })
    return blocks

# BERT FFN finder
def discover_bert_ffns(model, sample_inputs=None):
    """
    Works for both BertModel and BertForMaskedLM
    """
    core = getattr(model, "bert", model)
    blocks = []

    prefix = "bert." if hasattr(model, "bert") else ""

    for i, layer in enumerate(core.encoder.layer):
        fc_path   = f"{prefix}encoder.layer.{i}.intermediate.dense"
        proj_path = f"{prefix}encoder.layer.{i}.output.dense"
        blocks.append({
            "type":      "ffn",
            "fc_name":   fc_path,
            "proj_name": proj_path,
            "fc":        layer.intermediate.dense,
            "proj":      layer.output.dense,
        })
    return blocks


def discover_llama_ffns(model, sample_inputs=None):
    """
    LLaMA uses a gated GeLU FFN:
      hidden → gate_proj(x) ∘ up_proj(x) → down_proj
    We prune gate_proj & up_proj on their output dim, and down_proj on its input dim.
    """
    blocks = []
    # the causal‑LM wrapper holds the actual LLaMA module in .model
    core = getattr(model, "model", model)
    for i, layer in enumerate(core.layers):
        mlp = layer.mlp
        blocks.append({
            "type":      "gated",
            "gate_name": f"model.layers.{i}.mlp.gate_proj",
            "up_name":   f"model.layers.{i}.mlp.up_proj",
            "down_name": f"model.layers.{i}.mlp.down_proj",
            "gate":      mlp.gate_proj,
            "up":        mlp.up_proj,
            "down":      mlp.down_proj,
        })
    return blocks

def discover_dynamic(model, sample_inputs):
    from .pruner import Pruner  
    return Pruner._discover_mlp_blocks_dynamic(model, sample_inputs)

# registry mapping HF class‐names to finder fn
DISCOVERY_REGISTRY = {
    "GPT2Model":                     discover_gpt2_ffns,
    "GPT2LMHeadModel":               discover_gpt2_ffns,
    "BertModel":                     discover_bert_ffns,
    "BertForMaskedLM":               discover_bert_ffns,
    "LlamaModel":                    discover_llama_ffns,
    "LlamaForCausalLM":              discover_llama_ffns,
}

def default_discover(model, sample_inputs):
    return discover_ffns_model_agnostic(model)

def discover_ffns_model_agnostic(model, min_hidden_dim=128):
    all_mods = dict(model.named_modules())
    block_map = defaultdict(list)

    # group linears/conv1Ds by their parent
    for name, mod in all_mods.items():
        if isinstance(mod, (nn.Linear, Conv1D)):
            parts = name.split(".")
            for i in range(len(parts)-1, 0, -1):
                prefix = ".".join(parts[:i])
                if prefix in all_mods:
                    block_map[prefix].append(name)
                    break

    mlp_blocks = []
    for block_name, names in block_map.items():
        suffixes = [n.rsplit(".",1)[1] for n in names]

        if {"gate_proj", "up_proj", "down_proj"}.issubset(set(suffixes)):
            mlp_blocks.append({
                "type": "gated",
                "gate_name": f"{block_name}.gate_proj",
                "up_name":   f"{block_name}.up_proj",
                "down_name": f"{block_name}.down_proj",
                "gate":      all_mods[f"{block_name}.gate_proj"],
                "up":        all_mods[f"{block_name}.up_proj"],
                "down":      all_mods[f"{block_name}.down_proj"],
            })
            continue

        candidates = []
        for name in names:
            mod = all_mods[name]
            w0, w1 = mod.weight.shape
            if w0 != w1 and max(w0, w1) >= min_hidden_dim:
                candidates.append((name, mod.weight.numel()))

        if len(candidates) >= 2:
            candidates.sort(key=lambda x: x[1], reverse=True)
            fc_name, _   = candidates[0]
            proj_name, _ = candidates[1]
            mlp_blocks.append({
                "type":      "ffn",
                "fc_name":   fc_name,
                "proj_name": proj_name,
                "fc":        all_mods[fc_name],
                "proj":      all_mods[proj_name],
            })

    return mlp_blocks