import math
from typing import Optional

import numpy as np
import torch
from transformers import GenerationConfig
from transformers.generation import LogitsProcessorList

import kgen.models as models
import kgen.executor.tipo as tipo
from kgen.formatter import seperate_tags, apply_format
from kgen.sampling import (
    LogitsRecorder,
    DEFAULT_FORMAT,
    clone_kv,
)


def beam_search_sample(prompt, variations=7):
    # recorder
    total_gen = 0

    # Initialize with first tokens
    inputs = models.tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(next(models.text_model.parameters()).device)
    input_length = input_ids.shape[-1]

    beam_width = min(variations * 3, 32)  # Wider beam for more diversity
    active_beams = [(input_ids, None, 1.0)]  # (input_ids, past_kv, cumulative_score)
    completed_sequences = []

    # Setup logits processors for scoring and diversity
    recorder = LogitsRecorder()
    processors = LogitsProcessorList()
    processors.append(recorder)

    while len(completed_sequences) < variations:
        next_beams = []

        # Expand each active beam
        for beam_ids, beam_past_kv, beam_score in active_beams:
            # Generate next token distributions
            recorder.clean()
            gen_kwargs = {}
            gen_kwargs["max_new_tokens"] = 1
            gen_kwargs["return_dict_in_generate"] = True
            gen_kwargs["output_scores"] = True
            gen_kwargs["do_sample"] = True
            generation_config = GenerationConfig(**gen_kwargs)
            with torch.no_grad():
                generation_output = models.text_model.generate(
                    input_ids=beam_ids.clone(),
                    generation_config=generation_config,
                    logits_processor=processors,
                    past_key_values=beam_past_kv,
                )
            total_gen += 1

            # Sample multiple candidates from the distribution
            logits = recorder.scores[-1][0]  # Get logits for last token
            probs = torch.softmax(logits, dim=-1)

            # Sample without replacement for diversity
            num_samples = min(beam_width, 16)  # Limit samples per beam
            candidates = torch.multinomial(probs, num_samples, replacement=False)
            for candidate in candidates:
                new_ids = beam_ids.clone()
                new_ids = torch.cat(
                    [new_ids, candidate.unsqueeze(0).unsqueeze(0)], dim=-1
                )

                # Calculate sequence score
                candidate_prob = probs[candidate].item()
                new_score = beam_score + math.log(candidate_prob)

                # Check if sequence is complete
                if candidate.item() == models.tokenizer.eos_token_id:
                    decoded = models.tokenizer.decode(new_ids[0])
                    if decoded not in [seq[0] for seq in completed_sequences]:
                        print(
                            f"Completed: {len(completed_sequences)} ({new_score:.4f})"
                        )
                        completed_sequences.append((decoded, new_score))
                else:
                    next_beams.append(
                        (
                            new_ids,
                            clone_kv(generation_output.past_key_values),
                            new_score,
                        )
                    )

        # Prune and select top beams
        next_beams.sort(key=lambda x: x[2], reverse=True)
        active_beams = next_beams[:beam_width]

        # Break if no active beams
        if not active_beams:
            break

        # Avoid infinite loops
        if any(ids.shape[-1] > 1024 for ids, _, _ in active_beams):
            break

    print("Total output tokens:", total_gen)
    print("Completed sequences:", len(completed_sequences))
    # If we don't have enough sequences, return what we have
    return completed_sequences[:variations]


def stochastic_beam_search(prompt, variations=7, temperature=1.0):
    # recorder
    total_gen = 0

    # Initialize with first tokens
    inputs = models.tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(next(models.text_model.parameters()).device)
    input_length = input_ids.shape[-1]

    beam_width = min(variations * 3, 32)
    active_beams = [(input_ids, None, 0.0)]  # Using log probabilities
    completed_sequences = []

    # Setup processors
    recorder = LogitsRecorder()
    processors = LogitsProcessorList()
    processors.append(recorder)

    while len(completed_sequences) < variations:
        next_beams = []

        for beam_ids, beam_past_kv, log_score in active_beams:
            recorder.clean()
            gen_kwargs = {
                "max_new_tokens": 1,
                "return_dict_in_generate": True,
                "output_scores": True,
                "do_sample": True,
            }
            generation_config = GenerationConfig(**gen_kwargs)

            with torch.no_grad():
                generation_output = models.text_model.generate(
                    input_ids=beam_ids.clone(),
                    generation_config=generation_config,
                    logits_processor=processors,
                    past_key_values=beam_past_kv,
                )
            total_gen += 1

            # Get logits and apply temperature + min_p
            logits = recorder.scores[-1][0]
            logits = logits / temperature

            # Get log probabilities
            log_probs = torch.log_softmax(logits, dim=-1)
            probs = torch.softmax(logits, dim=-1)

            # Apply min_p filtering
            # mask = probs < min_p * probs.max()
            # log_probs[mask] = float("-inf")

            # Add Gumbel noise for stochastic sampling
            gumbel_noise = -torch.log(-torch.log(torch.rand_like(log_probs) + 1e-10))
            perturbed_logprobs = log_probs + gumbel_noise

            # Get top-k candidates using perturbed scores
            num_samples = min(beam_width, 16)
            topk_values, topk_indices = perturbed_logprobs.topk(num_samples)

            for value, token_id in zip(topk_values, topk_indices):
                new_ids = beam_ids.clone()
                new_ids = torch.cat(
                    [new_ids, token_id.unsqueeze(0).unsqueeze(0)], dim=-1
                )

                # Update sequence log probability
                new_log_score = log_score + log_probs[token_id].item()

                if token_id.item() == models.tokenizer.eos_token_id:
                    decoded = models.tokenizer.decode(new_ids[0])
                    if decoded not in [seq[0] for seq in completed_sequences]:
                        print(
                            f"Completed: {len(completed_sequences)} ({math.exp(new_log_score):.4f})"
                        )
                        completed_sequences.append((decoded, new_log_score))
                else:
                    next_beams.append(
                        (
                            new_ids,
                            clone_kv(generation_output.past_key_values),
                            new_log_score,
                        )
                    )

        # Select top beams based on perturbed scores
        next_beams.sort(key=lambda x: x[2], reverse=True)
        active_beams = next_beams[:beam_width]

        if not active_beams:
            break

        if any(ids.shape[-1] > 1024 for ids, _, _ in active_beams):
            break

    print("Total output tokens:", total_gen)
    print("Completed sequences:", len(completed_sequences))

    # Convert log probabilities back to probabilities for return
    return [(seq, math.exp(score)) for seq, score in completed_sequences[:variations]]


def diverse_beam_search(
    prompt,
    num_sequences: int = 16,  # Target number of sequences
    num_groups: Optional[int] = None,  # Optional: specify number of groups
    diversity_penalty: float = 0.75,
    temperature: float = 1.5,
):
    """
    Implements diverse beam search with automatic group/beam width calculation.

    Args:
        prompt: Input prompt text
        num_sequences: Target number of sequences to generate
        num_groups: Optional number of groups (if None, will be calculated)
        diversity_penalty: Penalty factor for repeated tokens across groups
        temperature: Temperature for logits scaling
    """
    # Calculate groups and beam width if not specified
    if num_groups is None:
        # Square root gives a balanced split between groups and beams
        num_groups = int(math.sqrt(num_sequences))

    # Calculate beam width needed per group to get desired sequences
    beam_width_per_group = math.ceil(num_sequences / num_groups)

    # Limit beam width to prevent memory issues
    beam_width_per_group = min(beam_width_per_group, 16)

    print(f"Using {num_groups} groups with {beam_width_per_group} beams each")

    # Initialize with first tokens
    inputs = models.tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to(next(models.text_model.parameters()).device)

    # Track beams per group
    group_beams = [
        [(input_ids, None, 0.0)] for _ in range(num_groups)
    ]  # (ids, past_kv, score)
    completed_sequences = []
    total_gen = 0

    # Setup recorder and processors
    recorder = LogitsRecorder()
    processors = LogitsProcessorList()
    processors.append(recorder)

    while True:
        next_group_beams = [[] for _ in range(num_groups)]

        # Process each group sequentially
        for group_id, active_beams in enumerate(group_beams):
            # Track tokens selected by previous groups for diversity
            prev_group_tokens = set()
            if group_id > 0:
                for prev_group in range(group_id):
                    for ids, _, _ in group_beams[prev_group]:
                        if ids.shape[-1] > 1:
                            prev_group_tokens.add(ids[0, -1].item())

            # Expand each beam in current group
            for beam_ids, beam_past_kv, log_score in active_beams:
                recorder.clean()
                gen_kwargs = {
                    "max_new_tokens": 1,
                    "return_dict_in_generate": True,
                    "output_scores": True,
                    "do_sample": True,
                }
                generation_config = GenerationConfig(**gen_kwargs)

                with torch.no_grad():
                    generation_output = models.text_model.generate(
                        input_ids=beam_ids.clone(),
                        generation_config=generation_config,
                        logits_processor=processors,
                        past_key_values=beam_past_kv,
                    )
                total_gen += 1

                # Get logits and apply temperature
                logits = recorder.scores[-1][0] / temperature

                # Apply diversity penalty
                if prev_group_tokens:
                    diversity_mask = torch.zeros_like(logits)
                    diversity_mask[list(prev_group_tokens)] = -diversity_penalty
                    logits = logits + diversity_mask

                # Get log probabilities
                log_probs = torch.log_softmax(logits, dim=-1)

                # Sample top candidates for this beam
                topk_values, topk_indices = log_probs.topk(beam_width_per_group)

                for value, token_id in zip(topk_values, topk_indices):
                    new_ids = beam_ids.clone()
                    new_ids = torch.cat(
                        [new_ids, token_id.unsqueeze(0).unsqueeze(0)], dim=-1
                    )
                    new_log_score = log_score + value.item()

                    if token_id.item() == models.tokenizer.eos_token_id:
                        decoded = models.tokenizer.decode(new_ids[0])
                        if decoded not in [seq[0] for seq in completed_sequences]:
                            score = math.exp(new_log_score)
                            print(
                                f"Completed (Group {group_id}): {len(completed_sequences)} ({score:.4f})"
                            )
                            completed_sequences.append(
                                (decoded, new_log_score, group_id)
                            )
                    else:
                        next_group_beams[group_id].append(
                            (
                                new_ids,
                                clone_kv(generation_output.past_key_values),
                                new_log_score,
                            )
                        )

            # Sort and prune beams within group
            next_group_beams[group_id].sort(key=lambda x: x[2], reverse=True)
            next_group_beams[group_id] = next_group_beams[group_id][
                :beam_width_per_group
            ]

        # Update all groups
        group_beams = next_group_beams

        # Check stopping conditions
        if not any(beams for beams in group_beams):
            break

        if any(
            any(ids.shape[-1] > 1024 for ids, _, _ in beams) for beams in group_beams
        ):
            break

        # Break if we have enough sequences
        if len(completed_sequences) >= num_sequences:
            break

    print("Total output tokens:", total_gen)
    print("Completed sequences:", len(completed_sequences))
    print("Groups represented:", len(set(g for _, _, g in completed_sequences)))

    # Convert log probabilities to probabilities and sort
    final_results = []
    for seq, score, group in completed_sequences:
        final_results.append((seq, math.exp(score)))

    # Sort by score and return exactly num_sequences results
    final_results.sort(key=lambda x: x[1], reverse=True)
    return final_results[:num_sequences]


if __name__ == "__main__":
    models.load_model(
        "KBlueLeaf/TIPO-100M",
        device="cuda",
    )

    meta, operations, general, prompt = tipo.parse_tipo_request(
        seperate_tags("scenery, wide shot, masterpiece, safe".split(",")),
        "",
    )
    mode, length, expand = operations[0]
    prompt = tipo.apply_tipo_prompt(meta, general, prompt, mode, length, expand)

    # results = beam_search_sample(prompt, 1024)
    results = diverse_beam_search(
        prompt,
        num_sequences=1024,  # Number of sequences to generate
        diversity_penalty=0.75,  # How strongly to encourage diversity
        temperature=1.5,  # Temperature for sampling
    )
    gen_per_prompt = [x[1] for x in results]
    print(sum(gen_per_prompt) / len(gen_per_prompt))
    with open("./test/div_beam_search.txt", "w", encoding="utf-8") as f:
        for result, gen in sorted(results):
            result = tipo.parse_tipo_result(result)
            formatted_output = apply_format(result, DEFAULT_FORMAT)
            f.write(formatted_output + "\n")
    # for result in sorted(results):
    #     print("=" * 20)
    #     print(result)
    # print("=" * 20)
