"""Transformer_Encoder LightningModule for sign language segmentation.

Inference-only model definition. All nn.Module attribute names match the
checkpoint state_dict prefixes exactly: encoder, mixer, left_conv, right_conv,
left_hand, right_hand, input_layer, pe, sign_layer, gloss_layer,
sign_loss_function.

Adapted from segmenterv2/segmenter/model/encoder.py and
segmenterv2/segmenter/model/temporal_conv.py.
"""

from __future__ import annotations

import lightning as L
import torch
import torch.nn as nn

from sltk.segmentation.transformer.layers import (
    PositionalEncoding,
    TransformerDecoderLayer,
    TransformerEncoder,
)

# Number of output classes: OUT=0, IN=1, BEGIN=2
NUM_OUTPUT_CLASSES = 3
DEFAULT_GLOSS_VOCAB_SIZE = 11265


# ---------------------------------------------------------------------------
# Temporal 1D convolution for hand feature processing
# ---------------------------------------------------------------------------


class TemporalConv(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, conv_type: int = 2):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.conv_type = conv_type

        if self.conv_type == 3:
            self.kernel_size = ["K3"]
        elif self.conv_type == 5:
            self.kernel_size = ["K5"]

        modules: list[nn.Module] = []
        for layer_idx, ks in enumerate(self.kernel_size):
            input_sz = self.input_size if layer_idx == 0 else self.hidden_size
            if ks[0] == "P":
                modules.append(nn.MaxPool1d(kernel_size=int(ks[1]), ceil_mode=False))
            elif ks[0] == "K":
                modules.append(
                    nn.Conv1d(
                        input_sz, self.hidden_size, kernel_size=int(ks[1]), stride=1, padding=0
                    )
                )
                modules.append(nn.BatchNorm1d(self.hidden_size))
                modules.append(nn.ReLU(inplace=True))
        self.temporal_conv = nn.Sequential(*modules)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.temporal_conv(x.permute(0, 2, 1))
        return x.permute(0, 2, 1)


# ---------------------------------------------------------------------------
# Loss placeholder (for checkpoint loading compatibility)
# ---------------------------------------------------------------------------


class _XentLossPlaceholder(nn.Module):
    """Minimal placeholder matching the XentLoss state_dict structure.

    The checkpoint has a key 'sign_loss_function.criterion.weight' which is
    an nn.NLLLoss weight tensor.
    """

    def __init__(self, weight: torch.Tensor | None = None):
        super().__init__()
        self.criterion = nn.NLLLoss(ignore_index=-100, reduction="sum", weight=weight)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor):
        raise NotImplementedError("Loss computation not available in inference-only mode")


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------


class Transformer_Encoder(L.LightningModule):
    """Two-stream Transformer encoder for sign language segmentation.

    Takes 192-dim WiLoR hand pose features (96 left + 96 right) and produces
    per-frame BIO labels (OUT=0, IN=1, BEGIN=2).
    """

    def __init__(
        self,
        model_config: dict,
        train_batch_size: int = 64,
        val_batch_size: int = 64,
        dataset=None,
        model_dir: str = "",
        fps: int = 50,
        make_model: bool = True,
    ):
        super().__init__()
        self.save_hyperparameters(ignore=["text_vocab", "gloss_vocab", "dataset"])

        self.mixer = None
        self.encoder = None
        self.decoder = None
        if not make_model:
            return

        self.input_type = model_config["input_type"]
        input_size = 192
        if dataset and hasattr(dataset, "train") and dataset.train.input_size is not None:
            input_size = dataset.train.input_size
        self.input_size = input_size

        self.make_encoder(model_config["encoder"])
        self.make_feature_mixer(model_config["mixer"])

        self.adapter_size = int(model_config["encoder"]["hidden_size"] / 2)
        hidden_size = model_config["encoder"]["hidden_size"]

        self.input_layer = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
        )

        two_stream_size = int(self.input_size / 2)
        self.temp_conv = model_config["encoder"]["temp_conv_kernel"]

        self.left_conv = TemporalConv(two_stream_size, self.adapter_size, conv_type=self.temp_conv)
        self.right_conv = TemporalConv(two_stream_size, self.adapter_size, conv_type=self.temp_conv)

        self.left_hand = nn.Sequential(
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
        )

        self.right_hand = nn.Sequential(
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
            nn.Linear(self.adapter_size, self.adapter_size),
            nn.ReLU(),
        )

        self.pe = PositionalEncoding(self.adapter_size)

        self.sign_layer = nn.Linear(hidden_size, NUM_OUTPUT_CLASSES)

        gloss_vocab_length = DEFAULT_GLOSS_VOCAB_SIZE
        if dataset and hasattr(dataset, "train") and hasattr(dataset.train, "gloss_vocab"):
            if len(dataset.train.gloss_vocab) > 1000:
                gloss_vocab_length = len(dataset.train.gloss_vocab)
        self.gloss_layer = nn.Linear(hidden_size, gloss_vocab_length)

        self._init_loss_placeholder(model_config)

    def _init_loss_placeholder(self, model_config: dict):
        weight = torch.tensor([1.0, 1.0, 1.0])
        self.sign_loss_function = _XentLossPlaceholder(weight=weight)

    def encode(self, src, src_length, src_mask, **kwargs):
        encoder_output = self.encoder(src, src_length, src_mask)
        sign_prediction = self.sign_layer(encoder_output)
        gloss_prediction = self.gloss_layer(encoder_output)
        return sign_prediction, gloss_prediction

    def model_forward(self, src, src_length, src_mask, **kwargs):
        src_dual_hand = src.reshape(src.shape[0], src.shape[1], 2, 96)
        temp_lh = self.left_conv(src_dual_hand[:, :, 0, :])
        temp_rh = self.right_conv(src_dual_hand[:, :, 1, :])
        lh = self.left_hand(temp_lh)
        rh = self.right_hand(temp_rh)
        all_h = torch.cat((lh, rh), dim=-1)
        x = self.input_layer(all_h)
        return self.encode(x, src_length, src_mask, **kwargs)

    def make_encoder(self, enc_cfg: dict):
        self.encoder = TransformerEncoder(
            **enc_cfg, emb_dropout=enc_cfg["embeddings"].get("dropout", 0.0)
        )

    def make_feature_mixer(self, mixer_cfg: dict):
        self.mixer = TransformerDecoderLayer(**mixer_cfg)
