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


class Model(nn.Module):
    """
    Invariant Point Attention (IPA) from AlphaFold2's structure module.
    Combines sequence-space attention with 3-D geometric queries/keys/values
    expressed in local residue frames.  The attention logits incorporate
    Euclidean distances in global frame, making it SE(3)-equivariant.
    """

    def __init__(self, node_dim, pair_dim, num_heads, num_query_points, num_value_points):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = node_dim // num_heads
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        assert node_dim % num_heads == 0

        d = self.head_dim
        self.q_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.k_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.v_proj = nn.Linear(node_dim, num_heads * d, bias=False)

        self.q_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.k_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.v_pts = nn.Linear(node_dim, num_heads * num_value_points * 3, bias=False)

        self.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)

        self.head_weights = nn.Parameter(torch.zeros(num_heads))

        out_features = num_heads * d + num_heads * num_value_points * 3
        self.out_proj = nn.Linear(out_features, node_dim)

        self.layer_norm = nn.LayerNorm(node_dim)

    def _apply_frame(self, pts, rotations, translations):
        """Transform local-frame points to global frame."""
        # pts: (B, N, P, 3), rotations: (B, N, 3, 3), translations: (B, N, 3)
        global_pts = torch.einsum("bnpc,bncd->bnpd", pts, rotations) + translations.unsqueeze(2)
        return global_pts

    def _apply_inv_frame(self, pts, rotations, translations):
        """Transform global-frame points to local frame."""
        centered = pts - translations.unsqueeze(2)
        local_pts = torch.einsum("bnpc,bndc->bnpd", centered, rotations)
        return local_pts

    def forward(
        self,
        node_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        translations: torch.Tensor,
        rotations: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            node_repr:    (B, N, node_dim)
            pair_repr:    (B, N, N, pair_dim)
            translations: (B, N, 3)
            rotations:    (B, N, 3, 3)
        Returns:
            updated node representation (B, N, node_dim)
        """
        B, N, D = node_repr.shape
        h = self.num_heads
        d = self.head_dim
        Qp = self.num_query_points
        Vp = self.num_value_points

        x = self.layer_norm(node_repr)

        q = self.q_proj(x).view(B, N, h, d)
        k = self.k_proj(x).view(B, N, h, d)
        v = self.v_proj(x).view(B, N, h, d)

        # Points in local frame, then transform to global
        q_pts_local = self.q_pts(x).view(B, N, h * Qp, 3)
        k_pts_local = self.k_pts(x).view(B, N, h * Qp, 3)
        v_pts_local = self.v_pts(x).view(B, N, h * Vp, 3)

        q_pts = self._apply_frame(q_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        k_pts = self._apply_frame(k_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        v_pts = self._apply_frame(v_pts_local, rotations, translations).view(B, N, h, Vp, 3)

        # Scalar attention logits
        attn_scalar = torch.einsum("bihd,bjhd->bhij", q, k) / math.sqrt(d)

        # Point attention: negative squared distance between query and key points
        # q_pts: (B, N_i, h, Qp, 3), k_pts: (B, N_j, h, Qp, 3)
        q_expand = q_pts.unsqueeze(2)  # (B, N, 1, h, Qp, 3)
        k_expand = k_pts.unsqueeze(1)  # (B, 1, N, h, Qp, 3)
        pt_dist = ((q_expand - k_expand) ** 2).sum(dim=(-1, -2))  # (B, N, N, h)

        w = F.softplus(self.head_weights)
        pt_attn = -(w.view(1, 1, 1, h) * pt_dist) / 2.0
        pt_attn = pt_attn.permute(0, 3, 1, 2)  # (B, h, N, N)

        # Pair bias
        pair_bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)  # (B, h, N, N)

        attn = attn_scalar + pt_attn + pair_bias
        attn = F.softmax(attn, dim=-1)  # (B, h, N_i, N_j)

        # Aggregate scalar values
        out_scalar = torch.einsum("bhij,bjhd->bihd", attn, v).reshape(B, N, h * d)

        # Aggregate point values in global frame, then transform back to local
        out_pts = torch.einsum("bhij,bjhpc->bihpc", attn, v_pts)  # (B, N, h, Vp, 3)
        out_pts = out_pts.reshape(B, N, h * Vp, 3)
        out_pts_local = self._apply_inv_frame(out_pts, rotations, translations)
        out_pts_flat = out_pts_local.reshape(B, N, h * Vp * 3)

        out = torch.cat([out_scalar, out_pts_flat], dim=-1)
        return self.out_proj(out)


node_dim = 64
pair_dim = 32
num_heads = 4
num_query_points = 4
num_value_points = 8
seq_len = 32
batch_size = 2


def get_inputs():
    p = popcorn_pri
    node_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), node_dim)
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    translations = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), 3)
    rotations = torch.zeros(p.jitter_int(batch_size), p.jitter_int(seq_len), 3, 3)
    for b in range(p.jitter_int(batch_size)):
        for n in range(p.jitter_int(seq_len)):
            rotations[b, n] = torch.eye(3)
    return [node_repr, pair_repr, translations, rotations]


def get_init_inputs():
    return [node_dim, pair_dim, num_heads, num_query_points, num_value_points]