"""Bottleneck of the renderer (Sec. 3.2, Appendix C.1): defocus-bias attention, defocus-aware attention blocks (DAABs)
and defocus residual groups (DRGs). Tokens are channels-last (b, h, w, c)."""
import torch
import torch.nn.functional as F
from timm.layers import DropPath
from torch import nn


class DWConv2d(nn.Module):
    """Depthwise convolution on channels-last (b, h, w, c) tensors."""

    def __init__(self, dim, kernel_size, padding):
        super().__init__()
        self.conv = nn.Conv2d(dim, dim, kernel_size, 1, padding, groups=dim)

    def forward(self, x):
        return self.conv(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)


class FeedForwardNetwork(nn.Module):
    def __init__(self, embed_dim, ffn_dim):
        super().__init__()
        self.fc1 = nn.Linear(embed_dim, ffn_dim)
        self.fc2 = nn.Linear(ffn_dim, embed_dim)

    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))


class DefocusBiasAttention(nn.Module):
    """Defocus-bias attention (Sec. 3.2, Appendix C.1), axial: along each row, then along each column.

    A_ij = softmax_j(q_i^T k_j / sqrt(d) + gamma_b * alpha * B_ij),  B_ij = -|T_D,i - T_D,j|,  T_D = phi_D(M),
    where phi_D maps the defocus magnitude M to fixed log features at several scales and then, with a learned linear
    layer (defocus_head_proj), to one value per head; alpha (defocus_alpha) is learned per head and
    gamma_b = exp(clip(w_b K + c_b, -4, 4)) (aperture_scale_proj); gamma_K = gamma_b * alpha in the main paper.
    The output is W_o(sum_j A_ij v_j + LePE(V)_i), with LePE a depthwise 5x5 convolution of the values."""

    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.key_dim = embed_dim // num_heads                                 # d
        self.scaling = self.key_dim ** -0.5
        self.q_proj = nn.Linear(embed_dim, embed_dim)                         # W_q
        self.k_proj = nn.Linear(embed_dim, embed_dim)                         # W_k
        self.v_proj = nn.Linear(embed_dim, embed_dim)                         # W_v
        self.lepe = DWConv2d(embed_dim, 5, 2)
        self.out_proj = nn.Linear(embed_dim, embed_dim)                       # W_o
        self.num_log_features = max(embed_dim // 4, 1)
        self.defocus_head_proj = nn.Linear(self.num_log_features, num_heads)  # linear layer of phi_D
        self.defocus_alpha = nn.Parameter(torch.ones(num_heads))              # alpha
        self.aperture_scale_proj = nn.Linear(1, num_heads)                    # w_b, c_b

    def _log_features(self, defocus):
        """M (b, 1, h, w) -> (b, h, w, n) features log(1 + M / s) / log(1 + 1 / s), scales s from 1/n to 1."""
        n = self.num_log_features
        min_scale = 1.0 / max(float(n), 16.0)
        scales = torch.logspace(torch.log10(defocus.new_tensor(min_scale)), torch.log10(defocus.new_tensor(1.0)),
                                steps=n, dtype=defocus.dtype, device=defocus.device).view(1, n, 1, 1)
        features = torch.log1p(defocus.clamp(0.0, 1.0) / scales) / torch.log1p(1.0 / scales)
        return features.permute(0, 2, 3, 1).contiguous()

    def _defocus_bias(self, defocus):
        """B along rows (b, h, heads, w, w) and along columns (b, w, heads, h, h)."""
        t_d = self.defocus_head_proj(self._log_features(defocus))           # T_D = phi_D(M): b, h, w, heads
        t_rows = t_d.permute(0, 1, 3, 2).contiguous()                         # b, h, heads, w
        t_cols = t_d.permute(0, 2, 3, 1).contiguous()                         # b, w, heads, h
        return (-torch.abs(t_rows.unsqueeze(-1) - t_rows.unsqueeze(-2)),
                -torch.abs(t_cols.unsqueeze(-1) - t_cols.unsqueeze(-2)))

    def _gamma_b(self, x, strength):
        raw = self.aperture_scale_proj(strength.to(device=x.device, dtype=x.dtype))
        return torch.exp(torch.clamp(raw, min=-4.0, max=4.0)).view(x.size(0), 1, self.num_heads, 1, 1)

    def forward(self, x, defocus, strength):
        b, h, w, _ = x.size()
        q = self.q_proj(x)
        k = self.k_proj(x) * self.scaling
        v = self.v_proj(x)
        lepe = self.lepe(v)
        q = q.view(b, h, w, self.num_heads, self.key_dim).permute(0, 3, 1, 2, 4)   # b, heads, h, w, d
        k = k.view(b, h, w, self.num_heads, self.key_dim).permute(0, 3, 1, 2, 4)

        bias_rows, bias_cols = self._defocus_bias(defocus)
        scale = self._gamma_b(x, strength) * self.defocus_alpha.view(1, 1, self.num_heads, 1, 1)   # gamma_b * alpha

        v = v.reshape(b, h, w, self.num_heads, -1).permute(0, 1, 3, 2, 4)            # b, h, heads, w, d
        attn = torch.softmax(q.transpose(1, 2) @ k.transpose(1, 2).transpose(-1, -2) + scale * bias_rows, dim=-1)
        v = (attn @ v).permute(0, 3, 2, 1, 4)                                         # b, w, heads, h, d
        attn = torch.softmax(q.permute(0, 3, 1, 2, 4) @ k.permute(0, 3, 1, 2, 4).transpose(-1, -2)
                             + scale * bias_cols, dim=-1)
        out = (attn @ v).permute(0, 3, 1, 2, 4).flatten(-2, -1)                       # b, h, w, c
        return self.out_proj(out + lepe)


class DefocusAwareAttentionBlock(nn.Module):
    """Defocus-aware attention block (DAAB, Fig. 3C): depthwise positional convolution, defocus-bias attention, FFN."""

    def __init__(self, embed_dim, num_heads, ffn_dim, drop_path=0.0):
        super().__init__()
        self.pos = DWConv2d(embed_dim, 3, 1)
        self.attention_layer_norm = nn.LayerNorm(embed_dim, eps=1e-6)
        self.attention = DefocusBiasAttention(embed_dim, num_heads)
        self.final_layer_norm = nn.LayerNorm(embed_dim, eps=1e-6)
        self.ffn = FeedForwardNetwork(embed_dim, ffn_dim)
        self.drop_path = DropPath(drop_path)

    def forward(self, x, defocus, strength):
        x = x + self.pos(x)
        x = x + self.drop_path(self.attention(self.attention_layer_norm(x), defocus, strength))
        return x + self.drop_path(self.ffn(self.final_layer_norm(x)))


class DAABStack(nn.Module):
    """The N DAABs of a DRG; the defocus magnitude M is resized to the token grid once for all of them."""

    def __init__(self, embed_dim, depth, num_heads, ffn_dim, drop_path):
        super().__init__()
        self.blocks = nn.ModuleList(
            [DefocusAwareAttentionBlock(embed_dim, num_heads, ffn_dim, drop_path[i]) for i in range(depth)])

    def forward(self, x, defocus, strength):
        defocus = F.interpolate(defocus.to(x.dtype), size=x.shape[1:3], mode="bilinear", align_corners=False)
        for daab in self.blocks:
            x = daab(x, defocus, strength)
        return x


class DefocusResidualGroup(nn.Module):
    """Defocus residual group (DRG, Fig. 3B): N DAABs, then a 3x3 convolution over the features and M_C, with a
    residual connection."""

    def __init__(self, embed_dim, cond_channels, depth, num_heads, ffn_dim, drop_path):
        super().__init__()
        self.residual_group = DAABStack(embed_dim, depth, num_heads, ffn_dim, drop_path)
        self.conv = nn.Conv2d(embed_dim + cond_channels, embed_dim, 3, 1, 1)

    def forward(self, x, cond, defocus, strength):
        y = self.residual_group(x, defocus, strength).permute(0, 3, 1, 2)
        return self.conv(torch.cat((y, cond), dim=1)).permute(0, 2, 3, 1) + x