from dataclasses import dataclass
from typing import List, Union

import torch
import transformers
from torch import Tensor
from transformers import set_seed

@dataclass
class SoftOptConfig:
    num_steps: int = 100
    optim_str_init: str = "x x x x x x x x x x x x x x x x x x x x"
    rand_init: bool = False
    num_tokens: int = 20
    lr: float = 0.0001
    early_stop_loss: float = None
    add_space_before_target: bool = False
    device: str = "cuda"
    seed: int = None
    verbose: bool = True

@dataclass
class SoftOptResult:
    losses: List[float]
    optim_embeds: Tensor
    input_embeds: Tensor

def mellowmax(t: Tensor, alpha=1.0, dim=-1):
   return 1.0 / alpha * (torch.logsumexp(alpha * t, dim=dim) - torch.log(torch.tensor(t.shape[-1], dtype=t.dtype, device=t.device)))
    
from accelerate import Accelerator

def run(
    model: transformers.PreTrainedModel,
    tokenizer: transformers.PreTrainedTokenizer,
    messages: Union[str, List[dict]],
    target: str,
    config: SoftOptConfig = None, 
) -> dict:
    """Generates a single optimized string using soft-prompt optimization with `accelerate` for distributed and mixed-precision support.

    Args:
        model: The model to optimize on.
        tokenizer: The model's tokenizer.
        messages: The conversation to use for optimization.
        target: The target generation.
        config: The configuration to use.
    
    Returns:
        A SoftOptResult object that contains losses and the optimized strings.
    """
    # Initialize accelerator
    accelerator = Accelerator(mixed_precision="bf16")
    
    if config is None:
        config = SoftOptConfig()
    
    if config.seed is not None:
        set_seed(config.seed)
        torch.use_deterministic_algorithms(True, warn_only=True)

    if isinstance(messages, str):
        messages = [{"role": "user", "content": messages}]
    
    if not any(["{optim_str}" in d["content"] for d in messages]):
        messages[-1]["content"] = messages[-1]["content"] + "{optim_str}"
    
    model = model.to(config.device)

    template = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    if tokenizer.bos_token and template.startswith(tokenizer.bos_token):
        template = template.replace(tokenizer.bos_token, "")
    before_str, after_str = template.split("{optim_str}")

    target = " " + target if config.add_space_before_target else target

    before_ids = tokenizer([before_str], padding=False)["input_ids"]
    after_ids = tokenizer([after_str], add_special_tokens=False)["input_ids"]
    target_ids = tokenizer([target], add_special_tokens=False)["input_ids"]

    before_ids, after_ids, target_ids = [torch.tensor(ids, device=config.device) for ids in (before_ids, after_ids, target_ids)]

    embedding_layer = model.get_input_embeddings()
    before_embeds, after_embeds, target_embeds = [embedding_layer(ids) for ids in (before_ids, after_ids, target_ids)]

    with torch.no_grad():
        output = model(inputs_embeds=before_embeds, use_cache=True)
        prefix_cache = output.past_key_values

    if not config.rand_init:
        optim_ids = tokenizer(config.optim_str_init, return_tensors="pt", add_special_tokens=False)["input_ids"].to(config.device)
        optim_embeds = embedding_layer(optim_ids).detach().clone().requires_grad_()
    else:
        optim_embeds = torch.randn((1, config.num_tokens, model.config.hidden_size), device=config.device).requires_grad_()

    optimizer = torch.optim.SGD([optim_embeds], lr=config.lr)
    
    for param in model.parameters():
        param.requires_grad = False

    # Prepare model, optimizer, and embeddings with accelerator
    model, optimizer = accelerator.prepare(model, optimizer)
    optim_embeds = accelerator.prepare(optim_embeds)

    losses = []
    for i in range(config.num_steps):
        optimizer.zero_grad()
        input_embeds = torch.cat([optim_embeds, after_embeds.detach(), target_embeds.detach()], dim=1)

        output = model(inputs_embeds=input_embeds, past_key_values=prefix_cache, output_hidden_states=True)
        logits = output.logits

        shift = input_embeds.shape[1] - target_ids.shape[1]
        shift_logits = logits[..., shift-1:-1, :].contiguous() 
        shift_labels = target_ids

        loss = torch.nn.functional.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        loss_float = loss.item()
        losses.append(loss_float)

        if config.verbose:
            print(f"Iter: {i} Loss: {loss_float}")
        
        if config.early_stop_loss and loss_float < config.early_stop_loss:
            break

        accelerator.backward(loss)  # Use accelerator for backpropagation
        optimizer.step()
    
    result = SoftOptResult(
        losses=losses,
        optim_embeds=optim_embeds.cpu(),
        input_embeds=torch.cat([before_embeds, optim_embeds, after_embeds], dim=1).cpu()
    )

    return result
