from __future__ import annotations

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

from transformers.modeling_utils import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
from transformers.configuration_utils import PretrainedConfig

from typing import Optional, Union

# ---------------------------------------------------------------------------
# Inlined config + model — fully self-contained for HF trust_remote_code.
# ---------------------------------------------------------------------------

class StreamMixerConfig(PretrainedConfig):
    model_type = "streammixer"

    def __init__(
        self,
        vocab_size=32768,
        n_embd=768,
        n_layer=16,
        n_streams=48,
        stream_dim=96,
        n_read_heads=6,
        max_sequence_length=2048,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.n_embd = n_embd
        self.n_layer = n_layer
        self.n_streams = n_streams
        self.stream_dim = stream_dim
        self.n_read_heads = n_read_heads
        self.max_sequence_length = max_sequence_length
        self.hidden_size = n_embd
        self.num_hidden_layers = n_layer
        self.num_attention_heads = n_read_heads
        self.intermediate_size = (int(8 * n_embd / 3) + 15) // 16 * 16

# ---------------------------------------------------------------------------
# Inlined from model.py — keeps modeling.py self-contained for HF cache.
# ---------------------------------------------------------------------------

_DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
def _detect_compute_dtype():
    env = os.environ.get("NANOCHAT_DTYPE")
    if env is not None:
        return _DTYPE_MAP[env]
    if torch.cuda.is_available():
        capability = torch.cuda.get_device_capability()
        if capability >= (8, 0):
            return torch.bfloat16
        return torch.float32
    return torch.float32

COMPUTE_DTYPE = _detect_compute_dtype()


class Linear(nn.Linear):
    def forward(self, x):
        b = None if self.bias is None else self.bias.to(dtype=x.dtype)
        return F.linear(x, self.weight.to(dtype=x.dtype), b)


class RMSNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
    def forward(self, x):
        return F.rms_norm(x, (x.size(-1),))


class StreamMixer(nn.Module):
    LOG_A_MAX_NEG = 0.5
    CHUNK_SIZE    = 128

    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
        super().__init__()
        self.M = n_streams
        self.D = stream_dim
        self.H = n_read_heads
        self.wv = Linear(n_embd, stream_dim, bias=False)
        self.wq = Linear(n_embd, n_read_heads * stream_dim, bias=False)
        self.wr = Linear(n_embd, n_streams, bias=False)
        self.wa = Linear(n_embd, n_streams, bias=True)
        self.wo_out = Linear(n_read_heads * stream_dim, n_embd, bias=False)

    @staticmethod
    def _rms_norm_last(x):
        return F.rms_norm(x, (x.size(-1),))

    def _read(self, q_flat, s):
        lead = q_flat.shape[:-1]
        q = q_flat.view(*lead, self.H, self.D)
        q_n = self._rms_norm_last(q)
        s_n = self._rms_norm_last(s)
        scores = torch.einsum('...hd,...md->...hm', q_n, s_n) * (self.D ** -0.5)
        weights = torch.sigmoid(scores)
        read = torch.einsum('...hm,...md->...hd', weights, s)
        return read.reshape(*lead, self.H * self.D)

    def _scan(self, log_a, bv):
        B, T, M, D = bv.shape
        C = self.CHUNK_SIZE
        pad = (C - T % C) % C
        if pad > 0:
            log_a = F.pad(log_a, (0, 0, 0, pad))
            bv    = F.pad(bv,    (0, 0, 0, 0, 0, pad))
        Tp = T + pad
        n_chunks = Tp // C

        log_a_c = log_a.view(B, n_chunks, C, M)
        bv_c    = bv.view(B, n_chunks, C, M, D)

        Z = log_a_c.cumsum(dim=2)
        eZ = torch.exp(Z).unsqueeze(-1)
        inv_eZ = torch.exp(-Z).unsqueeze(-1)

        s_intra = (bv_c * inv_eZ).cumsum(dim=2) * eZ

        chunk_decay     = eZ[:, :, -1]
        chunk_end_intra = s_intra[:, :, -1]

        incomings = []
        state = torch.zeros(B, M, D, device=bv.device, dtype=bv.dtype)
        for c in range(n_chunks):
            incomings.append(state)
            state = chunk_decay[:, c] * state + chunk_end_intra[:, c]
        incoming = torch.stack(incomings, dim=1)

        s = s_intra + eZ * incoming.unsqueeze(2)
        return s.view(B, Tp, M, D)[:, :T]

    def forward(self, x):
        out, _ = self.forward_with_state(x)
        return out

    def forward_with_state(self, x):
        v = self.wv(x)
        q = self.wq(x)
        r = F.softmax(self.wr(x), dim=-1)
        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x))
        bv = r.unsqueeze(-1) * v.unsqueeze(2)
        s = self._scan(log_a, bv)
        read = self._read(q, s)
        return self.wo_out(read), s[:, -1]

    def step(self, x, state):
        x_t = x.squeeze(1)
        v = self.wv(x_t)
        q = self.wq(x_t)
        r = F.softmax(self.wr(x_t), dim=-1)
        log_a = -self.LOG_A_MAX_NEG * torch.sigmoid(self.wa(x_t))
        a = torch.exp(log_a)
        write = r.unsqueeze(-1) * v.unsqueeze(1)
        s = a.unsqueeze(-1) * state + write
        read = self._read(q, s)
        return self.wo_out(read).unsqueeze(1), s


class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        hidden = (int(8 * n_embd / 3) + 15) // 16 * 16
        self.w_up   = Linear(n_embd, hidden, bias=False)
        self.w_down = Linear(n_embd, hidden, bias=False)
        self.w_out  = Linear(hidden, n_embd, bias=False)
    def forward(self, x):
        return self.w_out(F.relu(self.w_up(x)).square() * self.w_down(x))


class Block(nn.Module):
    def __init__(self, n_embd, n_streams, stream_dim, n_read_heads):
        super().__init__()
        self.ln1 = RMSNorm(n_embd)
        self.mix = StreamMixer(n_embd, n_streams, stream_dim, n_read_heads)
        self.ln2 = RMSNorm(n_embd)
        self.mlp = MLP(n_embd)
    def forward(self, x):
        x, _ = self.forward_with_state(x)
        return x
    def forward_with_state(self, x):
        mix_out, state = self.mix.forward_with_state(self.ln1(x))
        x = x + mix_out
        x = x + self.mlp(self.ln2(x))
        return x, state
    def step(self, x, state):
        mix_out, new_state = self.mix.step(self.ln1(x), state)
        x = x + mix_out
        x = x + self.mlp(self.ln2(x))
        return x, new_state


class GPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_layer, n_streams, stream_dim, n_read_heads):
        super().__init__()
        padded_vocab_size = ((vocab_size + 63) // 64) * 64
        self.config = dict(
            vocab_size=padded_vocab_size, n_embd=n_embd, n_layer=n_layer,
            n_streams=n_streams, stream_dim=stream_dim, n_read_heads=n_read_heads,
        )
        self.wte = nn.Embedding(padded_vocab_size, n_embd, dtype=COMPUTE_DTYPE)
        self.ln0 = RMSNorm(n_embd)
        self.blocks = nn.ModuleList([
            Block(n_embd, n_streams, stream_dim, n_read_heads) for _ in range(n_layer)
        ])
        self.lm_head = Linear(n_embd, padded_vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight

    def forward(self, token_ids):
        logits, _ = self.forward_with_states(token_ids)
        return logits

    def forward_with_states(self, token_ids):
        x = self.ln0(self.wte(token_ids))
        states = []
        for block in self.blocks:
            x, state = block.forward_with_state(x)
            states.append(state)
        return self.lm_head(x), states

    def step(self, token_ids, states):
        x = self.ln0(self.wte(token_ids))
        new_states = []
        for block, state in zip(self.blocks, states):
            x, new_state = block.step(x, state)
            new_states.append(new_state)
        return self.lm_head(x), new_states

    def initial_states(self, batch_size, device, dtype=COMPUTE_DTYPE):
        return [torch.zeros(batch_size, b.mix.M, b.mix.D, device=device, dtype=dtype)
                for b in self.blocks]

    def get_memory_footprint(self, return_buffers=True):
        mem = sum(p.nelement() * p.element_size() for p in self.parameters())
        if return_buffers:
            mem += sum(b.nelement() * b.element_size() for b in self.buffers())
        return mem

    @classmethod
    def from_config(cls, config):
        return cls(**config)

    def to_compute_dtype(self, dtype=None):
        dtype = dtype or COMPUTE_DTYPE
        self.to(dtype)
        return self


# ---------------------------------------------------------------------------
# HF Model wrappers
# ---------------------------------------------------------------------------

class StreamMixerPreTrainedModel(PreTrainedModel):
    config_class = StreamMixerConfig
    supports_gradient_checkpointing = False
    base_model_prefix = "model"
    _tied_weights_keys = []

    def __init__(self, config, *args, **kwargs):
        super().__init__(config, *args, **kwargs)
        self.all_tied_weights_keys = {}

    def _set_gradient_checkpointing(self, module, value=False):
        raise NotImplementedError("Gradient checkpointing is not supported")

    def _init_weights(self, module):
        pass


class StreamMixerModel(StreamMixerPreTrainedModel):
    def __init__(self, config: StreamMixerConfig, **kwargs):
        super().__init__(config, **kwargs)
        self.config = config
        self._tied_weights_keys = []
        inner_cfg = {
            'vocab_size': config.vocab_size,
            'n_embd': config.n_embd,
            'n_layer': config.n_layer,
            'n_streams': config.n_streams,
            'stream_dim': config.stream_dim,
            'n_read_heads': config.n_read_heads,
        }
        self.inner = GPT.from_config(inner_cfg)
        self.inner.to(COMPUTE_DTYPE)
        self.post_init()

    def get_input_embeddings(self):
        return self.inner.wte

    def set_input_embeddings(self, value):
        self.inner.wte = value

    def get_output_embeddings(self):
        return self.inner.lm_head

    def set_output_embeddings(self, value):
        self.inner.lm_head = value

    def _get_hidden_states(self, input_ids):
        x = self.inner.ln0(self.inner.wte(input_ids))
        for block in self.inner.blocks:
            x = block(x)
        return x.float()

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[tuple, BaseModelOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        hidden_states = self._get_hidden_states(input_ids)
        if not return_dict:
            return (hidden_states,)
        return BaseModelOutput(last_hidden_state=hidden_states)


class StreamMixerForCausalLM(StreamMixerPreTrainedModel):
    _keys_to_ignore_on_load_unexpected = set()

    def __init__(self, config: StreamMixerConfig, **kwargs):
        super().__init__(config, **kwargs)
        self.model = StreamMixerModel(config, **kwargs)
        self.vocab_size = config.vocab_size
        self.config = config
        self.post_init()

    def __getattr__(self, name):
        if name == "all_tied_weights_keys":
            return {}
        return super().__getattr__(name)

    def get_input_embeddings(self):
        return self.model.inner.wte

    def set_input_embeddings(self, value):
        self.model.inner.wte = value

    def get_output_embeddings(self):
        return self.model.inner.lm_head

    def set_output_embeddings(self, value):
        self.model.inner.lm_head = value

    def can_generate(self):
        return True

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Union[tuple, CausalLMOutput]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        hidden_states = self.model._get_hidden_states(input_ids)
        logits = self.model.inner.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

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