"""ECOACO — the rurtech.ai Mixture-of-Experts causal language model.

ECOACO (Sanskrit: wealth, prosperity, purpose) is rurtech.ai's vertical
Mixture-of-Experts LLM for banking.

This is a real, from-scratch sparse-MoE transformer (not a wrapper around
another model). Each transformer block replaces the dense feed-forward network
with a set of expert FFNs and a top-k router, so only a fraction of parameters
activate per token — the defining property of a Mixture-of-Experts model.

Design goals: small enough to train and run CPU-only in this environment, yet
architecturally faithful so the same code scales up by changing the config.

Key pieces:
  * RMSNorm + rotary position embeddings (RoPE)
  * Multi-head self-attention with a causal mask
  * Top-k token routing over N experts, with a load-balancing aux loss
  * Weight tying between the embedding and the LM head
"""

from __future__ import annotations

import math
from dataclasses import asdict, dataclass

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


@dataclass
class EcoacoConfig:
    version: str = "1.0"
    vocab_size: int = 4096
    dim: int = 256
    n_layers: int = 6
    n_heads: int = 8
    n_experts: int = 8
    n_experts_per_token: int = 2  # top-k routing
    ffn_hidden: int = 512
    max_seq_len: int = 512
    rope_theta: float = 10000.0
    aux_loss_coef: float = 0.01
    tie_embeddings: bool = True

    def to_dict(self) -> dict:
        d = asdict(self)
        d["model_type"] = "ecoaco"
        d["name"] = "ECOACO"
        d["architectures"] = ["EcoacoForCausalLM"]
        return d


def _rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    return emb.cos().to(dtype), emb.sin().to(dtype)


def _rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def _apply_rope(q, k, cos, sin):
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    q = (q * cos) + (_rotate_half(q) * sin)
    k = (k * cos) + (_rotate_half(k) * sin)
    return q, k


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight


class Attention(nn.Module):
    def __init__(self, cfg: EcoacoConfig):
        super().__init__()
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.dim // cfg.n_heads
        self.wq = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wk = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wv = nn.Linear(cfg.dim, cfg.dim, bias=False)
        self.wo = nn.Linear(cfg.dim, cfg.dim, bias=False)

    def forward(self, x, cos, sin):
        b, t, _ = x.shape
        q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        q, k = _apply_rope(q, k, cos, sin)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(b, t, -1)
        return self.wo(out)


class Expert(nn.Module):
    """A single SwiGLU feed-forward expert."""

    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden, bias=False)
        self.w2 = nn.Linear(hidden, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class MoELayer(nn.Module):
    """Top-k routed mixture of experts with a load-balancing auxiliary loss."""

    def __init__(self, cfg: EcoacoConfig):
        super().__init__()
        self.n_experts = cfg.n_experts
        self.top_k = cfg.n_experts_per_token
        self.gate = nn.Linear(cfg.dim, cfg.n_experts, bias=False)
        self.experts = nn.ModuleList(
            [Expert(cfg.dim, cfg.ffn_hidden) for _ in range(cfg.n_experts)]
        )
        self.aux_loss = torch.tensor(0.0)

    def forward(self, x):
        b, t, d = x.shape
        x_flat = x.reshape(-1, d)  # (tokens, dim)
        logits = self.gate(x_flat)  # (tokens, n_experts)
        probs = F.softmax(logits, dim=-1)

        topk_probs, topk_idx = probs.topk(self.top_k, dim=-1)
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)

        out = torch.zeros_like(x_flat)
        for slot in range(self.top_k):
            expert_ids = topk_idx[:, slot]
            weight = topk_probs[:, slot].unsqueeze(-1)
            for e in range(self.n_experts):
                mask = expert_ids == e
                if mask.any():
                    out[mask] += weight[mask] * self.experts[e](x_flat[mask])

        # Switch-Transformer load-balancing loss: encourage uniform expert use.
        importance = probs.mean(dim=0)  # fraction of routing mass per expert
        load = torch.zeros(self.n_experts, device=x.device)
        load.scatter_add_(0, topk_idx.reshape(-1), torch.ones_like(topk_idx.reshape(-1), dtype=load.dtype))
        load = load / load.sum().clamp(min=1)
        self.aux_loss = (importance * load).sum() * self.n_experts

        return out.view(b, t, d)


class Block(nn.Module):
    def __init__(self, cfg: EcoacoConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.dim)
        self.attn = Attention(cfg)
        self.moe_norm = RMSNorm(cfg.dim)
        self.moe = MoELayer(cfg)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.attn_norm(x), cos, sin)
        x = x + self.moe(self.moe_norm(x))
        return x


class EcoacoForCausalLM(nn.Module):
    def __init__(self, cfg: EcoacoConfig):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.dim)
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight
        self.apply(self._init)

    def _init(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, labels=None):
        b, t = input_ids.shape
        x = self.embed(input_ids)
        cos, sin = _rope_cache(t, self.cfg.dim // self.cfg.n_heads, self.cfg.rope_theta,
                               x.device, x.dtype)
        aux = torch.tensor(0.0, device=x.device)
        for block in self.blocks:
            x = block(x, cos, sin)
            aux = aux + block.moe.aux_loss
        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            ce = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            loss = ce + self.cfg.aux_loss_coef * (aux / self.cfg.n_layers)
        return {"logits": logits, "loss": loss, "aux_loss": aux}

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=64, temperature=0.8, top_k=40, eos_id=None):
        self.eval()
        for _ in range(max_new_tokens):
            ids = input_ids[:, -self.cfg.max_seq_len:]
            logits = self(ids)["logits"][:, -1, :]
            if temperature > 0:
                logits = logits / temperature
                if top_k:
                    v, _ = logits.topk(min(top_k, logits.size(-1)))
                    logits[logits < v[:, [-1]]] = -float("inf")
                probs = F.softmax(logits, dim=-1)
                nxt = torch.multinomial(probs, 1)
            else:
                nxt = logits.argmax(-1, keepdim=True)
            input_ids = torch.cat([input_ids, nxt], dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return input_ids

    def num_params(self) -> tuple[int, int]:
        total = sum(p.numel() for p in self.parameters())
        # Active params per token: everything except the non-selected experts.
        per_expert = sum(p.numel() for p in self.blocks[0].moe.experts[0].parameters())
        inactive = per_expert * (self.cfg.n_experts - self.cfg.n_experts_per_token) * self.cfg.n_layers
        return total, total - inactive

    @classmethod
    def from_pretrained(cls, model_id_or_path: str):
        """Load ECOACO from a local directory or a Hugging Face repo id.

            from modeling_ecoaco import EcoacoForCausalLM
            model, tok = EcoacoForCausalLM.from_pretrained("rurtech-ai/ECOACO")

        A local path is used as-is; anything else is fetched from the HF Hub
        (requires `pip install huggingface_hub`).
        """
        import json
        from pathlib import Path

        from safetensors.torch import load_model

        path = Path(model_id_or_path)
        if not path.exists():
            from huggingface_hub import snapshot_download

            path = Path(snapshot_download(model_id_or_path))

        cfg_d = json.loads((path / "config.json").read_text())
        for k in ("model_type", "architectures", "name"):
            cfg_d.pop(k, None)
        model = cls(EcoacoConfig(**cfg_d))
        load_model(model, str(path / "model.safetensors"))
        model.eval()

        # Load the paired tokenizer if present, and return it for convenience.
        tok = None
        tok_path = path / "tokenizer.json"
        if tok_path.exists():
            try:
                import sys

                sys.path.insert(0, str(path))
                from tokenizer import ByteBPETokenizer

                tok = ByteBPETokenizer.load(tok_path)
            except Exception:
                tok = None
        return model, tok

    @torch.no_grad()
    def chat(self, tokenizer, prompt: str, max_new_tokens: int = 80, temperature: float = 0.7) -> str:
        """One-line text generation given the paired tokenizer."""
        ids = torch.tensor([tokenizer.encode(prompt, add_bos=True)])
        out = self.generate(ids, max_new_tokens=max_new_tokens,
                            temperature=temperature, eos_id=tokenizer.eos_id)
        return tokenizer.decode(out[0].tolist())