"""Matched-budget PEFT baselines for the demo (GPT-2), built on tensorpeft's slot discovery."""
import math

import torch
import torch.nn as nn
from tensorpeft.layers import attach_hooks, find_slots


class _Adapter(nn.Module):
    lr = 1e-3

    def trainable(self):
        return [p for p in self.parameters() if p.requires_grad]

    def n_params(self):
        return sum(p.numel() for p in self.trainable())

    def logits(self, ids):
        return self.model[0](ids).logits


class LoRA(_Adapter):
    lr = 1e-3

    def __init__(self, model, rank=1, targets=("c_attn.q", "c_attn.v"), layers=None, seed=0):
        super().__init__()
        self.model = [model]
        self.slots = [s for s in find_slots(model, list(targets)) if layers is None or s.layer in layers]
        g = torch.Generator().manual_seed(seed)
        dev = next(model.parameters()).device
        self.A = nn.ParameterList([nn.Parameter(((torch.rand(s.shape[0], rank, generator=g) * 2 - 1)
                                                 / math.sqrt(s.shape[0])).to(dev)) for s in self.slots])
        self.B = nn.ParameterList([nn.Parameter(torch.zeros(rank, s.shape[1], device=dev)) for s in self.slots])
        self.scale = 2.0
        self._i = {id(s): i for i, s in enumerate(self.slots)}
        self._h = attach_hooks(self.slots, lambda s, x: (x @ self.A[self._i[id(s)]]) @ self.B[self._i[id(s)]] * self.scale)


class IA3(_Adapter):
    lr = 3e-3

    def __init__(self, model, kv=("c_attn.k", "c_attn.v"), ff=True):
        super().__init__()
        self.model = [model]
        self.out_slots = find_slots(model, list(kv))
        self.ff_slots = find_slots(model, ["mlp.c_proj"]) if ff else []
        dev = next(model.parameters()).device
        self.l_out = nn.ParameterList([nn.Parameter(torch.ones(s.shape[1], device=dev)) for s in self.out_slots])
        self.l_ff = nn.ParameterList([nn.Parameter(torch.ones(s.shape[0], device=dev)) for s in self.ff_slots])
        io = {id(s): i for i, s in enumerate(self.out_slots)}
        iff = {id(s): i for i, s in enumerate(self.ff_slots)}

        def d_out(s, x):
            y = x @ s.weight() + s.module.bias[s.c0:s.c1]
            return y * (self.l_out[io[id(s)]] - 1)

        def d_ff(s, x):
            return (x * (self.l_ff[iff[id(s)]] - 1)) @ s.weight()

        self._h = attach_hooks(self.out_slots, d_out) + attach_hooks(self.ff_slots, d_ff)


class Prompt(_Adapter):
    lr = 3e-2

    def __init__(self, model, n, seed=0):
        super().__init__()
        self.model = [model]
        emb = model.get_input_embeddings().weight
        g = torch.Generator().manual_seed(seed)
        ids = torch.randint(0, emb.shape[0], (n,), generator=g).to(emb.device)
        self.n = n
        self.P = nn.Parameter(emb[ids].detach().clone())

    def logits(self, ids):
        m = self.model[0]
        e = torch.cat([self.P.expand(ids.shape[0], -1, -1), m.get_input_embeddings()(ids)], dim=1)
        return m(inputs_embeds=e).logits[:, self.n:]


class LNTune(_Adapter):
    lr = 1e-3

    def __init__(self, model, layers=None, final=True):
        super().__init__()
        self.model = [model]
        self.params = []
        for name, mod in model.named_modules():
            if isinstance(mod, nn.LayerNorm):
                parts = name.split(".")
                layer = next((int(p) for p in parts if p.isdigit()), None)
                if (layer is None and final) or (layer is not None and (layers is None or layer in layers)):
                    for p in mod.parameters():
                        p.requires_grad_(True)
                        self.params.append(p)

    def trainable(self):
        return self.params