"""Parameter-tied recurrence with additive input injection."""

from __future__ import annotations

import torch
import torch.nn.functional as F

from benchmark import (
    ModelSpec,
    OptimizerBundle,
    OptimizerSpec,
    Submission,
    assert_model_state,
)
from model import (
    AttentionConfig,
    BlockConfig,
    InputInjectionConfig,
    LoopConfig,
    MixerConfig,
    ModelConfig,
    Transformer,
)


MODEL_DEPTH = 4


def build_model(spec: ModelSpec):
    model = Transformer(
        ModelConfig(
            max_seq_len=spec.max_seq_len,
            vocab_size=spec.vocab_size,
            num_classes=spec.num_classes,
            depth=1,
            block=BlockConfig(
                d_model=128,
                attention=AttentionConfig(num_heads=4),
                mixer=MixerConfig(expansion_factor=4.0, activation="gelu"),
            ),
            loop=LoopConfig(
                max_loops=MODEL_DEPTH,
                block_input_injection=InputInjectionConfig(kind="add"),
            ),
            position_embedding="learned",
        )
    )
    assert_model_state(model, spec)
    return model


def build_optimizer(model, spec: OptimizerSpec) -> OptimizerBundle:
    return OptimizerBundle(
        torch.optim.AdamW(
            model.parameters(),
            lr=1e-3,
            betas=(0.9, 0.95),
            weight_decay=0.1,
            capturable=spec.device_type == "cuda",
        )
    )


def training_loss(logits, labels, auxiliary):
    if logits.ndim == 2 and logits.shape[-1] > 1:
        loss = F.cross_entropy(logits, labels.long())
        return loss if auxiliary is None else loss + auxiliary
    labels = labels.float()
    bce = F.binary_cross_entropy_with_logits(logits, labels, reduction="none")
    target_probability = torch.where(
        labels.bool(), torch.sigmoid(logits), torch.sigmoid(-logits)
    )
    loss = ((1 - target_probability).square() * bce).mean()
    return loss if auxiliary is None else loss + auxiliary


SUBMISSION = Submission(
    build_model=build_model,
    build_optimizer=build_optimizer,
    training_loss=training_loss,
)