"""Ethoric: compact decoder-only transformer with RMSNorm, RoPE, GQA, SwiGLU,
depthwise Canon convolutions and Peri-LN residuals.

Each block normalises on the way into a sublayer and again on the way out,
which is what Peri-LN means here:

    h = RMSNorm_attn_in(x);  h = h + CanonA(h)
    x = x + RMSNorm_attn_out(Attention(h))
    h = RMSNorm_mlp_in(x);   h = h + CanonC(h)
    x = x + RMSNorm_mlp_out(SwiGLU(h))

Module and parameter names follow the usual Transformers layout. Works on
Transformers 4.x and 5.x; the v5 differences are marked inline.
"""
from __future__ import annotations

import math
import warnings
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

try:  # exported as a sibling file; importable either way
    from .attention import RMSNorm, EthoricAttention, build_flex_block_mask
    from .configuration_ethoric import EthoricConfig
except ImportError:  # pragma: no cover - direct script use
    from attention import RMSNorm, EthoricAttention, build_flex_block_mask
    from configuration_ethoric import EthoricConfig


# ----------------------------------------------------------------------------- canon

class Canon(nn.Module):
    """Causal depthwise conv1d, kernel k, no bias, used as a residual branch.

    Left-padded by k-1 so the output at t depends only on positions <= t.

    With `doc_ids` the convolution is also masked to the current document, so
    taps reaching back past an EOS contribute nothing. Leaving that off looks
    harmless -- only k-1 positions of the next document are touched -- but
    attention then spreads those positions across the rest of that document, so
    the whole of it is affected and `doc_masking` buys you nothing.
    """

    def __init__(self, d_model: int, kernel: int):
        super().__init__()
        self.kernel = int(kernel)
        self.weight = nn.Parameter(torch.empty(d_model, 1, self.kernel))
        nn.init.normal_(self.weight, mean=0.0, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        state: Optional[torch.Tensor] = None,
        doc_ids: Optional[torch.Tensor] = None,
        state_docs: Optional[torch.Tensor] = None,
    ):
        """x is [B, T, d]. `state` carries the previous k-1 inputs for decoding.

        Returns (out, new_state, new_state_docs) so a caller stepping one token
        at a time can thread the convolution window through the cache.
        """
        k = self.kernel
        xt = x.transpose(1, 2)                       # [B, d, T]
        xt = F.pad(xt, (k - 1, 0)) if state is None else torch.cat([state, xt], dim=2)
        new_state = xt[:, :, -(k - 1):] if k > 1 else xt[:, :, :0]

        if doc_ids is None:
            out = F.conv1d(xt, self.weight.to(xt.dtype), groups=xt.shape[1])
            return out.transpose(1, 2), new_state, None

        # -1 never matches a real document id, so the left padding masks out
        # exactly as zero-padding would.
        if state_docs is None:
            dpad = F.pad(doc_ids, (k - 1, 0), value=-1)
        else:
            dpad = torch.cat([state_docs, doc_ids], dim=1)
        new_state_docs = dpad[:, -(k - 1):] if k > 1 else dpad[:, :0]

        T = x.shape[1]
        w = self.weight.to(xt.dtype)
        out = torch.zeros_like(xt[:, :, :T])
        for i in range(k):
            seg = xt[:, :, i : i + T]
            keep = (dpad[:, i : i + T] == doc_ids).unsqueeze(1).to(seg.dtype)
            out = out + w[:, 0, i].view(1, -1, 1) * seg * keep
        return out.transpose(1, 2), new_state, new_state_docs


# ----------------------------------------------------------------------------- blocks

class SwiGLU(nn.Module):
    def __init__(self, d_model: int, hidden: int):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, hidden, bias=False)
        self.up_proj = nn.Linear(d_model, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class EthoricBlock(nn.Module):
    """One Peri-LN block: norm in, Canon, sublayer, norm out, residual add.

    The out-norms are the part that differs from a conventional pre-norm block.
    They sit inside the residual branch, so the residual stream itself is never
    normalised and the identity path stays exact.
    """

    def __init__(self, config: EthoricConfig, layer_idx: int):
        super().__init__()
        d, eps = config.d_model, config.norm_eps
        self.attn_in_norm = RMSNorm(d, eps)
        self.attn_out_norm = RMSNorm(d, eps)
        self.mlp_in_norm = RMSNorm(d, eps)
        self.mlp_out_norm = RMSNorm(d, eps)
        self.attn = EthoricAttention(config, layer_idx)
        self.mlp = SwiGLU(d, config.ffn_hidden)

        canon = config.canon or {}
        self.canon_enabled = bool(canon.get("enabled", False))
        canon_positions = list(canon.get("positions", []))
        k = int(canon.get("kernel", 4))
        self.canon_doc_masking = bool(canon.get("doc_masking", True))
        self.canon_a = Canon(d, k) if self.canon_enabled and "A" in canon_positions else None
        self.canon_c = Canon(d, k) if self.canon_enabled and "C" in canon_positions else None

    def _canon(self, module, name, h, doc_ids, layer_cache):
        docs = doc_ids if self.canon_doc_masking else None
        state = layer_cache.get(name) if layer_cache is not None else None
        state_docs = layer_cache.get(f"{name}_docs") if layer_cache is not None else None
        delta, new_state, new_docs = module(h, state, docs, state_docs)
        if layer_cache is not None:
            layer_cache[name] = new_state
            layer_cache[f"{name}_docs"] = new_docs
        return h + delta

    def forward(self, x, *, positions, doc_ids, layer_cache=None, block_mask=None):
        h = self.attn_in_norm(x)
        if self.canon_a is not None:
            h = self._canon(self.canon_a, "canon_a", h, doc_ids, layer_cache)
        a = self.attn(
            h,
            positions=positions,
            doc_ids=doc_ids,
            layer_cache=layer_cache,
            block_mask=block_mask,
        )
        x = x + self.attn_out_norm(a)

        h = self.mlp_in_norm(x)
        if self.canon_c is not None:
            h = self._canon(self.canon_c, "canon_c", h, doc_ids, layer_cache)
        return x + self.mlp_out_norm(self.mlp(h))


# ----------------------------------------------------------------------------- cache

class EthoricCache(list):
    """One dict per layer: extended keys, values, positions, Canon conv state.

    A plain list of dicts rather than a transformers Cache subclass, so the file
    stays self-contained across transformers versions.
    """

    @classmethod
    def empty(cls, n_layer: int) -> "EthoricCache":
        return cls({} for _ in range(n_layer))

    @property
    def seen_tokens(self) -> int:
        pos = self[0].get("positions") if len(self) else None
        return 0 if pos is None else int(pos.shape[1])


# ----------------------------------------------------------------------------- model

class EthoricPreTrainedModel(PreTrainedModel):
    config_class = EthoricConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["EthoricBlock"]

    def _init_weights(self, module):
        std = 0.02
        if isinstance(module, nn.Linear):
            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.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)


class EthoricModel(EthoricPreTrainedModel):
    def __init__(self, config: EthoricConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
        self.layers = nn.ModuleList(
            EthoricBlock(config, i) for i in range(config.n_layer)
        )
        self.final_norm = RMSNorm(config.d_model, config.norm_eps)
        self.post_init()

    def forward(
        self,
        input_ids: torch.LongTensor,
        *,
        doc_ids: Optional[torch.Tensor] = None,
        positions: Optional[torch.Tensor] = None,
        past_key_values: Optional[EthoricCache] = None,
        use_cache: bool = False,
    ):
        B, T = input_ids.shape
        device = input_ids.device

        if positions is None:
            offset = past_key_values.seen_tokens if past_key_values is not None else 0
            positions = torch.arange(offset, offset + T, device=device).expand(B, T)

        if doc_ids is None and self.config.doc_masking:
            doc_ids = document_ids(input_ids, eos_id=int(self.config.eos_token_id or 0))
            if past_key_values is not None and past_key_values.seen_tokens:
                prev = past_key_values[0].get("doc_ids")
                if prev is not None:
                    doc_ids = doc_ids + prev[:, -1:]
        if not self.config.doc_masking:
            doc_ids = None

        block_mask = None
        backend = getattr(self.layers[0].attn, "backend", "dense")
        if backend == "flex":
            k_pos = positions
            k_doc = doc_ids
            if past_key_values is not None and past_key_values.seen_tokens:
                k_pos = torch.cat([past_key_values[0]["positions"], positions], dim=1)
                if doc_ids is not None and past_key_values[0].get("doc_ids") is not None:
                    k_doc = torch.cat([past_key_values[0]["doc_ids"], doc_ids], dim=1)
            # Two masks: windowed layers additionally restrict to the local window.
            block_mask = {
                None: build_flex_block_mask(positions, k_pos, doc_ids, k_doc, None, self.config.n_head),
            }
            for layer in self.layers:
                w = layer.attn.attn_window
                if w is not None and w not in block_mask:
                    block_mask[w] = build_flex_block_mask(
                        positions, k_pos, doc_ids, k_doc, w, self.config.n_head
                    )

        x = self.embed_tokens(input_ids)
        for i, layer in enumerate(self.layers):
            cache = past_key_values[i] if (use_cache and past_key_values is not None) else None
            bm = None
            if block_mask is not None:
                w = layer.attn.attn_window
                bm = block_mask[w]
            x = layer(
                x,
                positions=positions,
                doc_ids=doc_ids,
                layer_cache=cache,
                block_mask=bm,
            )
        return self.final_norm(x)


# transformers 4.x expects a list of tied names; 5.x expects a
# {target: source} mapping. Support both so the export loads either way.
_TRANSFORMERS_MAJOR = int(transformers.__version__.split(".")[0])
_TIED_WEIGHTS_KEYS = (
    {"lm_head.weight": "model.embed_tokens.weight"}
    if _TRANSFORMERS_MAJOR >= 5
    else ["lm_head.weight"]
)


class EthoricForCausalLM(EthoricPreTrainedModel):
    _tied_weights_keys = _TIED_WEIGHTS_KEYS

    def __init__(self, config: EthoricConfig):
        super().__init__(config)
        self.model = EthoricModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.post_init()
        self._apply_residual_init()

    def _apply_residual_init(self):
        """Residual-output projections get scaled down by depth so the residual
        stream does not blow up at init (the GPT-2 trick).

        Runs after post_init(), so it overwrites what _init_weights wrote. Under
        from_pretrained the loaded weights land afterwards and win, which is why
        a trained checkpoint is unaffected by anything in here.
        """
        std = 0.02 / math.sqrt(2 * self.config.n_layer)
        for layer in self.model.layers:
            nn.init.normal_(layer.attn.o_proj.weight, mean=0.0, std=std)
            nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=std)
            if layer.attn.gate_proj is not None:
                # Zero weight, bias 2.0: the gate starts near-open at
                # sigmoid(2.0) ~= 0.88 and input-independent, so early training
                # sees plain attention and learns the gating later.
                nn.init.zeros_(layer.attn.gate_proj.weight)
                nn.init.constant_(layer.attn.gate_proj.bias, 2.0)

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head


    def set_attention_backend(self, backend: str):
        """'flex' or 'dense'. Dense is the CPU/MPS path and the reference."""
        if backend not in ("flex", "dense"):
            raise ValueError("backend must be 'flex' or 'dense'")
        for layer in self.model.layers:
            layer.attn.backend = backend
        return self

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask=None,
        labels: Optional[torch.LongTensor] = None,
        *,
        inputs_embeds=None,
        doc_ids=None,
        positions=None,
        past_key_values: Optional[EthoricCache] = None,
        use_cache: bool = False,
        z_loss: float = 0.0,
        return_dict: bool = True,
        **kwargs,
    ):
        if inputs_embeds is not None:
            raise ValueError("Ethoric requires input_ids; inputs_embeds is not supported")
        if input_ids is None:
            raise ValueError("Ethoric requires input_ids")

        if use_cache and past_key_values is None:
            past_key_values = EthoricCache.empty(self.config.n_layer)

        hidden = self.model(
            input_ids,
            doc_ids=doc_ids,
            positions=positions,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )
        logits = self.lm_head(hidden)

        loss = None
        if labels is not None:
            loss = self.loss_with_z(logits, labels, z_loss=z_loss)

        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )

    @staticmethod
    def loss_with_z(logits, labels, z_loss: float = 0.0):
        """Cross-entropy (labels already shifted by the loader) plus z-loss.

        Positions labelled -100 are ignored by both terms.
        """
        flat = logits.float().view(-1, logits.shape[-1])
        tgt = labels.reshape(-1)
        loss = F.cross_entropy(flat, tgt, ignore_index=-100)
        if z_loss:
            keep = tgt != -100
            if keep.any():
                lse = torch.logsumexp(flat[keep], dim=-1)
                loss = loss + z_loss * lse.pow(2).mean()
        return loss

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.LongTensor,
        max_new_tokens: int = 32,
        *,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        use_cache: bool = True,
        eos_token_id: Optional[int] = None,
        **kwargs,
    ):
        """Greedy / top-k sampling. Defined here rather than inherited so the
        exported model does not depend on a particular transformers Cache API."""
        self.eval()
        cache = EthoricCache.empty(self.config.n_layer) if use_cache else None
        ids = input_ids
        cur = input_ids
        for _ in range(int(max_new_tokens)):
            out = self.forward(
                cur if use_cache else ids,
                past_key_values=cache,
                use_cache=use_cache,
                return_dict=True,
            )
            logits = out.logits[:, -1, :].float()
            if do_sample:
                logits = logits / max(float(temperature), 1e-6)
                if top_k:
                    kth = logits.topk(int(top_k), dim=-1).values[:, -1:]
                    logits = logits.masked_fill(logits < kth, float("-inf"))
                nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
            else:
                nxt = logits.argmax(dim=-1, keepdim=True)
            ids = torch.cat([ids, nxt], dim=1)
            cur = nxt
            if eos_token_id is not None and bool((nxt == eos_token_id).all()):
                break
        return ids


# ----------------------------------------------------------------------------- helpers

def document_ids(input_ids: torch.Tensor, eos_id: int = 0) -> torch.Tensor:
    """Document id per position: EOS *ends* a document, so the id increments
    after the EOS token. The EOS itself belongs to the document it closes.

    Getting that boundary backwards puts each EOS in the following document,
    which quietly changes what every mask built from these ids allows.
    """
    is_eos = (input_ids == eos_id).to(torch.int32)
    return torch.cumsum(is_eos, dim=1) - is_eos


__all__ = [
    "EthoricConfig",
    "EthoricForCausalLM",
    "EthoricModel",
    "EthoricCache",
    "document_ids",
]