"""Load Qwen2.5-Omni-3B and extract causal audio-video features.

This research example does not generate text or speech: only the Thinker
(multimodal understanding model) is loaded.

Example:
    python qwen_omni_features.py --video clip.mp4 \
        --use-audio-in-video --output clip_features.pt
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Mapping

import torch
from qwen_omni_utils import process_mm_info
from transformers import (
    Qwen2_5OmniProcessor,
    Qwen2_5OmniThinkerForConditionalGeneration,
)


MODEL_ID = "Qwen/Qwen2.5-Omni-3B"


@dataclass
class OmniFeatures:
    """Features returned for one or more audio-video examples."""

    token_features: torch.Tensor  # [B, T, D]
    attention_mask: torch.Tensor  # [B, T]
    last_token: torch.Tensor  # [B, D]
    tail_mean: torch.Tensor  # [B, D]
    hidden_states: tuple[torch.Tensor, ...] | None = None


def choose_dtype(dtype: Literal["auto", "bfloat16", "float16", "float32"]):
    if dtype == "auto":
        return "auto"
    return getattr(torch, dtype)


def load_qwen_omni(
    model_id: str = MODEL_ID,
    *,
    dtype: Literal["auto", "bfloat16", "float16", "float32"] = "auto",
    device_map: str | Mapping[str, str | int] = "auto",
    use_flash_attention_2: bool = False,
) -> tuple[Qwen2_5OmniThinkerForConditionalGeneration, Qwen2_5OmniProcessor]:
    """Load the multimodal Thinker and processor without the Talker."""

    model_kwargs: dict[str, Any] = {
        "torch_dtype": choose_dtype(dtype),
        "device_map": device_map,
    }
    if use_flash_attention_2:
        model_kwargs["attn_implementation"] = "flash_attention_2"

    model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
        model_id, **model_kwargs
    )
    processor = Qwen2_5OmniProcessor.from_pretrained(model_id)
    model.eval()
    return model, processor


def build_conversation(
    *,
    video_path: str | Path | None = None,
    audio_path: str | Path | None = None,
    prompt: str = (
        "Observe the conversation up to the current time and represent the "
        "multimodal evidence useful for predicting the next speaking turn."
    ),
) -> list[dict[str, Any]]:
    """Build one Qwen chat containing video, optional audio, and instruction."""

    if video_path is None and audio_path is None:
        raise ValueError("At least one of video_path and audio_path is required")

    content: list[dict[str, str]] = []
    if video_path is not None:
        content.append({"type": "video", "video": str(Path(video_path).resolve())})
    if audio_path is not None:
        content.append({"type": "audio", "audio": str(Path(audio_path).resolve())})
    content.append({"type": "text", "text": prompt})

    return [
        {
            "role": "system",
            "content": [{
                "type": "text",
                "text": (
                    "You are a causal multimodal conversation observer. "
                    "Use only information available up to the current time."
                ),
            }],
        },
        {"role": "user", "content": content},
    ]


def prepare_inputs(
    processor: Qwen2_5OmniProcessor,
    conversations: list[list[dict[str, Any]]],
    *,
    use_audio_in_video: bool,
    model_device: torch.device,
    model_dtype: torch.dtype,
):
    """Convert a batch of conversations into model tensors."""

    texts = processor.apply_chat_template(
        conversations, add_generation_prompt=True, tokenize=False
    )
    audios, images, videos = process_mm_info(
        conversations, use_audio_in_video=use_audio_in_video
    )
    inputs = processor(
        text=texts,
        audio=audios,
        images=images,
        videos=videos,
        padding=True,
        return_tensors="pt",
        use_audio_in_video=use_audio_in_video,
    )
    # BatchFeature.to(dtype) casts floating tensors only; IDs/masks stay integer.
    return inputs.to(model_device).to(model_dtype)


def _last_valid_token(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    positions = torch.arange(states.shape[1], device=states.device).unsqueeze(0)
    indices = positions.masked_fill(~mask.bool(), -1).max(dim=1).values.clamp_min(0)
    batch_indices = torch.arange(states.shape[0], device=states.device)
    return states[batch_indices, indices]


def _tail_mean(
    states: torch.Tensor, mask: torch.Tensor, tail_tokens: int
) -> torch.Tensor:
    if tail_tokens <= 0:
        raise ValueError("tail_tokens must be positive")
    # Ranking valid tokens from the right supports both left and right padding.
    reverse_ranks = mask.long().flip(1).cumsum(dim=1).flip(1)
    tail_mask = mask.bool() & (reverse_ranks <= tail_tokens)
    weights = tail_mask.unsqueeze(-1).to(states.dtype)
    return (states * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1)


@torch.inference_mode()
def extract_features(
    model: Qwen2_5OmniThinkerForConditionalGeneration,
    processor: Qwen2_5OmniProcessor,
    conversations: list[list[dict[str, Any]]],
    *,
    use_audio_in_video: bool = False,
    layer: int = -1,
    tail_tokens: int = 8,
    return_all_layers: bool = False,
    output_device: str | torch.device = "cpu",
) -> OmniFeatures:
    """Extract fused Thinker states without generation.

    ``last_token`` is the recommended one-vector-per-window baseline. The
    Thinker is causal, so it summarizes all earlier prompt and AV tokens.
    """

    inputs = prepare_inputs(
        processor,
        conversations,
        use_audio_in_video=use_audio_in_video,
        model_device=model.device,
        model_dtype=model.dtype,
    )
    # A direct forward ends at the observation boundary; generate() would append
    # answer tokens and change which representation is being extracted.
    outputs = model(
        **inputs,
        use_audio_in_video=use_audio_in_video,
        output_hidden_states=True,
        use_cache=False,
        return_dict=True,
    )
    if outputs.hidden_states is None:
        raise RuntimeError("The model did not return hidden states")

    states = outputs.hidden_states[layer]
    mask = inputs["attention_mask"]
    return OmniFeatures(
        token_features=states.to(output_device),
        attention_mask=mask.to(output_device),
        last_token=_last_valid_token(states, mask).to(output_device),
        tail_mean=_tail_mean(states, mask, tail_tokens).to(output_device),
        hidden_states=(
            tuple(x.to(output_device) for x in outputs.hidden_states)
            if return_all_layers else None
        ),
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--video", type=Path)
    parser.add_argument("--audio", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--model-id", default=MODEL_ID)
    parser.add_argument("--prompt", default=None)
    parser.add_argument("--use-audio-in-video", action="store_true")
    parser.add_argument("--tail-tokens", type=int, default=8)
    parser.add_argument("--layer", type=int, default=-1)
    parser.add_argument(
        "--dtype",
        choices=("auto", "bfloat16", "float16", "float32"),
        default="auto",
    )
    parser.add_argument("--flash-attention-2", action="store_true")
    parser.add_argument("--device", default="cuda:0")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.video is None and args.audio is None:
        raise SystemExit("Pass --video, --audio, or both")
    if args.use_audio_in_video and args.video is None:
        raise SystemExit("--use-audio-in-video requires --video")

    model, processor = load_qwen_omni(
        args.model_id,
        dtype=args.dtype,
        device_map={"": args.device},
        use_flash_attention_2=args.flash_attention_2,
    )
    conversation_kwargs: dict[str, Any] = {
        "video_path": args.video,
        "audio_path": args.audio,
    }
    if args.prompt is not None:
        conversation_kwargs["prompt"] = args.prompt
    conversation = build_conversation(**conversation_kwargs)
    features = extract_features(
        model,
        processor,
        [conversation],
        use_audio_in_video=args.use_audio_in_video,
        layer=args.layer,
        tail_tokens=args.tail_tokens,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    torch.save({
        "model_id": args.model_id,
        "layer": args.layer,
        "token_features": features.token_features,
        "attention_mask": features.attention_mask,
        "last_token": features.last_token,
        "tail_mean": features.tail_mean,
    }, args.output)
    print(f"Saved features to {args.output}")
    print(f"token_features: {tuple(features.token_features.shape)}")
    print(f"last_token:     {tuple(features.last_token.shape)}")
    print(f"tail_mean:      {tuple(features.tail_mean.shape)}")


if __name__ == "__main__":
    main()