"""
Joint Fusion Model — nuFormer-style architecture combining:
  1. Transaction Transformer (pre-trained DomainTransformer) -> user embedding
  2. DCNv2 with PLR embeddings (tabular features) -> feature embedding
  3. Shared MLP head -> prediction

Architecture follows Nubank nuFormer (arXiv:2507.23267).
This is the fine-tuning architecture. Pre-training uses DomainTransformerForCausalLM alone.
"""

from typing import Dict, Optional

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

from .plr_embeddings import PeriodicLinearReLU


class DCNv2CrossLayer(nn.Module):
    """Single cross layer from Deep & Cross Network V2 (Wang et al. 2021).
    Computes: x_{l+1} = x_0 * (W_l * x_l + b_l) + x_l
    """
    def __init__(self, dim: int):
        super().__init__()
        self.weight = nn.Linear(dim, dim, bias=True)

    def forward(self, x0: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        return x0 * self.weight(x) + x


class DCNv2(nn.Module):
    """Deep & Cross Network V2 for tabular feature interaction."""
    def __init__(self, input_dim: int, cross_layers: int = 3, deep_layers: int = 2,
                 deep_dim: int = 256, dropout: float = 0.1):
        super().__init__()
        self.cross_layers = nn.ModuleList([DCNv2CrossLayer(input_dim) for _ in range(cross_layers)])
        layers = []
        in_dim = input_dim
        for _ in range(deep_layers):
            layers.extend([nn.Linear(in_dim, deep_dim), nn.ReLU(), nn.Dropout(dropout)])
            in_dim = deep_dim
        self.deep_network = nn.Sequential(*layers)
        self.output_dim = deep_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x0 = x
        cross_out = x
        for cross_layer in self.cross_layers:
            cross_out = cross_layer(x0, cross_out)
        return self.deep_network(cross_out)


class JointFusionModel(nn.Module):
    """nuFormer-style joint fusion: Transaction Transformer + DCNv2(PLR).

    Architecture:
        Transaction Sequence -> Pre-trained DomainTransformer -> user_embedding
        Tabular Features -> PLR -> flatten -> DCNv2 -> tab_embedding
        Concatenate -> MLP Head -> prediction
    """
    def __init__(self, transformer_model, n_tabular_features: int, n_classes: int = 1,
                 plr_frequencies: int = 64, plr_embedding_dim: int = 64,
                 dcn_cross_layers: int = 3, dcn_deep_layers: int = 2,
                 dcn_deep_dim: int = 256, head_hidden_dim: int = 256, dropout: float = 0.1):
        super().__init__()
        self.transformer = transformer_model
        transformer_dim = transformer_model.config.hidden_size

        self.plr = PeriodicLinearReLU(n_features=n_tabular_features,
                                       n_frequencies=plr_frequencies, embedding_dim=plr_embedding_dim)
        plr_flat_dim = n_tabular_features * plr_embedding_dim
        self.dcn = DCNv2(input_dim=plr_flat_dim, cross_layers=dcn_cross_layers,
                         deep_layers=dcn_deep_layers, deep_dim=dcn_deep_dim, dropout=dropout)

        combined_dim = transformer_dim + dcn_deep_dim
        self.head = nn.Sequential(
            nn.Linear(combined_dim, head_hidden_dim), nn.ReLU(),
            nn.Dropout(dropout), nn.Linear(head_hidden_dim, n_classes),
        )
        self.n_classes = n_classes

    def forward(self, input_ids, attention_mask=None, tabular_features=None, labels=None):
        user_embedding = self.transformer.get_user_embedding(input_ids=input_ids, attention_mask=attention_mask)
        tab_embedded = self.plr(tabular_features)
        tab_flat = tab_embedded.reshape(tab_embedded.size(0), -1)
        tab_output = self.dcn(tab_flat)

        combined = torch.cat([user_embedding, tab_output], dim=-1)
        logits = self.head(combined)

        loss = None
        if labels is not None:
            if self.n_classes == 1:
                loss = F.binary_cross_entropy_with_logits(logits.squeeze(-1), labels.float())
            else:
                loss = F.cross_entropy(logits, labels.long())

        return {"loss": loss, "logits": logits}