import torch
import torch.nn as nn


class QuantOHLCEmbedder(nn.Module):
    def __init__(
        self,
        num_features: int,
        sequence_length: int = 60,
        version_vocab_size: int = 4,
        hidden_dim: int = 320,
        num_layers: int = 3,
        num_heads: int = 8,
        output_dim: int = 1536,
        dtype: torch.dtype = torch.float16,
    ):
        super().__init__()
        self.num_features = num_features
        self.sequence_length = sequence_length
        self.output_dim = output_dim
        self.dtype = dtype

        self.feature_proj = nn.Sequential(
            nn.LayerNorm(num_features),
            nn.Linear(num_features, hidden_dim),
            nn.GELU(),
        )
        self.position_embedding = nn.Parameter(torch.zeros(1, sequence_length, hidden_dim))
        self.version_embedding = nn.Embedding(version_vocab_size, hidden_dim, padding_idx=0)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.output_head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.GELU(),
            nn.LayerNorm(hidden_dim * 2),
            nn.Linear(hidden_dim * 2, output_dim),
            nn.LayerNorm(output_dim),
        )
        self.to(dtype)

    def forward(
        self,
        feature_tokens: torch.Tensor,
        feature_mask: torch.Tensor,
        version_ids: torch.Tensor,
    ) -> torch.Tensor:
        if feature_tokens.ndim != 3:
            raise ValueError(f"Expected [B, T, F], got {feature_tokens.shape}")
        if feature_tokens.shape[1] != self.sequence_length:
            raise ValueError(f"Expected T={self.sequence_length}, got {feature_tokens.shape[1]}")
        if feature_tokens.shape[2] != self.num_features:
            raise ValueError(f"Expected F={self.num_features}, got {feature_tokens.shape[2]}")

        x = self.feature_proj(feature_tokens.to(self.dtype))
        version_embed = self.version_embedding(version_ids).unsqueeze(1)
        x = x + self.position_embedding[:, : x.shape[1], :].to(x.dtype) + version_embed
        key_padding_mask = ~(feature_mask > 0)
        x = self.encoder(x, src_key_padding_mask=key_padding_mask)

        mask = feature_mask.to(x.dtype).unsqueeze(-1)
        valid_any = (feature_mask.sum(dim=1, keepdim=True) > 0).to(x.dtype)
        denom = mask.sum(dim=1).clamp_min(1.0)
        pooled = (x * mask).sum(dim=1) / denom
        out = self.output_head(pooled)
        return out * valid_any