import torch
import torch.nn as nn

from trainDraft import DraftChain


class DraftDecoderLayer(nn.Module):

    def __init__(self, block):
        super().__init__()
        self.block = block

    def forward(self, hidden_states, position_embeddings=None, attention_mask=None,
                position_ids=None, past_key_values=None, use_cache=None, **kwargs):
        return self.block(hidden_states)


class _Identity(nn.Module):
    def forward(self, x):
        return x



def freeze_chain(chain):
    for b in chain.blocks:
        for name in ("fc1", "fc2"):
            wn = getattr(b, name)
            w = wn.weight().detach()
            lin = nn.Linear(w.shape[1], w.shape[0], bias=False, dtype=w.dtype)
            lin.weight.data.copy_(w)
            setattr(b, name, lin)
    return chain

def load_chain(path, dtype=None):
    ckpt = torch.load(path, weights_only=True)
    chain = DraftChain(ckpt["D"], ckpt["hidden"], ckpt["k"])
    chain.load_state_dict(ckpt["chain"])
    if dtype is not None:
        chain = chain.to(dtype)
    chain.eval()
    return freeze_chain(chain)


def apply_merge(model, chain):
    n = model.config.num_hidden_layers
    k = len(chain.blocks)
    for _ in range(k):
        del model.model.layers[-1] 
    for j in range(k):
        layer = DraftDecoderLayer(chain.blocks[j])
        layer.to(model.device) 
        model.model.layers.append(layer)
    model.model.norm = _Identity()
    return model