from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path

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


_compiled_block = None


def _rms(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    y = x.float()
    y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6)
    return y.to(x.dtype) * weight


def _block_forward(
    x: torch.Tensor,
    n1: torch.Tensor,
    qkv: torch.Tensor,
    qn: torch.Tensor,
    kn: torch.Tensor,
    out: torch.Tensor,
    n2: torch.Tensor,
    gate: torch.Tensor,
    up: torch.Tensor,
    down: torch.Tensor,
    rope_cos: torch.Tensor,
    rope_sin: torch.Tensor,
    heads: int,
) -> torch.Tensor:
    batch, length, dim = x.shape
    head_dim = dim // heads
    a = _rms(x, n1)
    q, k, v = F.linear(a, qkv).chunk(3, dim=-1)
    q = _rms(q.view(batch, length, heads, head_dim), qn).transpose(1, 2)
    k = _rms(k.view(batch, length, heads, head_dim), kn).transpose(1, 2)
    v = v.view(batch, length, heads, head_dim).transpose(1, 2)
    cos = rope_cos[:, :, :length].to(x.dtype)
    sin = rope_sin[:, :, :length].to(x.dtype)
    qa, qb = q[..., 0::2], q[..., 1::2]
    ka, kb = k[..., 0::2], k[..., 1::2]
    q = torch.stack((qa * cos - qb * sin, qa * sin + qb * cos), dim=-1).flatten(-2)
    k = torch.stack((ka * cos - kb * sin, ka * sin + kb * cos), dim=-1).flatten(-2)
    a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    x = x + F.linear(a.transpose(1, 2).reshape(batch, length, dim), out)
    a = _rms(x, n2)
    return x + F.linear(F.silu(F.linear(a, gate)) * F.linear(a, up), down)


@dataclass
class CascadeConfig:
    vocab_size: int = 256
    sequence_length: int = 1024
    local_dimension: int = 384
    local_layers: int = 4
    local_heads: int = 6
    local_intermediate_size: int = 1536
    dimension: int = 1024
    heads: int = 16
    intermediate_size: int = 12032
    route_paths: int = 3
    active_layers: int = 8
    max_patches: int = 192
    patch_rate: float = 0.125
    entropy_threshold: float = 5.0
    rope_theta: float = 10000.0
    patch_aux_weight: float = 0.25
    patch_mode: str = "byte_transition"
    tokenizer_name: str | None = None
    model_type: str = "cascade_byte_lm"
    architectures: tuple[str, ...] = ("CascadeForCausalLM",)

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(asdict(self), indent=2), encoding="utf-8")

    @classmethod
    def load(cls, path: str | Path) -> "CascadeConfig":
        values = json.loads(Path(path).read_text(encoding="utf-8"))
        values.pop("model_type", None)
        values.pop("architectures", None)
        return cls(**values)


class RMSNorm(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = x.float()
        y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6)
        return y.to(x.dtype) * self.weight


class Attention(nn.Module):
    def __init__(self, dim: int, heads: int, max_length: int, theta: float) -> None:
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.out = nn.Linear(dim, dim, bias=False)
        pos = torch.arange(max_length, dtype=torch.float32)
        inv = torch.exp(-math.log(theta) * torch.arange(0, self.head_dim, 2) / self.head_dim)
        angle = pos[:, None] * inv[None, :]
        self.register_buffer("rope_cos", angle.cos()[None, None], persistent=False)
        self.register_buffer("rope_sin", angle.sin()[None, None], persistent=False)

    def rotate(self, x: torch.Tensor) -> torch.Tensor:
        length = x.shape[-2]
        cos = self.rope_cos[:, :, :length].to(x.dtype)
        sin = self.rope_sin[:, :, :length].to(x.dtype)
        a, b = x[..., 0::2], x[..., 1::2]
        return torch.stack((a * cos - b * sin, a * sin + b * cos), dim=-1).flatten(-2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = self.q_norm(q.view(batch, length, self.heads, self.head_dim)).transpose(1, 2)
        k = self.k_norm(k.view(batch, length, self.heads, self.head_dim)).transpose(1, 2)
        v = v.view(batch, length, self.heads, self.head_dim).transpose(1, 2)
        q, k = self.rotate(q), self.rotate(k)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(y.transpose(1, 2).reshape(batch, length, -1))


class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden: int) -> None:
        super().__init__()
        self.gate = nn.Linear(dim, hidden, bias=False)
        self.up = nn.Linear(dim, hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, dim: int, heads: int, hidden: int, max_length: int, theta: float) -> None:
        super().__init__()
        self.n1 = RMSNorm(dim)
        self.attn = Attention(dim, heads, max_length, theta)
        self.n2 = RMSNorm(dim)
        self.ffn = FeedForward(dim, hidden)
        self.compiled = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.compiled:
            return _compiled_block(
                x,
                self.n1.weight,
                self.attn.qkv.weight,
                self.attn.q_norm.weight,
                self.attn.k_norm.weight,
                self.attn.out.weight,
                self.n2.weight,
                self.ffn.gate.weight,
                self.ffn.up.weight,
                self.ffn.down.weight,
                self.attn.rope_cos,
                self.attn.rope_sin,
                self.attn.heads,
            )
        x = x + self.attn(self.n1(x))
        return x + self.ffn(self.n2(x))


class CascadeForCausalLM(nn.Module):
    def __init__(self, config: CascadeConfig, surprise: torch.Tensor | None = None) -> None:
        super().__init__()
        self.config = config
        if surprise is None:
            shape = (config.vocab_size,) if config.patch_mode == "token_surprise" else (config.vocab_size, config.vocab_size)
            surprise = torch.full(shape, config.entropy_threshold)
        self.register_buffer("surprise", surprise.float())
        self.register_buffer("entropy_threshold", torch.tensor(config.entropy_threshold))
        self.embed = nn.Embedding(config.vocab_size, config.local_dimension)
        self.local_blocks = nn.ModuleList(
            Block(
                config.local_dimension,
                config.local_heads,
                config.local_intermediate_size,
                config.sequence_length,
                config.rope_theta,
            )
            for _ in range(config.local_layers)
        )
        self.to_core = nn.Linear(config.local_dimension, config.dimension, bias=False)
        self.paths = nn.ModuleList(
            nn.ModuleList(
                Block(
                    config.dimension,
                    config.heads,
                    config.intermediate_size,
                    config.max_patches,
                    config.rope_theta,
                )
                for _ in range(config.active_layers)
            )
            for _ in range(config.route_paths)
        )
        self.core_norm = RMSNorm(config.dimension)
        self.patch_head = nn.Linear(config.dimension, config.vocab_size, bias=False)
        self.from_core = nn.Linear(config.dimension, config.local_dimension, bias=False)
        self.mix = nn.Linear(2 * config.local_dimension, config.local_dimension, bias=False)
        self.norm = RMSNorm(config.local_dimension)
        self.head = nn.Linear(config.local_dimension, config.vocab_size, bias=False)
        self.head.weight = self.embed.weight
        self.apply(self._initialize)
        residual_std = 0.02 / math.sqrt(2 * (config.local_layers + config.active_layers))
        for module in self.modules():
            if isinstance(module, Attention):
                nn.init.normal_(module.out.weight, std=residual_std)
            elif isinstance(module, FeedForward):
                nn.init.normal_(module.down.weight, std=residual_std)

    @staticmethod
    def _initialize(module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def set_entropy(self, surprise: torch.Tensor, threshold: float) -> None:
        self.surprise.copy_(surprise)
        self.entropy_threshold.fill_(threshold)
        self.config.entropy_threshold = threshold

    def compile_blocks(self) -> None:
        global _compiled_block
        if _compiled_block is None:
            _compiled_block = torch.compile(_block_forward, mode="max-autotune-no-cudagraphs", fullgraph=True)
        for block in self.local_blocks:
            block.compiled = True
        for path in self.paths:
            for block in path:
                block.compiled = True

    def patch(self, x: torch.Tensor, h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        if self.config.patch_mode == "token_surprise":
            score = self.surprise[x[:, 1:]]
        else:
            score = self.surprise[x[:, :-1], x[:, 1:]]
        boundary = torch.cat((torch.ones_like(x[:, :1], dtype=torch.bool), score >= self.entropy_threshold), dim=1)
        raw_ids = boundary.long().cumsum(1).sub(1)
        totals = raw_ids[:, -1:].add(1)
        compacted = raw_ids.mul(self.config.max_patches).div(totals, rounding_mode="floor")
        ids = torch.where(totals.gt(self.config.max_patches), compacted, raw_ids)
        merged_boundary = torch.cat((torch.ones_like(ids[:, :1], dtype=torch.bool), ids[:, 1:].ne(ids[:, :-1])), dim=1)
        end = torch.cat((ids[:, :-1].ne(ids[:, 1:]), torch.ones_like(ids[:, -1:], dtype=torch.bool)), dim=1)
        patches = torch.zeros(x.shape[0], self.config.max_patches, h.shape[-1], device=h.device, dtype=h.dtype)
        patches.scatter_add_(1, ids.unsqueeze(-1).expand_as(h), h * end.unsqueeze(-1))
        counts = torch.zeros(x.shape[0], self.config.max_patches, device=x.device, dtype=torch.int32)
        counts.scatter_add_(1, ids, torch.ones_like(ids, dtype=torch.int32))
        first = torch.zeros(x.shape[0], self.config.max_patches, device=x.device, dtype=torch.long)
        first.scatter_add_(1, ids, x * merged_boundary)
        overflow = totals[:, 0].gt(self.config.max_patches)
        return patches, ids, counts, first, overflow

    def forward(self, x: torch.Tensor, route_id: int) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
        h = self.embed(x)
        for block in self.local_blocks:
            h = block(h)
        patches, ids, counts, first, overflow = self.patch(x, h)
        z = self.to_core(patches)
        for block in self.paths[route_id]:
            z = block(z)
        z = self.core_norm(z)
        previous = ids.sub(1).clamp_min(0)
        context = z.gather(1, previous.unsqueeze(-1).expand(-1, -1, z.shape[-1]))
        context = self.from_core(context) * ids.gt(0).unsqueeze(-1)
        gate = torch.sigmoid(self.mix(torch.cat((h, context), dim=-1)))
        logits = self.head(self.norm(h + gate * context))
        valid = counts[:, 1:].gt(0)
        auxiliary = F.cross_entropy(self.patch_head(z[:, :-1])[valid], first[:, 1:][valid])
        stats = {
            "patches": counts.gt(0).sum(),
            "overflow": overflow.sum(),
            "gate": gate.mean(),
        }
        return logits, auxiliary, stats

    def parameter_counts(self) -> dict[str, int]:
        paths = sum(p.numel() for p in self.paths.parameters())
        active_path = paths // self.config.route_paths
        total = sum(p.numel() for p in self.parameters())
        return {"total": total, "active": total - paths + active_path, "shared": total - paths, "active_path": active_path}