import time
import warnings
import logging

warnings.filterwarnings("ignore", message = "cudaGetDeviceCount.*")
logging.getLogger("torch.utils.flop_counter").setLevel(logging.ERROR)

import torch
import torch.nn as nn
import torch.nn.functional as F

from xlstm.xlstm_large.model import (
    xLSTMLargeConfig, xLSTMLargeBlockStack,
)

from transformers import PreTrainedModel, PretrainedConfig

class SnortleConfig(PretrainedConfig):
    model_type = "snortle_pancake_2"

    # these numbers must match the ones used in train.py
    def __init__(self, vocab_size = 12288, d_model = 768, num_blocks = 27, num_heads = 6, **kwargs):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_blocks = num_blocks
        self.num_heads = num_heads

        super().__init__(**kwargs)

class xLSTMStack(nn.Module):
    def __init__(self, vocab_size, num_heads, d_model, num_blocks):
        super().__init__()

        xlstm_config = xLSTMLargeConfig(
            num_blocks = num_blocks,
            embedding_dim = d_model,
            num_heads = num_heads,
            vocab_size = vocab_size,
            chunkwise_kernel = "chunkwise--native_autograd",
            sequence_kernel = "native_sequence__native",
            step_kernel = "native",
            mode = "inference",
            chunk_size = 64,
            return_last_states = True,
        )

        self.xlstm_stack = xLSTMLargeBlockStack(xlstm_config)

    def forward(self, x, state):
        out, state = self.xlstm_stack(x, state = state)

        return out, state

class Model(PreTrainedModel):

    config_class = SnortleConfig

    def __init__(self, config, tokenizer = None, vocab_size = None, context_length = 512):
        super().__init__(config)

        if vocab_size is None:
            vocab_size = config.vocab_size

        self.tokenizer = tokenizer
        self.d_model = config.d_model
        self.context_length = context_length

        self.embedding = nn.Embedding(
            num_embeddings = vocab_size,
            embedding_dim = self.d_model
        )

        self.norm1 = nn.RMSNorm(self.d_model)

        self.xlstm_stack = xLSTMStack(
            vocab_size = vocab_size,
            num_heads = config.num_heads,
            d_model = self.d_model,
            num_blocks = config.num_blocks,
        )

        self.norm2 = nn.RMSNorm(self.d_model)

        self.lm_head = nn.Linear(
            in_features = self.d_model,
            out_features = vocab_size,
            bias = False
        )

        self.lm_head.weight = self.embedding.weight

        self.post_init()

    # cache_position is ignored. It's only here so older code that passes it still works.
    def forward(self, x, state = None, cache_position = None):
        x = self.embedding(x)
        x = self.norm1(x)

        x, state = self.xlstm_stack(x, state)
        x = self.norm2(x)

        x = self.lm_head(x)

        return x, state

    def generate(self, tokenizer, user_input, state, max_tokens = 64, top_k = 40, temperature = 0.8, repeat_penalty = 1.0, repeat_window = 64):
        # keep the settings in a range that works
        top_k = max(1, min(top_k, self.embedding.num_embeddings))
        temperature = max(temperature, 0.05)

        with torch.inference_mode():
            output_chunks = []
            current_chunk = []

            start_time = time.perf_counter()

            token_ids = tokenizer.encode(user_input)

            # the tokenizer adds </s> to the end, but the model never learned to continue after it
            if token_ids[-1] == tokenizer.eos_token_id:
                token_ids = token_ids[:-1]

            input_tensor = torch.tensor(token_ids, dtype = torch.long).unsqueeze(0)
            out, state = self(input_tensor, state)

            think_time = time.perf_counter() - start_time

            print(f"\rThought for {think_time:.2f} seconds.")

            start_time = time.perf_counter()

            # Helper function to sample the next token
            def sample_next(logits, generated):
                logits = logits[:, -1, :] / temperature

                # Tokens used recently get a small penalty, so it doesn't repeat itself in loops
                if repeat_penalty > 0 and len(generated) > 0:
                    recent = list(set(generated[-repeat_window:]))
                    logits[0, recent] = logits[0, recent] - repeat_penalty

                values, indices = torch.topk(logits, top_k)
                probs = F.softmax(values, dim = -1)
                choice = torch.multinomial(probs, num_samples = 1)
                return indices.gather(-1, choice)

            generated = []

            # Sample first token
            tok_id = sample_next(out, generated)
            next_input = tok_id
            gen_token_count = 0

            chunk_size = 16 # Change to 32 if you want larger batches

            if tok_id.item() != tokenizer.eos_token_id:
                generated.append(tok_id.item())
                current_chunk.append(tok_id.item())
                gen_token_count += 1

                for t in range(max_tokens - 1):
                    out, state = self(next_input.to(torch.long), state)
                    tok_id = sample_next(out, generated)
                    next_input = tok_id
                    gen_token_count += 1

                    if tok_id.item() == tokenizer.eos_token_id:
                        break

                    generated.append(tok_id.item())
                    current_chunk.append(tok_id.item())

                    # When the chunk hits the target size, push it and flush
                    if len(current_chunk) >= chunk_size:
                        output_chunks.append(tokenizer.decode(current_chunk))
                        current_chunk = []

                # Flush any remaining tokens in the final chunk
                if current_chunk:
                    output_chunks.append(tokenizer.decode(current_chunk))

            generation_time = time.perf_counter() - start_time

            # Join chunks into the final string (or you can return output_chunks as a list if your UI handles it)
            output = "".join(output_chunks)
            print("\n")
            return output, think_time, generation_time, gen_token_count, state