#!/usr/bin/env python
"""Experiment 8-v2: stronger unified RGB+pose AfriSign encoder.

This file reuses the Exp8 data pipeline and training loop, but swaps in a
stronger shared encoder:

  - temporal convolutional pose stem before the Transformer
  - language/modality/task-conditioned residual adapters
  - stronger projection head for supervised contrastive learning

It is intentionally a thin wrapper around `exp8_unified_mixed_encoder.py` so the
same manifests, loaders, metrics, and aggregation scripts remain compatible.
"""

from __future__ import annotations

import sys
from pathlib import Path

import torch
import torch.nn as nn

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from experiments import exp8_unified_mixed_encoder as exp8  # noqa: E402


class TemporalPoseStem(nn.Module):
    """TransSLR-style temporal stem for landmark sequences."""

    def __init__(self, feature_dim: int, hidden_dim: int, dropout: float) -> None:
        super().__init__()
        self.in_proj = nn.Linear(feature_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LayerNorm(hidden_dim),
                    nn.Conv1d(hidden_dim, hidden_dim, kernel_size=5, padding=2, groups=hidden_dim),
                    nn.GELU(),
                    nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1),
                    nn.Dropout(dropout),
                )
                for _ in range(3)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.in_proj(x)
        for block in self.blocks:
            y = block[0](x).transpose(1, 2)
            y = block[1](y)
            y = block[2](y)
            y = block[3](y).transpose(1, 2)
            y = block[4](y)
            x = x + y
        return self.out_norm(x)


class ConditionalAdapter(nn.Module):
    """Small residual adapter conditioned by language/modality/level/task context."""

    def __init__(self, hidden_dim: int, bottleneck: int, dropout: float) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(hidden_dim)
        self.down = nn.Linear(hidden_dim, bottleneck)
        self.up = nn.Linear(bottleneck, hidden_dim)
        self.drop = nn.Dropout(dropout)
        self.gate = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.Sigmoid())

    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
        update = self.up(torch.nn.functional.gelu(self.down(self.norm(x))))
        gate = self.gate(context).unsqueeze(1)
        return x + self.drop(update * gate)


class StrongUnifiedAfriSignEncoder(nn.Module):
    def __init__(
        self,
        *,
        task_dims: dict[str, int],
        num_languages: int,
        hidden_dim: int,
        feature_dim: int,
        max_tokens: int,
        layers: int,
        heads: int,
        ff_dim: int,
        dropout: float,
        rgb_backbone: str,
        rgb_train_backbone: str,
        rgb_pretrained: bool,
    ) -> None:
        super().__init__()
        self.task_keys = list(task_dims)
        self.pose_stem = TemporalPoseStem(feature_dim, hidden_dim, dropout)
        if rgb_backbone == "small_cnn":
            self.rgb_cnn = exp8.SmallFrameCNN(hidden_dim)
        elif rgb_backbone == "efficientnet_b0":
            self.rgb_cnn = exp8.EfficientNetB0FrameEncoder(hidden_dim, rgb_train_backbone, rgb_pretrained)
        else:
            raise ValueError(f"Unknown rgb_backbone={rgb_backbone!r}")

        self.cls = nn.Parameter(torch.zeros(1, 1, hidden_dim))
        self.pos = nn.Embedding(max_tokens + 1, hidden_dim)
        self.lang_emb = nn.Embedding(num_languages, hidden_dim)
        self.modality_emb = nn.Embedding(len(exp8.MODALITY_IDS), hidden_dim)
        self.level_emb = nn.Embedding(len(exp8.LEVEL_IDS), hidden_dim)
        self.task_emb = nn.Embedding(len(self.task_keys), hidden_dim)

        self.pre_adapter = ConditionalAdapter(hidden_dim, max(32, hidden_dim // 4), dropout)
        enc_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=heads,
            dim_feedforward=ff_dim,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(enc_layer, layers, enable_nested_tensor=False)
        self.post_adapter = ConditionalAdapter(hidden_dim, max(32, hidden_dim // 4), dropout)
        self.norm = nn.LayerNorm(hidden_dim)
        self.drop = nn.Dropout(dropout)
        self.heads = nn.ModuleDict({key: nn.Linear(hidden_dim, dim) for key, dim in task_dims.items()})
        self.projector = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
        )
        nn.init.trunc_normal_(self.cls, std=0.02)

    def encode(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        if "pose" in batch:
            x = self.pose_stem(batch["pose"])
        elif "rgb" in batch:
            rgb = batch["rgb"]
            b, t, c, h, w = rgb.shape
            x = self.rgb_cnn(rgb.reshape(b * t, c, h, w)).reshape(b, t, -1)
        else:
            raise ValueError("Batch must contain pose or rgb")

        b, t, _ = x.shape
        positions = torch.arange(t + 1, device=x.device)
        context = (
            self.lang_emb(batch["lang_idx"])
            + self.modality_emb(batch["modality_idx"])
            + self.level_emb(batch["level_idx"])
            + self.task_emb(batch["task_idx"])
        )
        x = torch.cat([self.cls.expand(b, -1, -1), x], dim=1)
        x = x + self.pos(positions).unsqueeze(0) + context.unsqueeze(1)
        x = self.pre_adapter(x, context)
        x = self.encoder(x)
        x = self.post_adapter(x, context)
        return self.drop(self.norm(x[:, 0]))

    def forward(self, batch: dict[str, torch.Tensor], task_key: str) -> torch.Tensor:
        return self.heads[task_key](self.encode(batch))

    def contrast_features(self, features: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.normalize(self.projector(features), dim=1)


def main() -> None:
    exp8.UnifiedAfriSignEncoder = StrongUnifiedAfriSignEncoder
    exp8.main()


if __name__ == "__main__":
    main()