"""TS Encoder (T2.1): channel-independent patch embedding + lightweight Transformer.

Pipeline: ``[B, T, C]`` → Patchify → shared per-channel patch linear → mean-pool
over channels → +learnable positional embedding → 2-layer Transformer →
``[B, n_patches, d]``.

Design ref: ``docs/superpowers/specs/2026-06-29-ts-as-modality-design.md`` §2.

Spec clarifications (see ``tests/test_ts_encoder.py`` docstring):
- ``n_patches = (T - patch_len) // stride + 1``  (T=512,P=8,S=4 ⇒ 127).
- "Channel-independent" ⇒ shared per-channel patch embedding + mean-pool over
  channels, so the output token sequence is C-free. Channel identity is supplied
  separately by the ``[属性]`` text tokens in the multimodal splice (T2.3).
- Actual params ≈1.6M for the default compact config (design's "≈4M" was rough).
"""
from __future__ import annotations

import torch
from torch import nn


class Patchify(nn.Module):
    """Channel-independent patching along the time dimension.

    Input  : ``[B, T, C]``
    Output : ``[B, n_patches, C, patch_len]`` where
             ``n_patches = (T - patch_len) // stride + 1``.
    """

    def __init__(self, patch_len: int = 8, stride: int = 4) -> None:
        super().__init__()
        if patch_len <= 0 or stride <= 0:
            raise ValueError("patch_len and stride must be positive")
        self.patch_len = patch_len
        self.stride = stride

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, T, C] → unfold time into patches: [B, n_patches, C, patch_len]
        return x.unfold(1, self.patch_len, self.stride)

    def n_patches(self, T: int) -> int:
        if T < self.patch_len:
            raise ValueError(f"T={T} smaller than patch_len={self.patch_len}")
        return (T - self.patch_len) // self.stride + 1


class TSEncoder(nn.Module):
    """TS Encoder: Patchify → shared patch linear (per channel) → channel pool
    → positional embedding → 2-layer Transformer.

    Output is C-free: ``[B, n_patches, d]``.
    """

    def __init__(
        self,
        d: int = 256,
        layers: int = 2,
        heads: int = 4,
        patch_len: int = 8,
        stride: int = 4,
        ffn_mult: int = 4,
        dropout: float = 0.1,
        max_patches: int = 2048,
    ) -> None:
        super().__init__()
        if d % heads != 0:
            raise ValueError(f"d ({d}) must be divisible by heads ({heads})")
        self.d = d
        self.patch_len = patch_len
        self.stride = stride
        self.patchify = Patchify(patch_len=patch_len, stride=stride)

        # Shared per-channel patch embedding. Applied to each channel's patch
        # independently, then we mean-pool over channels → C-free tokens.
        self.patch_embed = nn.Linear(patch_len, d)

        # Learnable positional embedding, large enough for any foreseeable T.
        self.pos_embed = nn.Parameter(torch.zeros(max_patches, d))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d,
            nhead=heads,
            dim_feedforward=d * ffn_mult,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=layers)
        self.ln = nn.LayerNorm(d)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, T, C]
        B, T, C = x.shape
        patches = self.patchify(x)          # [B, n_patches, C, patch_len]
        n_patches = patches.shape[1]

        h = self.patch_embed(patches)       # [B, n_patches, C, d]
        h = h.mean(dim=2)                   # [B, n_patches, d]  (channel pool)
        h = h + self.pos_embed[:n_patches]  # add positional embedding
        h = self.transformer(h)             # [B, n_patches, d]
        return self.ln(h)