import torch
import torch.nn as nn

class SoftPrompt(nn.Module):
    def __init__(self, length, embed_dim):
        super().__init__()
        self.embedding = nn.Parameter(torch.randn(length, embed_dim) * 0.01)

    def get_prompt(self, batch_size, device):
        return self.embedding.unsqueeze(0).expand(batch_size, -1, -1).to(device)

class PromptTuningWrapper(nn.Module):
    def __init__(self, model, soft_prompt):
        super().__init__()
        self.model = model
        self.soft_prompt = soft_prompt
        self.prompt_length = soft_prompt.embedding.shape[0]
        self.model.get_input_embeddings().weight.requires_grad = False

    def forward(self, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, labels=None, past_key_values=None, use_cache=False, **kwargs):
        if past_key_values is not None:
            if position_ids is not None:
                position_ids = position_ids + self.prompt_length
            return self.model(input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, labels=labels, past_key_values=past_key_values, use_cache=use_cache, **kwargs)

        if inputs_embeds is None:
            assert input_ids is not None
            inputs_embeds = self.model.get_input_embeddings()(input_ids)

        batch_size = inputs_embeds.size(0)
        device = inputs_embeds.device
        prompt_embeds = self.soft_prompt.get_prompt(batch_size, device)
        combined = torch.cat([prompt_embeds, inputs_embeds], dim=1)

        if attention_mask is not None:
            prompt_mask = torch.ones(batch_size, prompt_embeds.size(1), device=device, dtype=attention_mask.dtype)
            attention_mask = torch.cat([prompt_mask, attention_mask], dim=1)

        if position_ids is not None:
            prompt_pos = torch.arange(prompt_embeds.size(1), device=device, dtype=position_ids.dtype).unsqueeze(0).expand(batch_size, -1)
            position_ids = torch.cat([prompt_pos, position_ids + prompt_embeds.size(1)], dim=1)

        if labels is not None:
            prompt_labels = torch.full((batch_size, prompt_embeds.size(1)), -100, device=device, dtype=labels.dtype)
            labels = torch.cat([prompt_labels, labels], dim=1)

        return self.model(inputs_embeds=combined, attention_mask=attention_mask, position_ids=position_ids, labels=labels, use_cache=use_cache, **kwargs)

class PromptTuning:
    def __init__(self, model, length=20):
        self.model = model
        self.length = length
        embed_dim = model.get_input_embeddings().embedding_dim
        self.prompt = SoftPrompt(length, embed_dim)
        self.wrapper = None

    def wrap(self):
        self.model.get_input_embeddings().weight.requires_grad = False
        self.wrapper = PromptTuningWrapper(self.model, self.prompt)

    def parameters(self):
        yield from self.prompt.parameters()

def inject_prompt_tuning(model, length=20):
    peft = PromptTuning(model, length=length)
    peft.wrap()
    return peft.wrapper, peft