from __future__ import annotations


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


from huggingface_hub import hf_hub_download
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput


try:
    from .configuration_transwave import TransWaveConfig
    from .transwave_core import (
        CTWModel,
        BYTE_VOCAB,
        BYTE_DIM,
        BYTES_PER_BLOCK,
        BYTE_DECODER_DIM,
    )
except ImportError:
    from configuration_transwave import TransWaveConfig
    from transwave_core import (
        CTWModel,
        BYTE_VOCAB,
        BYTE_DIM,
        BYTES_PER_BLOCK,
        BYTE_DECODER_DIM,
    )




EXPECTED_PARAMS = 96_130_040




class TransWaveForCausalLM(PreTrainedModel):
    config_class = TransWaveConfig
    base_model_prefix = "transwave"
    main_input_name = "input_ids"
    _tied_weights_keys = ["transwave.decoder.shared_embedding.weight"]


    def __init__(self, config):
        super().__init__(config)
        self.transwave = CTWModel()
        # Do not call post_init(): CTWModel already performs the architecture's
        # exact initialization. The trained checkpoint is loaded afterwards.


    def get_input_embeddings(self):
        return self.transwave.tokenizer.byte_embedding


    def set_input_embeddings(self, value):
        self.transwave.tokenizer.byte_embedding = value
        self.transwave.decoder.shared_embedding = value


    def get_output_embeddings(self):
        return self.transwave.decoder.shared_embedding


    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        revision = kwargs.pop("revision", None)
        cache_dir = kwargs.pop("cache_dir", None)
        token = kwargs.pop("token", None)
        local_files_only = kwargs.pop("local_files_only", False)
        torch_dtype = kwargs.pop("torch_dtype", None)
        kwargs.pop("trust_remote_code", None)
        kwargs.pop("use_safetensors", None)


        if config is None:
            config = TransWaveConfig.from_pretrained(
                pretrained_model_name_or_path,
                revision=revision,
                cache_dir=cache_dir,
                token=token,
                local_files_only=local_files_only,
            )


        model = cls(config)


        path = os.fspath(pretrained_model_name_or_path)
        if os.path.isdir(path):
            ckpt_path = os.path.join(path, "best_model.pt")
        else:
            ckpt_path = hf_hub_download(
                repo_id=path,
                filename="best_model.pt",
                revision=revision,
                cache_dir=cache_dir,
                token=token,
                local_files_only=local_files_only,
            )


        checkpoint = torch.load(
            ckpt_path,
            map_location="cpu",
            weights_only=False,
        )
        state = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint


        incompatible = model.transwave.load_state_dict(state, strict=True)
        if incompatible.missing_keys or incompatible.unexpected_keys:
            raise RuntimeError(
                f"Checkpoint mismatch. Missing={incompatible.missing_keys}, "
                f"unexpected={incompatible.unexpected_keys}"
            )


        n_params = sum(p.numel() for p in model.transwave.parameters())
        if n_params != EXPECTED_PARAMS:
            raise RuntimeError(f"Parameter count {n_params:,} != {EXPECTED_PARAMS:,}")


        if torch_dtype not in (None, "auto"):
            model = model.to(dtype=torch_dtype)


        model.eval()
        return model


    def nll(self, input_ids, checkpoint_layers=False):
        return self.transwave.nll(
            input_ids,
            checkpoint_layers=checkpoint_layers,
        )


    def contexts(self, input_ids, checkpoint_layers=False):
        return self.transwave.contexts(
            input_ids,
            checkpoint_layers=checkpoint_layers,
        )


    def teacher_forced_logits(self, input_ids):
        context, target_blocks = self.transwave.contexts(
            input_ids,
            checkpoint_layers=False,
        )
        decoder = self.transwave.decoder
        batch, blocks, _ = context.shape
        n = batch * blocks


        hidden = torch.tanh(
            decoder.context_proj(context)
        ).reshape(n, BYTE_DECODER_DIM)


        previous = decoder.bos[None, :].expand(n, -1)
        position_logits = []


        for r in range(BYTES_PER_BLOCK):
            decoder_input = previous + decoder.byte_position[r][None, :]
            hidden = decoder.gru(decoder_input, hidden)
            z = decoder.to_byte_space(decoder.norm(hidden))
            logits = F.linear(z, decoder.shared_embedding.weight)
            position_logits.append(
                logits.reshape(batch, blocks, BYTE_VOCAB)
            )
            target = target_blocks[:, :, r].reshape(-1).long()
            previous = decoder.shared_embedding(target)


        logits = torch.stack(position_logits, dim=2)
        return logits.reshape(
            batch,
            blocks * BYTES_PER_BLOCK,
            BYTE_VOCAB,
        )


    def forward(
        self,
        input_ids=None,
        labels=None,
        return_dict=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids are required")
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, tokens]")
        if input_ids.shape[1] % BYTES_PER_BLOCK != 0:
            raise ValueError("TransWave scoring length must be divisible by 8")


        if labels is not None:
            if labels.shape != input_ids.shape or not torch.equal(labels, input_ids):
                raise ValueError(
                    "TransWave performs the causal shift internally; "
                    "for exact scoring/training, labels must equal input_ids."
                )


        logits = self.teacher_forced_logits(input_ids)
        loss = self.transwave.nll(
            input_ids,
            checkpoint_layers=False,
        )


        if return_dict is False:
            return (loss, logits)


        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )


    @torch.no_grad()
    def generate_bytes(
        self,
        input_ids,
        max_new_tokens=32,
        do_sample=False,
        temperature=1.0,
        top_k=None,
    ):
        if input_ids.ndim != 2 or input_ids.shape[0] != 1:
            raise ValueError("generate_bytes currently supports batch size 1")
        out = input_ids.long().clone()
        device = next(self.parameters()).device
        out = out.to(device)


        for _ in range(int(max_new_tokens)):
            total = out.shape[1]
            complete = (total // BYTES_PER_BLOCK) * BYTES_PER_BLOCK
            r = total - complete


            if complete == 0:
                block_context = self.transwave.bos_context[None, :].to(device)
            else:
                _, blocks = self.transwave.encode_completed_blocks(
                    out[:, :complete],
                    checkpoint_layers=False,
                )
                h, _ = self.transwave.encode_completed_blocks(
                    out[:, :complete],
                    checkpoint_layers=False,
                )
                block_context = h[:, -1, :]


            decoder = self.transwave.decoder
            hidden = torch.tanh(decoder.context_proj(block_context))
            previous = decoder.bos[None, :]


            current = out[:, complete:]
            logits = None
            for j in range(r + 1):
                decoder_input = previous + decoder.byte_position[j][None, :]
                hidden = decoder.gru(decoder_input, hidden)
                z = decoder.to_byte_space(decoder.norm(hidden))
                logits = F.linear(z, decoder.shared_embedding.weight)
                if j < r:
                    previous = decoder.shared_embedding(current[:, j])


            scores = logits.float() / max(float(temperature), 1e-6)
            if top_k is not None and int(top_k) > 0:
                k = min(int(top_k), BYTE_VOCAB)
                threshold = torch.topk(scores, k, dim=-1).values[:, -1:]
                scores = scores.masked_fill(scores < threshold, float("-inf"))


            if do_sample:
                next_id = torch.multinomial(
                    torch.softmax(scores, dim=-1),
                    num_samples=1,
                )
            else:
                next_id = scores.argmax(dim=-1, keepdim=True)


            out = torch.cat([out, next_id], dim=1)


        return out