from typing import Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from diffusers.models.attention_processor import Attention


class LoRALinearLayer(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 4,
        network_alpha: Optional[float] = None,
        device: Optional[Union[torch.device, str]] = None,
        dtype: Optional[torch.dtype] = None,
        number=0,
        n_loras=1,
    ):
        super().__init__()
        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
        self.network_alpha = network_alpha
        self.rank = rank
        self.out_features = out_features
        self.in_features = in_features

        nn.init.normal_(self.down.weight, std=1 / rank)
        nn.init.zeros_(self.up.weight)

        self.number = number
        self.n_loras = n_loras

    def forward(self, hidden_states: torch.Tensor, cond_seq_len: int = None) -> torch.Tensor:
        orig_dtype = hidden_states.dtype
        dtype = self.down.weight.dtype

        batch_size = hidden_states.shape[0]
        cond_size = cond_seq_len

        block_size = hidden_states.shape[1] - cond_size * self.n_loras
        shape = (batch_size, hidden_states.shape[1], 3072)
        mask = torch.ones(shape, device=hidden_states.device, dtype=dtype)
        mask[:, : block_size + self.number * cond_size, :] = 0
        mask[:, block_size + (self.number + 1) * cond_size :, :] = 0
        hidden_states = mask * hidden_states

        down_hidden_states = self.down(hidden_states.to(dtype))
        up_hidden_states = self.up(down_hidden_states)

        if self.network_alpha is not None:
            up_hidden_states *= self.network_alpha / self.rank

        return up_hidden_states.to(orig_dtype)


class TextLoRALinearLayer(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 4,
        network_alpha: Optional[float] = None,
        device: Optional[Union[torch.device, str]] = None,
        dtype: Optional[torch.dtype] = None,
        token_length=512,
    ):
        super().__init__()
        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
        self.network_alpha = network_alpha
        self.rank = rank
        self.out_features = out_features
        self.in_features = in_features

        nn.init.normal_(self.down.weight, std=1 / rank)
        nn.init.zeros_(self.up.weight)

        self.token_length = token_length

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        orig_dtype = hidden_states.dtype
        dtype = self.down.weight.dtype

        batch_size, seq_len, feature_dim = hidden_states.shape
        if seq_len > self.token_length:
            mask = torch.ones((batch_size, seq_len, feature_dim), device=hidden_states.device, dtype=dtype)
            mask[:, self.token_length :, :] = 0
            hidden_states = mask * hidden_states

        down_hidden_states = self.down(hidden_states.to(dtype))
        up_hidden_states = self.up(down_hidden_states)

        if self.network_alpha is not None:
            up_hidden_states *= self.network_alpha / self.rank

        return up_hidden_states.to(orig_dtype)


class MultiSingleStreamBlockLoraProcessor(nn.Module):
    def __init__(
        self,
        dim: int,
        ranks=[],
        lora_weights=[],
        network_alphas=[],
        device=None,
        dtype=None,
        n_loras=1,
        text_lora_config=None,
    ):
        super().__init__()
        self.n_loras = n_loras
        if text_lora_config is not None:
            self.text_len = text_lora_config.get("token_length", 512)
        else:
            self.text_len = 512

        self.q_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.k_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.v_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.lora_weights = lora_weights

        if text_lora_config is not None:
            t_rank = text_lora_config.get("rank", 4)
            t_alpha = text_lora_config.get("alpha", None)

            self.text_q_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len
            )
            self.text_k_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len
            )
            self.text_v_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=self.text_len
            )

    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
        use_cond=False,
    ) -> torch.FloatTensor:
        batch_size, seq_len, _ = hidden_states.shape

        total_img_seq_len = seq_len - self.text_len
        assert total_img_seq_len % (1 + self.n_loras) == 0, (
            f"total_img_seq_len:{total_img_seq_len}, n_loras:{self.n_loras}, "
            f"seq_len:{seq_len}, text_len:{self.text_len}"
        )
        cond_seq_len = total_img_seq_len // (1 + self.n_loras)

        query = attn.to_q(hidden_states)
        key = attn.to_k(hidden_states)
        value = attn.to_v(hidden_states)

        for i in range(self.n_loras):
            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states, cond_seq_len=cond_seq_len)
            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states, cond_seq_len=cond_seq_len)
            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states, cond_seq_len=cond_seq_len)

        if getattr(self, "text_q_lora", None) is not None:
            query = query + self.text_q_lora(hidden_states)
        if getattr(self, "text_k_lora", None) is not None:
            key = key + self.text_k_lora(hidden_states)
        if getattr(self, "text_v_lora", None) is not None:
            value = value + self.text_v_lora(hidden_states)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads

        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        if image_rotary_emb is not None:
            from diffusers.models.embeddings import apply_rotary_emb

            query = apply_rotary_emb(query, image_rotary_emb)
            key = apply_rotary_emb(key, image_rotary_emb)

        cond_size = cond_seq_len
        block_size = hidden_states.shape[1] - cond_size * self.n_loras

        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        cond_hidden_states = hidden_states[:, block_size:, :]
        hidden_states = hidden_states[:, :block_size, :]

        return hidden_states if not use_cond else (hidden_states, cond_hidden_states)


class MultiDoubleStreamBlockLoraProcessor(nn.Module):
    def __init__(
        self,
        dim: int,
        ranks=[],
        lora_weights=[],
        network_alphas=[],
        device=None,
        dtype=None,
        n_loras=1,
        text_lora_config=None,
    ):
        super().__init__()
        self.n_loras = n_loras
        self.q_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.k_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.v_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.proj_loras = nn.ModuleList(
            [
                LoRALinearLayer(
                    dim, dim, ranks[i], network_alphas[i], device=device, dtype=dtype, number=i, n_loras=n_loras
                )
                for i in range(n_loras)
            ]
        )
        self.lora_weights = lora_weights
        if text_lora_config is not None:
            t_rank = text_lora_config.get("rank", 4)
            t_alpha = text_lora_config.get("alpha", None)
            t_len = text_lora_config.get("token_length", 512)

            self.text_q_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len
            )
            self.text_k_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len
            )
            self.text_v_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len
            )
            self.text_proj_lora = TextLoRALinearLayer(
                dim, dim, t_rank, t_alpha, device=device, dtype=dtype, token_length=t_len
            )

    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
        use_cond=False,
    ) -> torch.FloatTensor:
        batch_size, total_img_seq_len, _ = hidden_states.shape

        # `context` projections.
        inner_dim = 3072
        head_dim = inner_dim // attn.heads
        encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
        encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
        encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)

        if getattr(self, "text_q_lora", None) is not None:
            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj + self.text_q_lora(
                encoder_hidden_states
            )
        if getattr(self, "text_k_lora", None) is not None:
            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj + self.text_k_lora(encoder_hidden_states)
        if getattr(self, "text_v_lora", None) is not None:
            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj + self.text_v_lora(
                encoder_hidden_states
            )

        encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
            batch_size, -1, attn.heads, head_dim
        ).transpose(1, 2)
        encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
            batch_size, -1, attn.heads, head_dim
        ).transpose(1, 2)
        encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
            batch_size, -1, attn.heads, head_dim
        ).transpose(1, 2)

        if attn.norm_added_q is not None:
            encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
        if attn.norm_added_k is not None:
            encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)

        assert total_img_seq_len % (1 + self.n_loras) == 0, (
            f"total_img_seq_len:{total_img_seq_len}, n_loras:{self.n_loras}"
        )
        cond_seq_len = total_img_seq_len // (1 + self.n_loras)

        query = attn.to_q(hidden_states)
        key = attn.to_k(hidden_states)
        value = attn.to_v(hidden_states)

        for i in range(self.n_loras):
            query = query + self.lora_weights[i] * self.q_loras[i](hidden_states, cond_seq_len=cond_seq_len)
            key = key + self.lora_weights[i] * self.k_loras[i](hidden_states, cond_seq_len=cond_seq_len)
            value = value + self.lora_weights[i] * self.v_loras[i](hidden_states, cond_seq_len=cond_seq_len)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads
        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        # attention
        query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
        key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
        value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)

        if image_rotary_emb is not None:
            from diffusers.models.embeddings import apply_rotary_emb

            query = apply_rotary_emb(query, image_rotary_emb)
            key = apply_rotary_emb(key, image_rotary_emb)

        cond_size = cond_seq_len
        block_size = hidden_states.shape[1] - cond_size * self.n_loras

        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        encoder_hidden_states, hidden_states = (
            hidden_states[:, : encoder_hidden_states.shape[1]],
            hidden_states[:, encoder_hidden_states.shape[1] :],
        )

        hidden_states_input = hidden_states
        hidden_states = attn.to_out[0](hidden_states)
        for i in range(self.n_loras):
            hidden_states = hidden_states + self.lora_weights[i] * self.proj_loras[i](
                hidden_states_input, cond_seq_len=cond_seq_len
            )

        hidden_states = attn.to_out[1](hidden_states)

        encoder_input = encoder_hidden_states
        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
        if getattr(self, "text_proj_lora", None) is not None:
            encoder_hidden_states = encoder_hidden_states + self.text_proj_lora(encoder_input)

        cond_hidden_states = hidden_states[:, block_size:, :]
        hidden_states = hidden_states[:, :block_size, :]

        return (
            (hidden_states, encoder_hidden_states, cond_hidden_states)
            if use_cond
            else (encoder_hidden_states, hidden_states)
        )