"""Patch tokenisation with gated feature projection, resolution and temporal-gap embeddings."""
from __future__ import annotations

import math

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


class GatedProjection(nn.Module):
    """x -> W2(GELU(W1 x) * sigmoid(Wg x)) : a gated linear unit over the flattened patch features."""

    def __init__(self, d_in: int, d_model: int, dropout: float = 0.0):
        super().__init__()
        self.lin = nn.Linear(d_in, 2 * d_model)
        self.out = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a, g = self.lin(x).chunk(2, dim=-1)
        return self.out(self.drop(F.gelu(a) * torch.sigmoid(g)))


def sinusoidal(positions: torch.Tensor, d: int) -> torch.Tensor:
    """positions: [..., ] float in arbitrary units -> [..., d] sinusoidal features."""
    half = d // 2
    freq = torch.exp(-math.log(10000.0) * torch.arange(half, device=positions.device, dtype=torch.float32) / half)
    ang = positions.unsqueeze(-1).float() * freq
    return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)


class PatchEmbed(nn.Module):
    """Groups `patch` consecutive bars (each with f features) into one token: [B, L, f] -> [B, L/patch, d_model].
    Adds: learned resolution embedding, sinusoidal position within the window (in patch units), temporal-gap
    embedding (max log gap inside the patch), and a 'padded / missing fraction' scalar."""

    def __init__(self, f_bar: int, patch: int, d_model: int, n_resolutions: int = 3, dropout: float = 0.1):
        super().__init__()
        self.patch, self.f_bar, self.d = patch, f_bar, d_model
        self.proj = GatedProjection(patch * f_bar + 2, d_model, dropout)
        self.res_emb = nn.Embedding(n_resolutions, d_model)
        self.gap_proj = nn.Linear(d_model, d_model)
        self.pos_proj = nn.Linear(d_model, d_model)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor, mask: torch.Tensor, res_id: int, gap_feat_index: int = 10,
                missing_feat_index: int = 8) -> tuple[torch.Tensor, torch.Tensor]:
        B, L, f = x.shape
        P = L // self.patch
        xp = x[:, :P * self.patch].reshape(B, P, self.patch * f)
        mp = mask[:, :P * self.patch].reshape(B, P, self.patch).float()
        frac_present = mp.mean(-1, keepdim=True)
        gap = x[:, :P * self.patch, gap_feat_index].reshape(B, P, self.patch).max(-1, keepdim=True).values
        tok = self.proj(torch.cat([xp, frac_present, gap], dim=-1))
        pos = torch.arange(P, device=x.device, dtype=torch.float32).flip(0)  # distance to origin in patches
        tok = tok + self.pos_proj(sinusoidal(pos, self.d)).unsqueeze(0)
        tok = tok + self.gap_proj(sinusoidal(gap.squeeze(-1) * 8.0, self.d))
        tok = tok + self.res_emb.weight[res_id]
        token_valid = (frac_present.squeeze(-1) > 0)   # a patch with no present bar is still a token (masked by flag)
        return self.norm(tok), token_valid