"""Fusion of the three resolutions: (a) shared backbone over the concatenated token sequence with resolution
embeddings; (b) separate small per-scale encoders + cross-attention fusion into the 1-minute stream."""
from __future__ import annotations

import torch
import torch.nn as nn

from .encoder import Encoder


class SharedFusion(nn.Module):
    def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float):
        super().__init__()
        self.enc = Encoder(d_model, n_layers, n_heads, d_ff, dropout)

    def forward(self, streams: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
        tok = torch.cat([t for t, _ in streams], dim=1)
        valid = torch.cat([v for _, v in streams], dim=1)
        return self.enc(tok, key_padding_mask=~valid), valid


class CrossAttentionFusion(nn.Module):
    def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float, n_streams: int = 3):
        super().__init__()
        per = max(1, n_layers // 2)
        self.encs = nn.ModuleList([Encoder(d_model, per, n_heads, d_ff, dropout) for _ in range(n_streams)])
        self.cross = nn.ModuleList([nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
                                    for _ in range(n_streams - 1)])
        self.norms = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(n_streams - 1)])
        self.post = Encoder(d_model, max(1, n_layers - per), n_heads, d_ff, dropout)

    def forward(self, streams: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
        encoded = [enc(t, key_padding_mask=~v) for enc, (t, v) in zip(self.encs, streams)]
        h, v0 = encoded[0], streams[0][1]
        for k in range(1, len(encoded)):
            kv, vk = encoded[k], streams[k][1]
            kpm = ~vk
            if kpm.all(dim=1).any():
                kpm = kpm.clone(); kpm[kpm.all(dim=1), 0] = False
            a, _ = self.cross[k - 1](self.norms[k - 1](h), kv, kv, key_padding_mask=kpm, need_weights=False)
            h = h + a
        return self.post(h, key_padding_mask=~v0), v0