import torch
from torch import nn
from collections import defaultdict
from transformers.modeling_utils import Conv1D
from .mlp_discovery import DISCOVERY_REGISTRY, default_discover

"""
todo: 

implement dataset averaging for activations 
make work for t5 transformer models

"""

class Pruner:
    def __init__(self, model: nn.Module, pruning_strategy=None):
        self.model = model
        self.activations = {}
        self.pruning_strategy = pruning_strategy or self._compute_topk_neurons
        self.saved_target_modules = self._precompute_target_modules()

    def _precompute_target_modules(self):
        target_keywords = [
            "c_fc", "c_proj",                         # GPT
            "gate_proj", "up_proj", "down_proj",      # LLaMA
            "intermediate.dense", "output.dense",     # BERT
            "wi", "wo", "w0",                         # T5, Flan-T5
            "mlp.dense_h_to_4h", "mlp.dense_4h_to_h",  # OPT, MosaicML
            "dense", "fc1", "fc2",                    # fallback/legacy
        ]
        return [
            name for name, module in self.model.named_modules()
            if isinstance(module, (nn.Linear, Conv1D)) and any(kw in name for kw in target_keywords)
        ]


    def get_prunable_linear_modules(self):
        return self.saved_target_modules

    def _hook_activations(self, layer_name: str):
        layer = dict(self.model.named_modules())[layer_name]
        def forward_hook(module, inp, out):
            self.activations[layer_name] = out.detach()
        layer.register_forward_hook(forward_hook)

    def _accumulate_activations(self, layer_name, acts):
        if layer_name not in self.activations:
            self.activations[layer_name] = acts.detach().cpu()
        else:
            self.activations[layer_name] = torch.cat([self.activations[layer_name], acts.detach().cpu()], dim=0)
    
    def accumulate_over_dataset(self, dataloader, device="cuda"):
        self.model.eval()
        handles = []

        # Register hooks
        for name in self._active_hook_names:
            layer = dict(self.model.named_modules())[name]
            handle = layer.register_forward_hook(lambda m, i, o, n=name: self._accumulate_activations(n, o))
            handles.append(handle)

        with torch.no_grad():
            for batch in dataloader:
                inputs = {k: v.to(device) for k, v in batch.items()}
                _ = self.model(**inputs)

        for h in handles:
            h.remove()


    def _compute_topk_neurons(self, activations, sparsity: float):
        # handle [B,S,D] or [B,D] or [S,B,D]
        if activations.dim() == 3:
            mags = activations.abs().mean(dim=(0,1))
        elif activations.dim() == 2:
            mags = activations.abs().mean(dim=0)
        else:
            raise ValueError(f"Bad act shape {activations.shape}")
        total = mags.size(0)
        k = int((1.0 - sparsity) * total)
        return torch.topk(mags, k=k).indices, total

    def _prune_output_dim(self, layer, idx):
        """
        Slice away the unwanted output neurons on a Linear or Conv1D layer,
        and make sure Conv1D.forward() knows about the new output size.
        """
        W = layer.weight.data
        B = layer.bias.data if layer.bias is not None else None

        if isinstance(layer, Conv1D):
            new_W = W[:, idx]  
            layer.weight = nn.Parameter(new_W.to(layer.weight.device))
            if B is not None:
                new_B = B[idx]
                layer.bias = nn.Parameter(new_B.to(layer.bias.device))
            layer.nf = new_W.shape[1]

        else:
            layer.weight.data = W[idx, :]
            layer.out_features = layer.weight.size(0)
            if B is not None:
                layer.bias.data = B[idx]

    def _prune_input_dim(self, layer, idx):
        W = layer.weight.data
        if isinstance(layer, Conv1D):
            # W: (in, out)
            layer.weight.data = W[idx, :]
            layer.in_features = layer.weight.size(0)
        else:
            # Linear: W: (out, in)
            layer.weight.data = W[:, idx]
            layer.in_features = layer.weight.size(1)

    @staticmethod
    def _discover_mlp_blocks(model: nn.Module, sample_inputs: dict):
        """
        Choose a model‐specific finder if registered, otherwise fall
        back to the generic dynamic hook approach.
        """
        cls_name = type(model).__name__
        finder   = DISCOVERY_REGISTRY.get(cls_name, default_discover)
        # both discovery fns accept (model, sample_inputs)
        return finder(model, sample_inputs)

    def prune_all_mlp_layers(self, inputs, sparsity: float = 0.3):
        self.model.eval()
        device = next(self.model.parameters()).device
        print(f"Pruning all MLP layers with {sparsity*100:.0f}% sparsity")

        blocks = Pruner._discover_mlp_blocks(self.model, sample_inputs=inputs)
        print(f"Discovered {len(blocks)} MLP blocks.")

        for i, blk in enumerate(blocks):
            if blk["type"] == "gated":
                gate, up, down = blk["gate"], blk["up"], blk["down"]
                hook_name = blk["gate_name"]
                self._hook_activations(hook_name)

                with torch.no_grad():
                    _ = self.model(**{k: v.to(device) for k, v in inputs.items()})

                acts = self.activations.pop(hook_name)
                idx, orig = self.pruning_strategy(acts, sparsity)

                # prune both gate_proj and up_proj on *output* dim
                self._prune_output_dim(gate, idx)
                self._prune_output_dim(up,   idx)

                # prune down_proj on *input* dim
                self._prune_input_dim(down, idx)

                print(f"Block {i} (gated): pruned {orig} → {len(idx)}")

            else:
                fc, proj = blk["fc"], blk["proj"]
                hook_name = blk["fc_name"]
                self._hook_activations(hook_name)

                with torch.no_grad():
                    _ = self.model(**{k: v.to(device) for k, v in inputs.items()})

                acts = self.activations.pop(hook_name)
                idx, orig = self._compute_topk_neurons(acts, sparsity)

                # figure out if fc’s orig dim is its input or output
                w = fc.weight.data.shape
                if (isinstance(fc, Conv1D) and w[0] == orig) or (not isinstance(fc, Conv1D) and w[1] == orig):
                    self._prune_input_dim(fc, idx)
                else:
                    self._prune_output_dim(fc, idx)

                w2 = proj.weight.data.shape
                if (isinstance(proj, Conv1D) and w2[0] == orig) or (not isinstance(proj, Conv1D) and w2[1] == orig):
                    self._prune_input_dim(proj, idx)
                else:
                    self._prune_output_dim(proj, idx)

                print(f"Block {i}: pruned {orig} → {len(idx)}")

        print("All MLP layers pruned successfully!")

    def _accumulate_activations(self, layer_name, acts):
        if layer_name not in self.activations:
            self.activations[layer_name] = acts.detach().cpu()
        else:
            self.activations[layer_name] = torch.cat([self.activations[layer_name], acts.detach().cpu()], dim=0)

    def accumulate_over_dataset(self, dataloader, device="cuda"):
        self.model.eval()
        handles = []

        # Register hooks
        for name in self._active_hook_names:
            layer = dict(self.model.named_modules())[name]
            handle = layer.register_forward_hook(lambda m, i, o, n=name: self._accumulate_activations(n, o))
            handles.append(handle)

        with torch.no_grad():
            for batch in dataloader:
                inputs = {k: v.to(device) for k, v in batch.items()}
                _ = self.model(**inputs)

        for h in handles:
            h.remove()

