# modeling_fwkv.py
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, AutoModel, AutoConfig
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation import GenerationMixin
from transformers.cache_utils import Cache  # needed to detect Cache objects
from .configuration_fwkv import FWKVConfig


# ════════════════════════════════════════════════════════════════════════════
# utility functions
# ════════════════════════════════════════════════════════════════════════════
def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
    """REFERENCE implementation (Hillis-Steele). y_t = a_t + W * y_{t-1}.
    Kept for parity/testing with decay_scan; not used at inference time."""
    W = W.to(dtype=a.dtype)
    val = a
    T = a.shape[1]
    d = 1
    while d < T:
        shifted = F.pad(val[:, :-d, :], (0, 0, d, 0))
        val = val + (W ** d) * shifted
        d *= 2
    return val


def decay_scan(a: torch.Tensor, W: torch.Tensor, state: torch.Tensor | None = None) -> torch.Tensor:
    """y_t = a_t + W * y_{t-1}, y_{-1} = state, as a causal FFT convolution in fp32.

    a: (B, T, d), W: (d,) with 0 < W < 1, state: (B, d) or None.
    Matches the training-time implementation exactly (same math as
    parallel_scan_decay, but computed as an FFT convolution for speed).
    """
    with torch.autocast(device_type=a.device.type, enabled=False):
        a = a.float()
        W = W.float()
        B, T, d = a.shape
        if state is not None:
            first = a[:, :1] + W.view(1, 1, d) * state.float().unsqueeze(1)
            a = torch.cat([first, a[:, 1:]], dim=1)
        if T == 1:
            return a
        n = 1 << (2 * T - 1).bit_length()          # >= 2T-1, avoids circular wrap
        t = torch.arange(T, device=a.device, dtype=a.dtype)
        K = torch.exp(t.unsqueeze(1) * torch.log(W).unsqueeze(0))   # (T, d) = W^t
        Fa = torch.fft.rfft(a, n=n, dim=1)
        Fk = torch.fft.rfft(K, n=n, dim=0)
        y = torch.fft.irfft(Fa * Fk.unsqueeze(0), n=n, dim=1)[:, :T]
        return y.contiguous()


def chunked_cross_entropy(x, weight, labels, chunk_size=512, ignore_index=-100):
    x_2d = x.reshape(-1, x.size(-1))
    labels_1d = labels.reshape(-1)
    N = x_2d.size(0)
    total_loss = x_2d.new_zeros((), dtype=torch.float32)
    total_count = labels_1d.new_zeros(())
    for start in range(0, N, chunk_size):
        end = min(start + chunk_size, N)
        label_chunk = labels_1d[start:end]
        count = (label_chunk != ignore_index).sum()
        if count == 0:
            continue
        logits_chunk = F.linear(x_2d[start:end], weight)
        loss_chunk = F.cross_entropy(logits_chunk.float(), label_chunk, ignore_index=ignore_index, reduction="sum")
        total_loss = total_loss + loss_chunk
        total_count = total_count + count
    return total_loss / total_count.clamp(min=1)


# ════════════════════════════════════════════════════════════════════════════
# modules — mirror the trained Myosotis v2 / FWKVBlock exactly
# ════════════════════════════════════════════════════════════════════════════
class FactorizedTiedHead(nn.Module):
    def __init__(self, vocab_size, d_model, d_emb):
        super().__init__()
        self.d_model = d_model
        self.d_emb = d_emb
        self.weight = nn.Parameter(torch.empty(vocab_size, d_emb))
        self.proj = nn.Linear(d_emb, d_model, bias=False)
        nn.init.normal_(self.weight, std=0.02)
        nn.init.normal_(self.proj.weight, std=0.02)

    def embed(self, input_ids):
        return self.proj(F.embedding(input_ids, self.weight))

    # kept as an alias so either calling convention works
    def forward(self, input_ids):
        return self.embed(input_ids)

    def to_emb_space(self, x):
        return F.linear(x, self.proj.weight.t())

    def logits(self, x_emb):
        return F.linear(x_emb, self.weight)

    def resize_vocab(self, new_vocab_size):
        old = self.weight.shape[0]
        if new_vocab_size == old:
            return
        new_w = torch.empty(new_vocab_size, self.d_emb, device=self.weight.device, dtype=self.weight.dtype)
        nn.init.normal_(new_w, std=0.02)
        new_w[:old] = self.weight.data
        self.weight = nn.Parameter(new_w)


class FWKVBlock(nn.Module):
    """Pre-norm block: token-shifted WKV mixer + token-shifted FFN.

    Per-layer recurrent state = (wkv_state, shift_state_timemix, shift_state_chanmix),
    each (B, d). Must match the training-time block bit-for-bit (module names and
    math), since checkpoints are loaded by state_dict key.
    """

    def __init__(self, d, layer_id, n_layers, ffn_mult=4, floor=0.1, ceil=0.999,
                 head_size=64, tau_min=2.0, tau_max=1000.0):
        super().__init__()
        assert d % head_size == 0, "d_model must be divisible by head_size"
        self.d = d
        self.floor = floor
        self.ceil = ceil

        self.ln1 = nn.LayerNorm(d)
        self.ln2 = nn.LayerNorm(d)
        self.proj_k = nn.Linear(d, d, bias=False)
        self.proj_v = nn.Linear(d, d, bias=False)
        self.proj_r = nn.Linear(d, d, bias=False)
        self.proj_out = nn.Linear(d, d, bias=False)
        self.proj_out._is_out_proj = True
        self.ln_x = nn.GroupNorm(d // head_size, d)  # per-head norm on mixer output
        self.ffn = nn.Sequential(
            nn.Linear(d, ffn_mult * d, bias=False),
            nn.GELU(),
            nn.Linear(ffn_mult * d, d, bias=False),
        )
        self.ffn[2]._is_out_proj = True

        # Decay: channels span timescales tau in [tau_min, tau_max] (W = 1 - 1/tau).
        tau = torch.exp(torch.linspace(math.log(tau_min), math.log(tau_max), d))
        W0 = (1.0 - 1.0 / tau).clamp(floor + 1e-3, ceil - 1e-4)
        p = (W0 - floor) / (ceil - floor)
        self.w = nn.Parameter(torch.log(p / (1.0 - p)))

        # Token-shift mix coefficients (RWKV-4 style init).
        ratio_1_to_almost0 = 1.0 - (layer_id / n_layers)
        ratio_0_to_1 = layer_id / max(1, n_layers - 1)
        ramp = torch.arange(d, dtype=torch.float32) / d
        self.mu_k = nn.Parameter(ramp ** ratio_1_to_almost0)
        self.mu_v = nn.Parameter(ramp ** ratio_1_to_almost0 + 0.3 * ratio_0_to_1)
        self.mu_r = nn.Parameter(ramp ** (0.5 * ratio_1_to_almost0))
        self.mu_c = nn.Parameter(ramp ** ratio_1_to_almost0)

    @property
    def W(self):
        return self.floor + (self.ceil - self.floor) * torch.sigmoid(self.w)

    def forward(self, x, state=None):
        B, T, d = x.shape
        if state is None:
            wkv_s = None
            sh_tm = x.new_zeros(B, d)
            sh_cm = x.new_zeros(B, d)
        else:
            wkv_s, sh_tm, sh_cm = state

        # ── time mix (WKV) ────────────────────────────────────────────────
        h = self.ln1(x)
        hp = torch.cat([sh_tm.unsqueeze(1).to(h.dtype), h[:, :-1]], dim=1)   # x_{t-1}
        xk = h * self.mu_k + hp * (1 - self.mu_k)
        xv = h * self.mu_v + hp * (1 - self.mu_v)
        xr = h * self.mu_r + hp * (1 - self.mu_r)
        k = self.proj_k(xk)
        v = self.proj_v(xv)
        r = torch.sigmoid(self.proj_r(xr))

        W = self.W
        gain = torch.sqrt(1.0 - W * W)                 # unit stationary variance per channel
        a = k.float() * v.float() * gain
        y = decay_scan(a, W, wkv_s)                    # (B, T, d) fp32
        new_wkv = y[:, -1].detach()
        yn = self.ln_x(y.reshape(B * T, d)).reshape(B, T, d)
        x = x + self.proj_out(r * yn.to(r.dtype))

        # ── channel mix (FFN) ────────────────────────────────────────────
        h2 = self.ln2(x)
        h2p = torch.cat([sh_cm.unsqueeze(1).to(h2.dtype), h2[:, :-1]], dim=1)
        xc = h2 * self.mu_c + h2p * (1 - self.mu_c)
        x = x + self.ffn(xc)

        new_state = (new_wkv, h[:, -1].detach(), h2[:, -1].detach())
        return x, new_state


# ════════════════════════════════════════════════════════════════════════════
# HF wrapper
# ════════════════════════════════════════════════════════════════════════════
class FWKVPreTrainedModel(PreTrainedModel):
    config_class = FWKVConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if getattr(module, "_is_out_proj", False):
                std = 0.02 / math.sqrt(2 * self.config.n_layers)
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)
        elif isinstance(module, FactorizedTiedHead):
            pass


class FWKVLanguageModel(FWKVPreTrainedModel, GenerationMixin):
    # This model does NOT use the standard KV-Cache machinery (DynamicCache/
    # StaticCache). Without this, newer transformers versions (4.45+/5.x)
    # auto-wrap past_key_values in a Cache object during generate(), which this
    # custom recurrent-state architecture can't consume.
    _supports_cache_class = False
    _supports_static_cache = False
    _supports_quantized_cache = False

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.amp_dtype = None    # optionally set by a training/inference driver
        self.shared = FactorizedTiedHead(config.vocab_size, config.d_model, config.d_emb)
        self.ln0 = nn.LayerNorm(config.d_model)
        self.blocks = nn.ModuleList([
            FWKVBlock(config.d_model, i, config.n_layers, config.ffn_mult,
                      config.wkv_floor, config.wkv_ceil, config.head_size,
                      config.tau_min, config.tau_max)
            for i in range(config.n_layers)
        ])
        self.norm = nn.LayerNorm(config.d_model)
        self.post_init()

    def get_input_embeddings(self):
        return self.shared

    def set_input_embeddings(self, value):
        self.shared = value

    def resize_vocab(self, new_vocab_size):
        self.shared.resize_vocab(new_vocab_size)
        self.config.vocab_size = new_vocab_size

    def count_params(self):
        return sum(p.numel() for p in self.parameters())

    def forward(
        self,
        input_ids=None,
        past_key_values=None,
        attention_mask=None,
        labels=None,
        use_cache=True,
        **kwargs
    ):
        if input_ids is None:
            raise ValueError("input_ids must be provided")
        if not isinstance(input_ids, torch.Tensor):
            input_ids = torch.tensor(input_ids)
        if input_ids.dtype != torch.long:
            input_ids = input_ids.long()
        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)

        x = self.ln0(self.shared.embed(input_ids))

        # If transformers handed us a Cache object (DynamicCache, etc.) instead
        # of our custom list/tuple of per-layer state tensors, treat it as
        # "no state yet" rather than trying to iterate/index into it.
        if past_key_values is None or isinstance(past_key_values, Cache):
            states_in = [None] * len(self.blocks)
        else:
            states_in = past_key_values

        states_out = []
        for block, state in zip(self.blocks, states_in):
            x, new_state = block(x, state)
            states_out.append(new_state)

        x = self.norm(x)
        x_emb = self.shared.to_emb_space(x)

        loss = None
        logits = None
        if labels is not None:
            if not isinstance(labels, torch.Tensor):
                labels = torch.tensor(labels)
            if labels.dtype != torch.long:
                labels = labels.long()
            loss = chunked_cross_entropy(
                x_emb, self.shared.weight, labels,
                chunk_size=self.config.ce_chunk, ignore_index=-100
            )
        else:
            logits = self.shared.logits(x_emb)

        past_key_values_out = tuple(states_out) if use_cache else None
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values_out,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:] if input_ids.dim() > 1 else input_ids[-1:]

        # Don't coerce a Cache object into a tuple — pass it through as-is,
        # forward() knows how to treat it as "empty" state.
        if past_key_values is not None and not isinstance(past_key_values, (tuple, Cache)):
            past_key_values = tuple(past_key_values)

        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "use_cache": True,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        if past_key_values is None or isinstance(past_key_values, Cache):
            return past_key_values
        reordered = []
        for state in past_key_values:
            if state is None:
                reordered.append(None)
                continue
            # each per-layer state is itself a 3-tuple (wkv_state, sh_tm, sh_cm)
            reordered.append(tuple(s.index_select(0, beam_idx) for s in state))
        return tuple(reordered)


# Register the model with AutoModel
AutoConfig.register("fwkv", FWKVConfig)
AutoModel.register(FWKVConfig, FWKVLanguageModel)