# style_flow_matching_1d.py
# Stabilny i uproszczony model dla Flow Matching przewidujący 3 style
# (pitch, energy, duration), każdy wektorem 1x64, czyli wejście/wyjście (B, 3, 64).
#
# - Architektura: dekoderowy Transformer nad 3 tokenami stylu z Cross-Attn do
#   tokenów kondycjonujących (embedding). Stabilne pre-normy (RMSNorm), SiLU,
#   opcjonalny RoPE, DropPath (stochastic depth), dropout tokenów pamięci.
# - Czas: sinusoidalny embedding czasu + MLP dodawany do tokenów stylu i pamięci.
# - CFG: wbudowane przez FixedEmbedding (maskowanie embeddingu + skala).
# - Strata: MSE lub SmoothL1 (Huber).
# - Zależności: torch, einops.
#
# Oczekiwane kształty:
#   x: (B, 3, 64)             -> style: pitch, energy, duration
#   embedding: (B, L, D)      -> tokeny kondycjonujące
#   time t: (B,) w [0, 1]
# Wyjście sieci: (B, 3, 64)
#
# Dla małego zbioru (~50k) polecane:
#   d_model=128, num_layers=4, num_heads=4, head_features=32, multiplier=3
#   dropout=0.1, attn_dropout=0.1, ff_dropout=0.1
#   mem_token_keep_prob=0.8, drop_path_prob=0.1
#   FlowMatching1D(loss_type="smooth_l1", huber_delta=0.02)

from __future__ import annotations

from math import pi
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from torch import Tensor, einsum


# -------------------------
# Utils
# -------------------------


def exists(x) -> bool:
    return x is not None


def rand_bool(shape, proba: float, device=None) -> Tensor:
    if proba <= 0.0:
        return torch.zeros(shape, dtype=torch.bool, device=device)
    if proba >= 1.0:
        return torch.ones(shape, dtype=torch.bool, device=device)
    return torch.rand(shape, device=device) < proba


def rearrange_many(tensors, pattern: str, **kwargs):
    return tuple(rearrange(t, pattern, **kwargs) for t in tensors)


@torch.no_grad()
def _time_warp(u: Tensor, kind: str = "cos") -> Tensor:
    # u w [0, 1]
    if kind == "linear":
        return u
    if kind == "cos":
        # 0.5 * (1 - cos(pi * u))
        return 0.5 * (1.0 - torch.cos(pi * u))
    raise ValueError(f"Unknown time_scheduler: {kind}")


class DropPath(nn.Module):
    """
    Stochastic depth (per-sample). Zera cały residual branch z prawdopodobieństwem
    drop_prob w trakcie treningu. W ewaluacji: identity.
    """

    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = float(drop_prob)

    def forward(self, x: Tensor) -> Tensor:
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep = 1.0 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask = x.new_empty(shape).bernoulli_(keep)
        return x * mask / keep


# -------------------------
# Norms
# -------------------------


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-8, elementwise_affine: bool = True):
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_buffer("weight", torch.ones(dim))

    def forward(self, x: Tensor) -> Tensor:
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        return x / rms * self.weight


def _make_norm(norm_type: str, dim: int) -> nn.Module:
    if norm_type is None or norm_type == "none":
        return nn.Identity()
    if norm_type == "layer":
        return nn.LayerNorm(dim)
    if norm_type == "rms":
        return RMSNorm(dim)
    raise ValueError(f"Unknown norm_type: {norm_type}")


# -------------------------
# Time embeddings
# -------------------------


class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0, "time embedding dim must be even"
        self.dim = dim

    def forward(self, t: Tensor) -> Tensor:
        # t: (B,)
        half = self.dim // 2
        device = t.device
        exponents = torch.arange(half, device=device, dtype=t.dtype)
        freqs = torch.exp(
            -torch.log(torch.tensor(10000.0, device=device)) * exponents / half
        )
        args = t[:, None] * freqs[None, :]
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)


def TimePositionalEmbedding(out_features: int, time_embed_dim: int = 128) -> nn.Module:
    return nn.Sequential(
        SinusoidalTimeEmbedding(time_embed_dim),
        nn.Linear(time_embed_dim, out_features),
        nn.SiLU(),
        nn.Linear(out_features, out_features),
    )


# -------------------------
# Fixed embedding for CFG
# -------------------------


class FixedEmbedding(nn.Module):
    """
    Learned positional embedding o długości 'max_length' i wymiarze 'features'.
    Używany jako embedding bezwarunkowy (CFG).
    """

    def __init__(self, max_length: int, features: int):
        super().__init__()
        self.max_length = max_length
        self.features = features
        if features > 0:
            self.embedding = nn.Embedding(max_length, features)
        else:
            self.register_buffer("dummy", torch.zeros(1))

    def forward(self, x_like: Tensor) -> Tensor:
        # x_like: (B, L, D) - wykorzystywane B i L
        batch_size, length = x_like.shape[0], x_like.shape[1]
        assert length <= self.max_length, "L must be <= max_length"
        device = x_like.device
        if self.features == 0:
            return x_like.new_zeros(batch_size, length, 0)
        pos = torch.arange(length, device=device)
        fixed = self.embedding(pos)  # (L, D)
        fixed = fixed.unsqueeze(0).expand(batch_size, -1, -1)  # (B, L, D)
        return fixed


# -------------------------
# Rotary embeddings (optional)
# -------------------------


class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 2048, base: int = 10000):
        super().__init__()
        assert dim % 2 == 0, "RoPE head dim must be even"
        self.dim = dim
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def _build_sin_cos(self, n: int, device: torch.device, dtype: torch.dtype):
        positions = torch.arange(n, device=device, dtype=dtype).unsqueeze(1)
        angles = positions * self.inv_freq.to(device=device, dtype=dtype)[None, :]
        sin = torch.sin(angles)
        cos = torch.cos(angles)
        sin = torch.stack([sin, sin], dim=-1).reshape(n, self.dim)
        cos = torch.stack([cos, cos], dim=-1).reshape(n, self.dim)
        return sin, cos

    @staticmethod
    def rotate_half(x: Tensor) -> Tensor:
        x1, x2 = x[..., ::2], x[..., 1::2]
        return torch.stack((-x2, x1), dim=-1).reshape_as(x)

    # Self-attn: Nq == Nk
    def apply_rotary_same(self, q: Tensor, k: Tensor) -> Tuple[Tensor, Tensor]:
        n = q.shape[-2]
        device, dtype = q.device, q.dtype
        sin, cos = self._build_sin_cos(n, device, dtype)
        sin = sin.unsqueeze(0).unsqueeze(0)
        cos = cos.unsqueeze(0).unsqueeze(0)
        q = (q * cos) + (self.rotate_half(q) * sin)
        k = (k * cos) + (self.rotate_half(k) * sin)
        return q, k

    # Cross-attn: Nq może różnić się od Nk
    def apply_rotary_qk(self, q: Tensor, k: Tensor) -> Tuple[Tensor, Tensor]:
        n_q, n_k = q.shape[-2], k.shape[-2]
        # q
        sin_q, cos_q = self._build_sin_cos(n_q, q.device, q.dtype)
        sin_q = sin_q.unsqueeze(0).unsqueeze(0)
        cos_q = cos_q.unsqueeze(0).unsqueeze(0)
        q = (q * cos_q) + (self.rotate_half(q) * sin_q)
        # k
        sin_k, cos_k = self._build_sin_cos(n_k, k.device, k.dtype)
        sin_k = sin_k.unsqueeze(0).unsqueeze(0)
        cos_k = cos_k.unsqueeze(0).unsqueeze(0)
        k = (k * cos_k) + (self.rotate_half(k) * sin_k)
        return q, k


# -------------------------
# Attention primitives
# -------------------------


class FeedForward(nn.Module):
    def __init__(self, features: int, multiplier: int):
        super().__init__()
        mid = features * multiplier
        self.net = nn.Sequential(
            nn.Linear(features, mid),
            nn.SiLU(),
            nn.Linear(mid, features),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)


class SelfAttention(nn.Module):
    def __init__(
        self,
        features: int,
        *,
        head_features: int,
        num_heads: int,
        use_rope: bool = False,
        rope_max_seq_len: int = 512,
        attn_dropout: float = 0.0,
        out_dropout: float = 0.0,
        norm_type: str = "rms",
    ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = head_features**-0.5
        mid = head_features * num_heads

        self.norm = _make_norm(norm_type, features)
        self.to_qkv = nn.Linear(features, mid * 3, bias=False)

        self.to_out = nn.Linear(mid, features)
        self.attn_dropout = (
            nn.Dropout(attn_dropout) if attn_dropout > 0.0 else nn.Identity()
        )
        self.out_dropout = (
            nn.Dropout(out_dropout) if out_dropout > 0.0 else nn.Identity()
        )
        self.rotary = (
            RotaryEmbedding(head_features, max_seq_len=rope_max_seq_len)
            if use_rope
            else None
        )

    def forward(self, x: Tensor) -> Tensor:
        # x: (B, N, F)
        x_n = self.norm(x) if not isinstance(self.norm, nn.Identity) else x
        q, k, v = self.to_qkv(x_n).chunk(3, dim=-1)
        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)
        if self.rotary is not None:
            q, k = self.rotary.apply_rotary_same(q, k)
        sim = einsum("b h n d, b h m d -> b h n m", q, k) * self.scale
        attn = self.attn_dropout(sim.softmax(dim=-1))
        out = einsum("b h n m, b h m d -> b h n d", attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.out_dropout(self.to_out(out))
        return out


class CrossAttention(nn.Module):
    def __init__(
        self,
        features_q: int,
        features_kv: int,
        *,
        head_features: int,
        num_heads: int,
        use_rope: bool = False,
        rope_max_seq_len: int = 512,
        attn_dropout: float = 0.0,
        out_dropout: float = 0.0,
        norm_type: str = "rms",
    ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = head_features**-0.5
        mid = head_features * num_heads

        self.norm_q = _make_norm(norm_type, features_q)
        self.norm_kv = _make_norm(norm_type, features_kv)

        self.to_q = nn.Linear(features_q, mid, bias=False)
        self.to_kv = nn.Linear(features_kv, mid * 2, bias=False)

        self.to_out = nn.Linear(mid, features_q)
        self.attn_dropout = (
            nn.Dropout(attn_dropout) if attn_dropout > 0.0 else nn.Identity()
        )
        self.out_dropout = (
            nn.Dropout(out_dropout) if out_dropout > 0.0 else nn.Identity()
        )
        self.rotary = (
            RotaryEmbedding(head_features, max_seq_len=rope_max_seq_len)
            if use_rope
            else None
        )

    def forward(self, x: Tensor, mem: Tensor) -> Tensor:
        # x: (B, Nq, Fq), mem: (B, Nk, Fkv)
        x_n = self.norm_q(x) if not isinstance(self.norm_q, nn.Identity) else x
        m_n = self.norm_kv(mem) if not isinstance(self.norm_kv, nn.Identity) else mem

        q = self.to_q(x_n)
        k, v = self.to_kv(m_n).chunk(2, dim=-1)

        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)

        if self.rotary is not None:
            # Nq i Nk mogą się różnić -> osobne sin/cos
            q, k = self.rotary.apply_rotary_qk(q, k)

        sim = einsum("b h n d, b h m d -> b h n m", q, k) * self.scale
        attn = self.attn_dropout(sim.softmax(dim=-1))
        out = einsum("b h n m, b h m d -> b h n d", attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.out_dropout(self.to_out(out))
        return out


# -------------------------
# Transformer style blocks
# -------------------------


class StyleBlock(nn.Module):
    """
    Jeden blok:
      x = x + DropPath(SelfAttn(x))
      x = x + DropPath(CrossAttn(x, mem))
      x = x + DropPath(FF(x))
    """

    def __init__(
        self,
        d_model: int,
        d_mem: int,
        *,
        num_heads: int,
        head_features: int,
        multiplier: int,
        use_rope_sa: bool = True,
        use_rope_ca: bool = True,
        rope_max_seq_len: int = 512,
        dropout: float = 0.0,
        attn_dropout: float = 0.0,
        ff_dropout: float = 0.0,
        norm_type: str = "rms",
        drop_path: Optional[nn.Module] = None,
    ):
        super().__init__()
        self.self_attn = SelfAttention(
            features=d_model,
            head_features=head_features,
            num_heads=num_heads,
            use_rope=use_rope_sa,
            rope_max_seq_len=rope_max_seq_len,
            attn_dropout=attn_dropout,
            out_dropout=dropout,
            norm_type=norm_type,
        )
        self.cross_attn = (
            CrossAttention(
                features_q=d_model,
                features_kv=d_mem,
                head_features=head_features,
                num_heads=num_heads,
                use_rope=use_rope_ca,
                rope_max_seq_len=rope_max_seq_len,
                attn_dropout=attn_dropout,
                out_dropout=dropout,
                norm_type=norm_type,
            )
            if d_mem > 0
            else None
        )
        self.ff_norm = _make_norm(norm_type, d_model)
        self.ff = FeedForward(d_model, multiplier)
        self.ff_dropout = nn.Dropout(ff_dropout) if ff_dropout > 0.0 else nn.Identity()
        self.drop_path = drop_path if drop_path is not None else nn.Identity()

    def forward(self, x: Tensor, mem: Optional[Tensor]) -> Tensor:
        x = x + self.drop_path(self.self_attn(x))
        if self.cross_attn is not None and exists(mem) and mem.size(-1) > 0:
            x = x + self.drop_path(self.cross_attn(x, mem))
        ff_in = self.ff_norm(x) if not isinstance(self.ff_norm, nn.Identity) else x
        x = x + self.drop_path(self.ff_dropout(self.ff(ff_in)))
        return x


# -------------------------
# StyleTransformer1d (velocity net)
# -------------------------


class StyleTransformer1d(nn.Module):
    """
    Model do przewidywania 3 wektorów stylu (B, 3, 64) w Flow Matching.
    3 tokeny stylu są dekodowane z pomocą self-attn i cross-attn do tokenów
    kondycjonujących (embedding).
    """

    def __init__(
        self,
        *,
        style_count: int = 3,  # 3: pitch, energy, duration
        channels: int = 64,  # rozmiar wektora stylu
        context_embedding_features: int = 256,  # D dla tokenów kondycjonujących
        d_model: int = 256,
        num_layers: int = 6,
        num_heads: int = 8,
        head_features: int = 32,
        multiplier: int = 4,
        use_rope_sa: bool = True,
        use_rope_ca: bool = True,
        rope_max_seq_len: int = 512,
        time_embed_dim: int = 128,
        dropout: float = 0.0,
        attn_dropout: float = 0.0,
        ff_dropout: float = 0.0,
        norm_type: str = "rms",
        embedding_max_length: int = 512,
        mem_token_keep_prob: float = 1.0,
        drop_path_prob: float = 0.0,
    ):
        super().__init__()
        assert head_features % 2 == 0, "head_features must be even for RoPE"
        self.style_count = style_count
        self.channels = channels
        self.context_embedding_features = context_embedding_features
        self.d_model = d_model
        self.mem_token_keep_prob = float(mem_token_keep_prob)
        assert 0.0 < self.mem_token_keep_prob <= 1.0, "keep_prob in (0,1]"
        self.drop_path_prob = float(drop_path_prob)

        # Projekcje wejść
        self.x_in = nn.Linear(channels, d_model)
        self.mem_in = (
            nn.Linear(context_embedding_features, d_model)
            if context_embedding_features > 0
            else None
        )

        # Embedding czasu dla x i pamięci
        self.time_mlp_x = TimePositionalEmbedding(d_model, time_embed_dim)
        self.time_mlp_mem = (
            TimePositionalEmbedding(d_model, time_embed_dim)
            if context_embedding_features > 0
            else None
        )

        # Embedding ID stylu (0..T-1)
        self.style_id_emb = nn.Embedding(style_count, d_model)

        # Token-type embedding (0: style, 1: memory)
        self.type_emb = nn.Embedding(2, d_model)

        # Dropout wejściowy
        self.input_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()

        # DropPath schedule per-blok
        dps = []
        if num_layers > 0 and self.drop_path_prob > 0.0:
            steps = [i / max(1, num_layers - 1) for i in range(num_layers)]
            dps = [DropPath(self.drop_path_prob * s) for s in steps]
        else:
            dps = [nn.Identity() for _ in range(num_layers)]

        # Bloki transformera
        self.blocks = nn.ModuleList()
        for li in range(num_layers):
            self.blocks.append(
                StyleBlock(
                    d_model=d_model,
                    d_mem=d_model if context_embedding_features > 0 else 0,
                    num_heads=num_heads,
                    head_features=head_features,
                    multiplier=multiplier,
                    use_rope_sa=use_rope_sa,
                    use_rope_ca=use_rope_ca,
                    rope_max_seq_len=rope_max_seq_len,
                    dropout=dropout,
                    attn_dropout=attn_dropout,
                    ff_dropout=ff_dropout,
                    norm_type=norm_type,
                    drop_path=dps[li],
                )
            )

        # Wyjście per token stylu
        self.to_out = nn.Linear(d_model, channels)

        # Fixed (unconditional) embedding dla CFG
        self.fixed_embedding = FixedEmbedding(
            max_length=embedding_max_length, features=context_embedding_features
        )

    def _encode_style_tokens(self, x: Tensor, t: Tensor) -> Tensor:
        # x: (B, T, C), t: (B,)
        assert (
            x.size(1) == self.style_count and x.size(2) == self.channels
        ), "x must be (B, 3, 64)"
        b, t_len, _ = x.shape
        x_tok = self.x_in(x)  # (B, T, d_model)

        # Style-id emb
        style_ids = torch.arange(t_len, device=x.device).unsqueeze(0).expand(b, -1)
        x_tok = x_tok + self.style_id_emb(style_ids)

        # Token type (0)
        x_tok = x_tok + self.type_emb.weight[0]

        # Time bias
        time_bias = self.time_mlp_x(t)  # (B, d_model)
        x_tok = x_tok + time_bias.unsqueeze(1)
        return x_tok

    def _encode_memory_tokens(self, mem: Tensor, t: Tensor) -> Optional[Tensor]:
        # mem: (B, L, D)
        if self.context_embedding_features == 0 or mem.size(-1) == 0:
            return None
        mem_tok = self.mem_in(mem)  # (B, L, d_model)
        mem_tok = mem_tok + self.type_emb.weight[1]  # token-type 1
        mem_tok = mem_tok + self.time_mlp_mem(t).unsqueeze(1)
        return mem_tok

    def run(self, x: Tensor, time: Tensor, embedding: Tensor) -> Tensor:
        """
        x: (B, 3, 64)
        time: (B,)
        embedding: (B, L, D)
        return: (B, 3, 64)
        """
        # Encode
        x_tok = self._encode_style_tokens(x, time)
        mem_tok = self._encode_memory_tokens(embedding, time)

        # Dropout wejścia
        x_tok = self.input_dropout(x_tok)
        if exists(mem_tok):
            mem_tok = self.input_dropout(mem_tok)
            # Dropout tokenów pamięci (per-token scaling)
            if self.training and self.mem_token_keep_prob < 1.0:
                keep = self.mem_token_keep_prob
                b, l, _ = mem_tok.shape
                mask = mem_tok.new_empty(b, l, 1).bernoulli_(keep)
                mem_tok = mem_tok * mask / keep

        # Bloki
        for blk in self.blocks:
            x_tok = blk(x_tok, mem_tok)

        # Wyjście (B, T, C)
        out = self.to_out(x_tok)
        return out

    def forward(
        self,
        x: Tensor,  # (B, 3, 64)
        time: Tensor,  # (B,)
        embedding_mask_proba: float = 0.1,
        embedding: Optional[Tensor] = None,  # (B, L, D)
        embedding_scale: float = 1.0,
    ) -> Tensor:
        b, device = x.shape[0], x.device

        # Gdy D=0 i embedding niepodany, tworzymy placeholder (nieużywany)
        if not exists(embedding):
            if self.context_embedding_features == 0:
                embedding = x.new_zeros(b, 1, 0)
            else:
                raise AssertionError("embedding must be provided when D > 0")

        fixed_embedding = self.fixed_embedding(embedding)
        if embedding_mask_proba > 0.0:
            mask = rand_bool((b, 1, 1), proba=embedding_mask_proba, device=device)
            embedding = torch.where(mask, fixed_embedding, embedding)

        if embedding_scale != 1.0:
            out = self.run(x, time=time, embedding=embedding)
            out_u = self.run(x, time=time, embedding=fixed_embedding)
            return out_u + (out - out_u) * embedding_scale
        else:
            return self.run(x, time=time, embedding=embedding)


# -------------------------
# Flow Matching (rectified-like)
# -------------------------


class FlowMatching1DFR(nn.Module):
    """
    Flow Matching z StyleTransformer1d jako siecią prędkości.
    Ścieżka:
      x_t = ((1 - (1 - sigma) * t) * z + t * x), z ~ N(0, I), t w (0, 1]
    Cel:
      v*(x_t, t) = x - (1 - sigma) * z
    Sieć przewiduje v_theta(x_t, t, cond). Strata: MSE lub SmoothL1.
    """

    def __init__(
        self,
        net: nn.Module,  # StyleTransformer1d
        *,
        sigma: float = 1e-5,
        time_scheduler: str = "cos",  # "linear" | "cos"
        embedding_mask_proba: float = 0.1,
        loss_type: str = "smooth_l1",  # "mse" | "smooth_l1"
        huber_delta: float = 0.02,
    ):
        super().__init__()
        self.net = net
        self.sigma = float(sigma)
        self.time_scheduler = time_scheduler
        self.embedding_mask_proba = float(embedding_mask_proba)
        assert loss_type in ("mse", "smooth_l1")
        self.loss_type = loss_type
        self.huber_delta = float(huber_delta)

    @torch.no_grad()
    def forward_diffusion(
        self, x: Tensor, t: Tensor, noise: Optional[Tensor] = None
    ) -> Tuple[Tensor, Tensor, Tensor]:
        if noise is None:
            noise = torch.randn_like(x)
        t_view = t.view(-1, 1, 1)
        x_t = ((1.0 - (1.0 - self.sigma) * t_view) * noise) + (t_view * x)
        return x_t, noise, t

    def _velocity_target(self, x: Tensor, z: Tensor) -> Tensor:
        return x - (1.0 - self.sigma) * z

    def velocity(
        self,
        x_t: Tensor,
        t: Tensor,
        *,
        embedding: Tensor,
        cfg: float = 1.0,
        rescale_cfg: float = 0.0,
    ) -> Tensor:
        # Szybka ścieżka (CFG wewnątrz modelu)
        if cfg == 1.0 and rescale_cfg == 0.0:
            return self.net(
                x_t,
                time=t,
                embedding=embedding,
                embedding_mask_proba=0.0,
                embedding_scale=1.0,
            )

        # Jawny CFG + rescale
        v = self.net.run(x_t, time=t, embedding=embedding)
        fixed_emb = self.net.fixed_embedding(embedding)
        v_u = self.net.run(x_t, time=t, embedding=fixed_emb)
        v_cfg = v + cfg * (v - v_u)

        if rescale_cfg > 0.0:
            s_pos = v.std(dim=(-1, -2), keepdim=True).clamp_min(1e-8)
            s_cfg = v_cfg.std(dim=(-1, -2), keepdim=True).clamp_min(1e-8)
            v_rescaled = v_cfg * (s_pos / s_cfg)
            v_cfg = rescale_cfg * v_rescaled + (1.0 - rescale_cfg) * v_cfg
        return v_cfg

    def denoise_fn(self, x_t: Tensor, *, t: Tensor, embedding: Tensor) -> Tensor:
        v = self.net(
            x_t,
            time=t,
            embedding=embedding,
            embedding_mask_proba=0.0,
            embedding_scale=1.0,
        )
        t_view = t.view(-1, 1, 1)
        x_hat = x_t + (1.0 - t_view) * v
        return x_hat

    def forward(
        self,
        x: Tensor,  # (B, 3, 64)
        *,
        embedding: Tensor,  # (B, L, D)
        x_mask: Optional[Tensor] = None,  # opcjonalnie (B, 3)
        noise: Optional[Tensor] = None,
    ) -> Tuple[Tensor, Tensor]:
        B = x.size(0)
        device = x.device

        # Losowanie czasu
        u = torch.rand(B, device=device).clamp_(1e-5, 1.0)
        t = _time_warp(u, kind=self.time_scheduler)

        # Dyfuzja
        x_t, z, t = self.forward_diffusion(x, t, noise=noise)
        v_target = self._velocity_target(x, z)

        # Przewidywanie prędkości
        v_pred = self.net(
            x_t,
            time=t,
            embedding=embedding,
            embedding_mask_proba=self.embedding_mask_proba,
            embedding_scale=1.0,
        )

        # Maskowana MSE / SmoothL1
        if x_mask is not None:
            m = x_mask[..., None].float()
        else:
            m = torch.ones_like(v_pred)

        if self.loss_type == "mse":
            loss = F.mse_loss(v_pred, v_target, reduction="none")
        else:
            loss = F.smooth_l1_loss(
                v_pred, v_target, reduction="none", beta=self.huber_delta
            )

        loss = (loss * m).mean()

        with torch.no_grad():
            x_hat = self.denoise_fn(x_t, t=t, embedding=embedding)
        return loss, x_hat


# -------------------------
# Przykład użycia
# -------------------------
if __name__ == "__main__":
    # Konfiguracja „small-data” pod ~50k próbek
    net = StyleTransformer1d(
        style_count=3,
        channels=64,
        context_embedding_features=256,  # D twojego conditioningu
        d_model=128,
        num_layers=4,
        num_heads=4,
        head_features=32,
        multiplier=3,
        use_rope_sa=True,
        use_rope_ca=True,
        rope_max_seq_len=512,
        time_embed_dim=128,
        dropout=0.1,
        attn_dropout=0.1,
        ff_dropout=0.1,
        norm_type="rms",
        embedding_max_length=512,
        mem_token_keep_prob=0.8,  # dropout tokenów pamięci
        drop_path_prob=0.1,  # stochastic depth
    )

    fm = FlowMatching1D(
        net=net,
        sigma=1e-5,
        time_scheduler="cos",
        embedding_mask_proba=0.2,  # CFG uczone maskowaniem
        loss_type="smooth_l1",
        huber_delta=0.02,
    )

    B, L, D = 6, 14, 256
    x = torch.randn(B, 3, 64)
    cond = torch.randn(B, L, D)

    loss, x_hat = fm(x, embedding=cond)
    print("loss:", float(loss.item()), "x_hat:", tuple(x_hat.shape))