from __future__ import annotations

from contextlib import nullcontext
from functools import partial
from typing import Callable, Optional

import torch
import torch.nn as nn

from .oat_transformer import Fp32LayerNorm


def str_to_dtype(dtype_str: Optional[str]) -> Optional[torch.dtype]:
    if dtype_str is None:
        return None
    if dtype_str in {"float16", "fp16"}:
        return torch.float16
    if dtype_str in {"bfloat16", "bf16"}:
        return torch.bfloat16
    if dtype_str in {"float32", "fp32"}:
        return torch.float32
    raise ValueError(f"Invalid dtype string representation: {dtype_str}")


def get_autocast_context(x: torch.Tensor, dtype_override: Optional[torch.dtype]):
    if dtype_override is None:
        return nullcontext()
    return torch.amp.autocast(
        device_type=x.device.type,
        dtype=dtype_override,
        enabled=dtype_override != torch.float32,
    )


class SampleEmbedder(nn.Module):
    def __init__(
        self,
        dim_in: int,
        dim_out: Optional[int] = None,
        weight_init_style: str = "xavier",
    ):
        super().__init__()
        self.dim_in = dim_in
        self.dim_out = dim_out
        if dim_out is None:
            self.proj = nn.Identity()
            self.dim_out = dim_in
        else:
            self.proj = nn.Linear(dim_in, dim_out, bias=True)
        self.weight_init_style = weight_init_style
        self._init_weights()

    def _init_weights(self) -> None:
        if isinstance(self.proj, nn.Identity):
            return
        if self.weight_init_style == "zero":
            nn.init.constant_(self.proj.weight, 0)
        elif self.weight_init_style == "xavier":
            nn.init.xavier_uniform_(self.proj.weight)
        elif self.weight_init_style == "trunc_normal":
            nn.init.trunc_normal_(self.proj.weight, std=0.02)
        else:
            raise ValueError(f"Unsupported weight init: {self.weight_init_style}")
        if self.proj.bias is not None:
            nn.init.constant_(self.proj.bias, 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x)


class LinearLayer(nn.Module):
    def __init__(
        self,
        dim_in: int,
        dim_out: int,
        weight_init_style: str = "xavier",
        dtype_override: Optional[str] = None,
        proj_bias: bool = True,
    ):
        super().__init__()
        self.dtype_override = str_to_dtype(dtype_override)
        self.proj = nn.Linear(dim_in, dim_out, bias=proj_bias)
        self.weight_init_style = weight_init_style
        self._init_weights()

    def _init_weights(self) -> None:
        if self.weight_init_style == "zero":
            nn.init.constant_(self.proj.weight, 0)
        elif self.weight_init_style == "xavier":
            nn.init.xavier_uniform_(self.proj.weight)
        elif self.weight_init_style == "trunc_normal":
            nn.init.trunc_normal_(self.proj.weight, std=0.02)
        else:
            raise ValueError(f"Unsupported weight init: {self.weight_init_style}")
        if self.proj.bias is not None:
            nn.init.constant_(self.proj.bias, 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with get_autocast_context(x, self.dtype_override):
            return self.proj(x)


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1.0 + scale) + shift


def _expand_emb_from_seq(emb: torch.Tensor, seq_len: int) -> torch.Tensor:
    if emb.ndim == 2:
        emb = emb.unsqueeze(1)
    if emb.shape[1] == seq_len:
        return emb
    return emb.expand(-1, seq_len, -1)


class LinearHead(nn.Module):
    def __init__(
        self,
        dim: int,
        dim_out: int,
        weight_init_style: str = "zero",
        norm_layer: Optional[partial] = partial(
            Fp32LayerNorm, bias=False, elementwise_affine=False
        ),
        dtype_override: Optional[str] = None,
        use_adaLN: bool = False,
        adaLN_bias: bool = True,
        proj_bias: bool = True,
    ):
        super().__init__()
        self.dtype_override = str_to_dtype(dtype_override)
        self.norm = norm_layer(dim) if norm_layer is not None else nn.Identity()
        self.use_adaLN = use_adaLN
        if use_adaLN:
            self.adaLN_modulation = nn.Sequential(
                nn.SiLU(),
                nn.Linear(dim, 2 * dim, bias=adaLN_bias),
            )
        self.proj = nn.Linear(dim, dim_out, bias=proj_bias)
        self.weight_init_style = weight_init_style
        self._init_weights()

    def _init_weights(self) -> None:
        for name, m in self.named_modules():
            if isinstance(m, nn.Linear):
                if self.weight_init_style == "zero" or "adaLN_modulation" in name:
                    nn.init.constant_(m.weight, 0)
                elif self.weight_init_style == "xavier":
                    nn.init.xavier_uniform_(m.weight)
                elif self.weight_init_style == "trunc_normal":
                    nn.init.trunc_normal_(m.weight, std=0.02)
                else:
                    raise ValueError(
                        f"Unsupported weight init: {self.weight_init_style}"
                    )
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                if m.weight is not None:
                    nn.init.constant_(m.weight, 1.0)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(
        self,
        x: torch.Tensor,
        adaLN_emb: Optional[torch.Tensor] = None,
        adaLN_packing_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
    ) -> torch.Tensor:
        with get_autocast_context(x, self.dtype_override):
            x = self.norm(x)
            if self.use_adaLN and adaLN_emb is not None:
                packed = self.adaLN_modulation(adaLN_emb)
                packed = (
                    adaLN_packing_fn(packed)
                    if adaLN_packing_fn is not None
                    else _expand_emb_from_seq(packed, x.shape[1])
                )
                shift, scale = packed.chunk(2, dim=-1)
                x = modulate(x, shift, scale)
            x = self.proj(x)
        return x
