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

from .utils import split_feature, merge_splits


def gather_view_tensor(tensor, view_indices):
    view_indices = view_indices.to(device=tensor.device, dtype=torch.long)
    batch_indices = torch.arange(tensor.size(0), device=tensor.device)
    batch_indices = batch_indices.reshape(
        tensor.size(0), *([1] * (view_indices.dim() - 1))
    )
    return tensor[batch_indices, view_indices]


def non_self_view_indices(num_views, device):
    view_ids = torch.arange(num_views, device=device)
    return torch.stack(
        [torch.cat((view_ids[:i], view_ids[i + 1 :])) for i in range(num_views)],
        dim=0,
    )


def target_view_indices(batch_size, num_views, nn_matrix, device):
    if nn_matrix is not None:
        return nn_matrix[..., 1:].to(device=device, dtype=torch.long)
    return non_self_view_indices(num_views, device).unsqueeze(0).expand(
        batch_size, -1, -1
    )


def flatten_batch_view(tensor):
    # Batch-major flatten: [B, V, ...] -> [B*V, ...].
    return tensor.flatten(0, 1)


def flatten_view_batch(tensor):
    # View-major flatten: [B, V, ...] -> [V*B, ...].
    return tensor.transpose(0, 1).reshape(
        tensor.size(0) * tensor.size(1), *tensor.shape[2:]
    )


def unflatten_view_batch(tensor, batch_size, num_views):
    # Inverse of flatten_view_batch.
    return tensor.reshape(num_views, batch_size, *tensor.shape[1:]).transpose(0, 1)


def single_head_full_attention(q, k, v):
    # q, k, v: [B, L, C]
    assert q.dim() == k.dim() == v.dim() == 3

    return F.scaled_dot_product_attention(
        q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1)
    ).squeeze(1)


def generate_shift_window_attn_mask(
    input_resolution,
    window_size_h,
    window_size_w,
    shift_size_h,
    shift_size_w,
    device=torch.device("cuda"),
):
    # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
    # Build the shifted-window attention mask used by Swin-style local attention.
    h, w = input_resolution
    img_mask = torch.zeros((1, h, w, 1)).to(device)  # 1 H W 1
    h_slices = (
        slice(0, -window_size_h),
        slice(-window_size_h, -shift_size_h),
        slice(-shift_size_h, None),
    )
    w_slices = (
        slice(0, -window_size_w),
        slice(-window_size_w, -shift_size_w),
        slice(-shift_size_w, None),
    )
    cnt = 0
    for h in h_slices:
        for w in w_slices:
            img_mask[:, h, w, :] = cnt
            cnt += 1

    mask_windows = split_feature(
        img_mask, num_splits=input_resolution[-1] // window_size_w, channel_last=True
    )

    mask_windows = mask_windows.view(-1, window_size_h * window_size_w)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
        attn_mask == 0, float(0.0)
    )

    return attn_mask


def single_head_split_window_attention(
    q,
    k,
    v,
    num_splits=1,
    with_shift=False,
    h=None,
    w=None,
    attn_mask=None,
):
    # Ref: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
    # q, k, v: [B, L, C] for self-attention or pairwise cross-attention.
    # For multi-view cross-attention, q: [B, L, C], k/v: [B, M, L, C].

    # Multi-view cross-attention, where M is the number of target views per query.
    if not (q.dim() == k.dim() == v.dim() == 3):
        assert k.dim() == v.dim() == 4
        assert h is not None and w is not None
        assert q.size(1) == h * w

        m = k.size(1)

        b, _, c = q.size()

        b_new = b * num_splits * num_splits

        window_size_h = h // num_splits
        window_size_w = w // num_splits

        q = q.view(b, h, w, c)  # [B, H, W, C]
        k = k.view(b, m, h, w, c)  # [B, M, H, W, C]
        v = v.view(b, m, h, w, c)

        if with_shift:
            assert attn_mask is not None  # compute once
            shift_size_h = window_size_h // 2
            shift_size_w = window_size_w // 2

            q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
            k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3))
            v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(2, 3))

        q = split_feature(
            q, num_splits=num_splits, channel_last=True
        )  # [B*K*K, H/K, W/K, C]
        k = split_feature(
            k.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1),
            num_splits=num_splits,
            channel_last=True,
        )  # [B*K*K, H/K, W/K, C*M]
        v = split_feature(
            v.permute(0, 2, 3, 4, 1).reshape(b, h, w, -1),
            num_splits=num_splits,
            channel_last=True,
        )  # [B*K*K, H/K, W/K, C*M]

        k = (
            k.view(b_new, h // num_splits, w // num_splits, c, m)
            .permute(0, 1, 2, 4, 3)
            .reshape(b_new, -1, c)
        )  # [B*K*K, H/K*W/K*M, C]
        v = (
            v.view(b_new, h // num_splits, w // num_splits, c, m)
            .permute(0, 1, 2, 4, 3)
            .reshape(b_new, -1, c)
        )  # [B*K*K, H/K*W/K*M, C]

        attn_bias = None
        if with_shift:
            attn_bias = attn_mask.repeat(b, 1, m).unsqueeze(1).to(q.dtype)

        out = F.scaled_dot_product_attention(
            q.view(b_new, -1, c).unsqueeze(1),
            k.unsqueeze(1),
            v.unsqueeze(1),
            attn_mask=attn_bias,
        ).squeeze(1)

        out = merge_splits(
            out.view(b_new, h // num_splits, w // num_splits, c),
            num_splits=num_splits,
            channel_last=True,
        )  # [B, H, W, C]

        # shift back
        if with_shift:
            out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))

        out = out.view(b, -1, c)
    else:
        # Self-attention or pairwise cross-attention.
        assert q.dim() == k.dim() == v.dim() == 3

        assert h is not None and w is not None
        assert q.size(1) == h * w

        b, _, c = q.size()

        b_new = b * num_splits * num_splits

        window_size_h = h // num_splits
        window_size_w = w // num_splits

        q = q.view(b, h, w, c)  # [B, H, W, C]
        k = k.view(b, h, w, c)
        v = v.view(b, h, w, c)

        if with_shift:
            assert attn_mask is not None  # compute once
            shift_size_h = window_size_h // 2
            shift_size_w = window_size_w // 2

            q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
            k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
            v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))

        q = split_feature(
            q, num_splits=num_splits, channel_last=True
        )  # [B*K*K, H/K, W/K, C]
        k = split_feature(k, num_splits=num_splits, channel_last=True)
        v = split_feature(v, num_splits=num_splits, channel_last=True)

        attn_bias = None
        if with_shift:
            attn_bias = attn_mask.repeat(b, 1, 1).unsqueeze(1).to(q.dtype)

        out = F.scaled_dot_product_attention(
            q.view(b_new, -1, c).unsqueeze(1),
            k.view(b_new, -1, c).unsqueeze(1),
            v.view(b_new, -1, c).unsqueeze(1),
            attn_mask=attn_bias,
        ).squeeze(1)

        out = merge_splits(
            out.view(b_new, h // num_splits, w // num_splits, c),
            num_splits=num_splits,
            channel_last=True,
        )  # [B, H, W, C]

        # shift back
        if with_shift:
            out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))

        out = out.view(b, -1, c)

    return out


def multi_head_split_window_attention(
    q,
    k,
    v,
    num_splits=1,
    with_shift=False,
    h=None,
    w=None,
    attn_mask=None,
    num_head=1,
):
    """Windowed multi-head scaled dot-product attention.

    Args:
        q: [B, L, C]
        k: [B, S, C]
        v: [B, S, C]
    Returns:
        out: [B, L, C]
    """

    assert h is not None and w is not None
    assert q.size(1) == h * w

    b, _, c = q.size()

    b_new = b * num_splits * num_splits

    window_size_h = h // num_splits
    window_size_w = w // num_splits

    q = q.view(b, h, w, c)  # [B, H, W, C]
    k = k.view(b, h, w, c)
    v = v.view(b, h, w, c)

    assert c % num_head == 0

    if with_shift:
        assert attn_mask is not None  # compute once
        shift_size_h = window_size_h // 2
        shift_size_w = window_size_w // 2

        q = torch.roll(q, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
        k = torch.roll(k, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))
        v = torch.roll(v, shifts=(-shift_size_h, -shift_size_w), dims=(1, 2))

    q = split_feature(q, num_splits=num_splits)  # [B*K*K, H/K, W/K, C]
    k = split_feature(k, num_splits=num_splits)
    v = split_feature(v, num_splits=num_splits)

    # Split channels into attention heads for PyTorch SDPA.
    head_dim = c // num_head
    q = q.view(b_new, -1, num_head, head_dim).transpose(1, 2)
    k = k.view(b_new, -1, num_head, head_dim).transpose(1, 2)
    v = v.view(b_new, -1, num_head, head_dim).transpose(1, 2)

    attn_bias = None
    if with_shift:
        attn_bias = attn_mask.unsqueeze(1).repeat(b, num_head, 1, 1).to(q.dtype)

    out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias)

    out = merge_splits(
        out.transpose(1, 2).reshape(b_new, h // num_splits, w // num_splits, c),
        num_splits=num_splits,
    )  # [B, H, W, C]

    # shift back
    if with_shift:
        out = torch.roll(out, shifts=(shift_size_h, shift_size_w), dims=(1, 2))

    out = out.view(b, -1, c)

    return out


class TransformerLayer(nn.Module):
    def __init__(
        self,
        d_model=256,
        nhead=1,
        attention_type="swin",
        no_ffn=False,
        ffn_dim_expansion=4,
        with_shift=False,
        add_per_view_attn=False,
        **kwargs,
    ):
        super(TransformerLayer, self).__init__()

        self.dim = d_model
        self.nhead = nhead
        self.attention_type = attention_type
        self.no_ffn = no_ffn
        self.add_per_view_attn = add_per_view_attn

        self.with_shift = with_shift

        # Linear projections used by both single-head and multi-head attention.
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)

        self.merge = nn.Linear(d_model, d_model, bias=False)

        self.norm1 = nn.LayerNorm(d_model)

        # no ffn after self-attn, with ffn after cross-attn
        if not self.no_ffn:
            in_channels = d_model * 2
            self.mlp = nn.Sequential(
                nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False),
                nn.GELU(),
                nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False),
            )

            self.norm2 = nn.LayerNorm(d_model)

    def forward(
        self,
        source,
        target,
        height=None,
        width=None,
        shifted_window_attn_mask=None,
        attn_num_splits=None,
        **kwargs,
    ):
        if "attn_type" in kwargs:
            attn_type = kwargs["attn_type"]
        else:
            attn_type = self.attention_type

        # source: [B, L, C]
        # target: [B, L, C] for pairwise attention or [B, M, L, C] for multi-view.
        query, key, value = source, target, target

        # Project tokens before dispatching to local/global attention.
        query = self.q_proj(query)  # [B, L, C]
        key = self.k_proj(key)  # [B, L, C] or [B, M, L, C]
        value = self.v_proj(value)  # [B, L, C] or [B, M, L, C]

        if attn_type == "swin" and attn_num_splits > 1:
            if self.nhead > 1:
                message = multi_head_split_window_attention(
                    query,
                    key,
                    value,
                    num_splits=attn_num_splits,
                    with_shift=self.with_shift,
                    h=height,
                    w=width,
                    attn_mask=shifted_window_attn_mask,
                    num_head=self.nhead,
                )
            else:
                if self.add_per_view_attn:
                    assert query.dim() == 3 and key.dim() == 4 and value.dim() == 4
                    b, l, c = query.size()
                    query = query.unsqueeze(1).expand(
                        -1, key.size(1), -1, -1
                    ).reshape(-1, l, c)  # [B*M, L, C]
                    key = key.reshape(-1, l, c)
                    value = value.reshape(-1, l, c)
                    message = single_head_split_window_attention(
                        query,
                        key,
                        value,
                        num_splits=attn_num_splits,
                        with_shift=self.with_shift,
                        h=height,
                        w=width,
                        attn_mask=shifted_window_attn_mask,
                    )
                    # Sum per-target-view messages back to [B, L, C].
                    message = message.view(b, -1, l, c).sum(1)
                else:
                    message = single_head_split_window_attention(
                        query,
                        key,
                        value,
                        num_splits=attn_num_splits,
                        with_shift=self.with_shift,
                        h=height,
                        w=width,
                        attn_mask=shifted_window_attn_mask,
                    )
        else:
            message = single_head_full_attention(query, key, value)  # [B, L, C]

        message = self.merge(message)  # [B, L, C]
        message = self.norm1(message)

        if not self.no_ffn:
            message = self.mlp(torch.cat([source, message], dim=-1))
            message = self.norm2(message)

        return source + message


class TransformerBlock(nn.Module):
    """self attention + cross attention + FFN"""

    def __init__(
        self,
        d_model=256,
        nhead=1,
        attention_type="swin",
        ffn_dim_expansion=4,
        with_shift=False,
        add_per_view_attn=False,
        no_cross_attn=False,
        **kwargs,
    ):
        super(TransformerBlock, self).__init__()

        self.no_cross_attn = no_cross_attn

        if no_cross_attn:
            self.self_attn = TransformerLayer(
                d_model=d_model,
                nhead=nhead,
                attention_type=attention_type,
                ffn_dim_expansion=ffn_dim_expansion,
                with_shift=with_shift,
                add_per_view_attn=add_per_view_attn,
            )
        else:
            self.self_attn = TransformerLayer(
                d_model=d_model,
                nhead=nhead,
                attention_type=attention_type,
                no_ffn=True,
                ffn_dim_expansion=ffn_dim_expansion,
                with_shift=with_shift,
            )

            self.cross_attn_ffn = TransformerLayer(
                d_model=d_model,
                nhead=nhead,
                attention_type=attention_type,
                ffn_dim_expansion=ffn_dim_expansion,
                with_shift=with_shift,
                add_per_view_attn=add_per_view_attn,
            )

    def forward(
        self,
        source,
        target,
        height=None,
        width=None,
        shifted_window_attn_mask=None,
        attn_num_splits=None,
        **kwargs,
    ):
        # source, target: [B, L, C]
        # self attention
        source = self.self_attn(
            source,
            source,
            height=height,
            width=width,
            shifted_window_attn_mask=shifted_window_attn_mask,
            attn_num_splits=attn_num_splits,
            **kwargs,
        )

        if self.no_cross_attn:
            return source

        # cross attention and ffn
        source = self.cross_attn_ffn(
            source,
            target,
            height=height,
            width=width,
            shifted_window_attn_mask=shifted_window_attn_mask,
            attn_num_splits=attn_num_splits,
            **kwargs,
        )

        return source


def batch_features(features, nn_matrix=None):
    # Construct transformer inputs from tensor views.
    # features: [B, V, C, H, W] or [B, V, H*W, C].
    # Returns view-major query/target tensors: [V*B, ...] and [V*B, M, ...].
    assert features.dim() in (4, 5)
    batch_size, num_views = features.shape[:2]
    indices = target_view_indices(
        batch_size, num_views, nn_matrix, features.device
    )

    q = flatten_view_batch(features)
    kv = gather_view_tensor(features, indices)
    kv = flatten_view_batch(kv)

    return q, kv


class MultiViewFeatureTransformer(nn.Module):
    def __init__(
        self,
        num_layers=6,
        d_model=128,
        nhead=1,
        attention_type="swin",
        ffn_dim_expansion=4,
        add_per_view_attn=False,
        no_cross_attn=False,
        **kwargs,
    ):
        super(MultiViewFeatureTransformer, self).__init__()

        self.attention_type = attention_type

        self.d_model = d_model
        self.nhead = nhead

        self.layers = nn.ModuleList(
            [
                TransformerBlock(
                    d_model=d_model,
                    nhead=nhead,
                    attention_type=attention_type,
                    ffn_dim_expansion=ffn_dim_expansion,
                    with_shift=(
                        True if attention_type == "swin" and i % 2 == 1 else False
                    ),
                    add_per_view_attn=add_per_view_attn,
                    no_cross_attn=no_cross_attn,
                )
                for i in range(num_layers)
            ]
        )

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

        # zero init layers beyond 6
        if num_layers > 6:
            for i in range(6, num_layers):
                self.layers[i].self_attn.norm1.weight.data.zero_()
                self.layers[i].self_attn.norm1.bias.data.zero_()
                self.layers[i].cross_attn_ffn.norm2.weight.data.zero_()
                self.layers[i].cross_attn_ffn.norm2.bias.data.zero_()

    def forward(
        self,
        multi_view_features,
        attn_num_splits=None,
        **kwargs,
    ):
        nn_matrix = kwargs.pop("nn_matrix", None)

        # multi_view_features: [B, V, C, H, W].
        assert multi_view_features.dim() == 5
        b, num_views, c, h, w = multi_view_features.shape
        assert self.d_model == c

        if self.attention_type == "swin" and attn_num_splits > 1:
            # global and refine use different number of splits
            window_size_h = h // attn_num_splits
            window_size_w = w // attn_num_splits

            # compute attn mask once
            shifted_window_attn_mask = generate_shift_window_attn_mask(
                input_resolution=(h, w),
                window_size_h=window_size_h,
                window_size_w=window_size_w,
                shift_size_h=window_size_h // 2,
                shift_size_w=window_size_w // 2,
                device=multi_view_features.device,
            )  # [K*K, H/K*W/K, H/K*W/K]
        else:
            shifted_window_attn_mask = None

        # View-major query/target tensors: [V*B, C, H, W], [V*B, M, C, H, W].
        concat0, concat1 = batch_features(multi_view_features, nn_matrix=nn_matrix)
        concat0 = concat0.reshape(num_views * b, c, -1).permute(
            0, 2, 1
        )  # [V*B, H*W, C]
        c1_v = num_views - 1 if nn_matrix is None else nn_matrix.shape[-1] - 1
        concat1 = concat1.reshape(num_views * b, c1_v, c, -1).permute(
            0, 1, 3, 2
        )  # [V*B, M, H*W, C]

        for i, layer in enumerate(self.layers):
            concat0 = layer(
                concat0,
                concat1,
                height=h,
                width=w,
                shifted_window_attn_mask=shifted_window_attn_mask,
                attn_num_splits=attn_num_splits,
            )

            if i < len(self.layers) - 1:
                # Rebuild target-view groups after each layer.
                features = unflatten_view_batch(concat0, b, num_views)
                concat0, concat1 = batch_features(features, nn_matrix=nn_matrix)

        features = unflatten_view_batch(concat0, b, num_views)
        features = features.reshape(b, num_views, h, w, c)
        features = features.permute(0, 1, 4, 2, 3).contiguous()

        return features


def batch_features_camera_parameters(
    features,
    intrinsics,
    extrinsics,
    nn_matrix=None,
):
    # Construct batch-major inputs for plane-sweep warping.
    # features: [B, V, C, H, W].
    # intrinsics/extrinsics: [B, V, 3, 3] and [B, V, 4, 4].
    # Returns [B*V, ...] query tensors and [B*V, M, ...] target tensors.
    assert features.dim() == 5
    assert intrinsics.shape[-2:] == (3, 3)
    assert extrinsics.shape[-2:] == (4, 4)

    batch_size, num_views, c, h, w = features.shape
    indices = target_view_indices(
        batch_size, num_views, nn_matrix, features.device
    )
    num_selected_views = indices.size(-1)

    q = features
    q_intrinsics = intrinsics
    q_extrinsics = extrinsics
    kv = gather_view_tensor(features, indices)
    kv_intrinsics = gather_view_tensor(intrinsics, indices)
    kv_extrinsics = gather_view_tensor(extrinsics, indices)

    q = flatten_batch_view(q)  # [BV, C, H, W]
    q_intrinsics = flatten_batch_view(q_intrinsics)  # [BV, 3, 3]
    q_extrinsics = flatten_batch_view(q_extrinsics)  # [BV, 4, 4]
    kv = kv.reshape(-1, num_selected_views, c, h, w)  # [BV, V-1, C, H, W]
    kv_intrinsics = kv_intrinsics.reshape(-1, num_selected_views, 3, 3)
    kv_extrinsics = kv_extrinsics.reshape(-1, num_selected_views, 4, 4)

    return q, q_intrinsics, q_extrinsics, kv, kv_intrinsics, kv_extrinsics