import gradio as gr
import torch
import concurrent.futures
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "HuggingFaceTB/SmolLM2-360M-Instruct"
MAX_NEW_TOKENS = 64

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if getattr(torch.backends, "mps", None) is not None
    and torch.backends.mps.is_available()
    else "cpu"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
model.eval()

if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

torch.set_num_threads(2)


def encode_prompt(prompt):
    messages = [{"role": "user", "content": prompt}]
    if getattr(tokenizer, "chat_template", None):
        encoded = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        )
    else:
        encoded = tokenizer(prompt, return_tensors="pt")
    return {key: value.to(device) for key, value in encoded.items()}


def decode_generated(sequence_ids, prompt_length):
    generated_ids = sequence_ids[prompt_length:]
    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()


def min_p_sampling(logits, pbase=0.1):
    """
    Perform min-p sampling on the logits. As described in
    https://arxiv.org/abs/2407.01082

    Args:
        logits (torch.Tensor): 1D tensor of logits for the next token.
        pbase (float): Base probability to scale pmax.

    Returns:
        int: The sampled token index.
    """
    # Convert logits to probabilities.
    probs = torch.softmax(logits, dim=-1)

    # 1. Find maximum probability.
    pmax = probs.max()

    # 2. Compute the dynamic threshold.
    pscaled = pbase * pmax

    # 3. Create a mask of tokens with probability >= pscaled.
    mask = probs >= pscaled
    # In the unlikely event that no token meets the threshold, use the full distribution.
    if mask.sum() == 0:
        mask = torch.ones_like(probs, dtype=torch.bool)

    probs_filtered = probs * mask.float()

    # 4. Normalize and sample.
    probs_normalized = probs_filtered / probs_filtered.sum()
    sampled_index = torch.multinomial(probs_normalized, num_samples=1)

    return sampled_index.item()


def generate_laconic_completion(
    prompt: str, n: int = 5, max_new_tokens: int = MAX_NEW_TOKENS
):
    # Generate n sampled completions and return the shortest decoded result.
    with torch.no_grad():
        encoded = encode_prompt(prompt)
        input_ids = encoded["input_ids"]
        attention_mask = encoded.get("attention_mask")

        outputs = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            num_return_sequences=n,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
        )
        completions = [
            decode_generated(output, input_ids.size(1)) for output in outputs
        ]
        return min(completions, key=len)


def generate_with_confidence(input_ids, max_new_tokens):
    """
    Generate a sequence using greedy decoding while returning the scores.
    """
    outputs = model.generate(
        input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        output_scores=True,
        return_dict_in_generate=True,
        pad_token_id=tokenizer.pad_token_id,
    )
    return outputs


def compute_answer_confidence(outputs):
    """
    Compute the answer confidence over the generated tokens.
    For each generated token, compute the difference between the top-1 and top-2 logits.
    Returns the average difference.
    """
    diffs = []
    for score in outputs.scores:
        # Get top-2 logit values
        top2 = torch.topk(score[0], 2)
        diff = top2.values[0] - top2.values[1]
        diffs.append(diff.item())

    return sum(diffs) / len(diffs) if diffs else 0.0


def cot_decoding(prompt, k=5, max_new_tokens=MAX_NEW_TOKENS):
    """
    Perform Chain-of-Thought (CoT) decoding by exploring top-k alternative paths.
    """
    encoded = encode_prompt(prompt)
    input_ids = encoded["input_ids"]
    prompt_length = input_ids.size(1)

    # Get logits for the next token
    with torch.no_grad():
        outputs = model(input_ids)
    logits = outputs.logits[0, -1, :]

    # Get top-k candidate tokens
    topk = torch.topk(logits, k)
    candidate_tokens = topk.indices

    paths = []
    for token in candidate_tokens:
        # Append the candidate token to the prompt
        new_input_ids = torch.cat([input_ids, token.view(1, 1)], dim=1)

        # Generate a full sequence with output scores
        gen_outputs = generate_with_confidence(
            new_input_ids, max_new_tokens=max_new_tokens
        )

        # Decode the generated sequence
        generated_text = decode_generated(gen_outputs.sequences[0], prompt_length)

        # Compute answer confidence
        confidence = compute_answer_confidence(gen_outputs)

        paths.append({"text": generated_text, "confidence": confidence})

    return max(paths, key=lambda x: x["confidence"])["text"]


def generate_completion(prompt, strategy, params):
    """
    Generate a complete answer using model.generate with specified parameters.
    """
    with torch.no_grad():
        encoded = encode_prompt(prompt)
        input_ids = encoded["input_ids"]
        attention_mask = encoded.get("attention_mask")

        output_ids = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=MAX_NEW_TOKENS,
            pad_token_id=tokenizer.pad_token_id,
            **params,
        )
    return decode_generated(output_ids[0], input_ids.size(1))


def generate_min_p_completion(prompt, pbase=0.1, max_new_tokens=MAX_NEW_TOKENS):
    encoded = encode_prompt(prompt)
    input_ids = encoded["input_ids"]
    prompt_length = input_ids.size(1)
    past = None
    with torch.no_grad():
        for _ in range(max_new_tokens):
            # Only pass the last token if past is available
            outputs = (
                model(input_ids[:, -1:], past_key_values=past)
                if past is not None
                else model(input_ids)
            )
            past = outputs.past_key_values
            logits = outputs.logits[:, -1, :]

            next_token = min_p_sampling(logits.squeeze(0), pbase=pbase)
            next_token_ids = torch.tensor([[next_token]], device=device)
            input_ids = torch.cat([input_ids, next_token_ids], dim=-1)
            if next_token == tokenizer.eos_token_id:
                break
    return decode_generated(input_ids[0], prompt_length)


def sequence_log_probability(sequence_ids, prompt_length):
    if sequence_ids.size(1) <= prompt_length:
        return float("-inf")

    with torch.no_grad():
        outputs = model(sequence_ids)
        logits = outputs.logits[:, prompt_length - 1 : -1, :]
        target_ids = sequence_ids[:, prompt_length:]
        log_probs = torch.log_softmax(logits, dim=-1)
        token_log_probs = log_probs.gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
        return token_log_probs.sum().item()


def sample_base_sequence(input_ids, attention_mask, max_new_tokens):
    with torch.no_grad():
        return model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            top_p=0.95,
            temperature=0.8,
            pad_token_id=tokenizer.pad_token_id,
        )


def generate_power_sampling_completion(
    prompt,
    beta=2.0,
    mcmc_steps=4,
    max_new_tokens=MAX_NEW_TOKENS,
):
    """
    Approximate the Karan-Du power distribution q(y) proportional to p(y)^beta.

    We use an independence Metropolis sampler: proposals are complete continuations
    from the base model p(y), while the acceptance rule corrects toward p(y)^beta.
    """
    encoded = encode_prompt(prompt)
    input_ids = encoded["input_ids"]
    attention_mask = encoded.get("attention_mask")
    prompt_length = input_ids.size(1)

    current = sample_base_sequence(input_ids, attention_mask, max_new_tokens)
    current_logp = sequence_log_probability(current, prompt_length)
    best = current
    best_logp = current_logp

    for _ in range(mcmc_steps):
        candidate = sample_base_sequence(input_ids, attention_mask, max_new_tokens)
        candidate_logp = sequence_log_probability(candidate, prompt_length)
        log_acceptance = (beta - 1.0) * (candidate_logp - current_logp)

        if torch.log(torch.rand((), device=device)).item() < min(0.0, log_acceptance):
            current = candidate
            current_logp = candidate_logp

        if current_logp > best_logp:
            best = current
            best_logp = current_logp

    return decode_generated(best[0], prompt_length)


def generate_all(prompt):
    """
    Run multiple decoding strategies concurrently and yield updates as each completes.
    """
    # Define each decoding strategy and its parameters.
    methods = {
        "Greedy": {"type": "default", "params": {"do_sample": False}},
        "Top-k Sampling": {
            "type": "default",
            "params": {"do_sample": True, "top_k": 100},
        },
        "Top-p Sampling": {
            "type": "default",
            "params": {"do_sample": True, "top_p": 0.95},
        },
        "Beam Search": {
            "type": "default",
            "params": {"num_beams": 5, "early_stopping": True},
        },
        "Eta Sampling": {
            "type": "default",
            "params": {"do_sample": True, "eta_cutoff": 0.3},
        },
        "Epsilon Sampling": {
            "type": "default",
            "params": {"do_sample": True, "epsilon_cutoff": 0.2},
        },
        "Min-p Sampling": {"type": "min_p", "pbase": 0.1},
        "laconic": {
            "type": "laconic",
            "params": {"n": 5},
        },
        "Power Sampling": {"type": "power_sampling", "beta": 2.0, "mcmc_steps": 4},
        "COT Decoding": {
            "type": "cot_decoding",
            "params": {"k": 5, "max_new_tokens": MAX_NEW_TOKENS},
        },
    }

    # Define the order for display.
    method_order = [
        "Greedy",
        "Top-k Sampling",
        "Top-p Sampling",
        "Beam Search",
        "Min-p Sampling",
        "Eta Sampling",
        "Epsilon Sampling",
        "laconic",
        "Power Sampling",
        "COT Decoding",
    ]
    results = {method: None for method in methods}

    # Yield an initial placeholder state.
    yield tuple("Processing..." for _ in method_order)

    # Use a thread pool to run each generation concurrently.
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_method = {}
        for method, info in methods.items():
            if info["type"] == "default":
                future = executor.submit(
                    generate_completion, prompt, method, info["params"]
                )
            elif info["type"] == "min_p":
                future = executor.submit(
                    generate_min_p_completion, prompt, info["pbase"]
                )
            elif info["type"] == "laconic":
                future = executor.submit(
                    generate_laconic_completion, prompt, **info["params"]
                )
            elif info["type"] == "power_sampling":
                future = executor.submit(
                    generate_power_sampling_completion,
                    prompt,
                    beta=info["beta"],
                    mcmc_steps=info["mcmc_steps"],
                )
            elif info["type"] == "cot_decoding":
                future = executor.submit(cot_decoding, prompt, **info["params"])

            future_to_method[future] = method

        # As each future completes, update its result and yield the current state.
        for future in concurrent.futures.as_completed(future_to_method):
            method = future_to_method[future]
            try:
                result = future.result()
            except Exception as exc:
                result = f"Error: {exc}"
            results[method] = result

            # Yield the results in the pre-defined order; pending methods show "Processing..."
            yield tuple(
                results[m] if results[m] is not None else "Processing..."
                for m in method_order
            )


# Create the Gradio interface.
interface = gr.Interface(
    fn=generate_all,
    inputs=gr.Textbox(lines=3, placeholder="Enter your prompt here...", label="Prompt"),
    outputs=[
        gr.Textbox(label="Greedy"),
        gr.Textbox(label="Top-k Sampling"),
        gr.Textbox(label="Top-p Sampling"),
        gr.Textbox(label="Beam Search"),
        gr.Textbox(label="Min-p Sampling (as in https://arxiv.org/abs/2407.01082)"),
        gr.Textbox(label="Eta Sampling"),
        gr.Textbox(label="Epsilon Sampling"),
        gr.Textbox(
            label=(
                "laconic decoding "
                "(Alex Dimakis: https://x.com/AlexGDimakis/status/1885447830120362099)"
            )
        ),
        gr.Textbox(
            label="Power Sampling (Karan & Du, 2025; beta=2, 4 MCMC steps)"
        ),
        gr.Textbox(
            label="COT Decoding (Chain-of-Thought Reasoning without Prompting, Wang, Zhou, 2024)"
        ),
    ],
    title="Decoding Methods Comparison",
    description=(
        "Each decoding method's final answer is printed as soon as it is done. "
        f"Model used: {MODEL_NAME}."
    ),
)

if __name__ == "__main__":
    interface.launch()