"""Particle Transformer (ParT)

Paper: "Particle Transformer for Jet Tagging" - https://arxiv.org/abs/2202.03772

We have to do three things to build a LLoCa-ParT
- Construct a LLoCaAttention module for the whole transformer that preprocesses the frames
  and is passed to each attention block during initialization.
- When evaluating attention, use the LLoCaAttention module.

More comments:
- We also added an extra clamp in to_ptrapphim to avoid numerical issues from log(0). This case
  might not happen in the original ParT, but it can happen with LLoCa for highly boosted frames.
- We use LLoCaAttention only for the self-attention blocks, and use default attention (corresponds
  to scalar messages only) for the class attention blocks.

Use git diff --no-index experiments/baselines/particletransformer.py lloca/nn/particletransformer.py
to see the changes required to include frame-to-frame transformations
"""

import math
import random
import copy
from functools import partial
from typing import Optional, Tuple, Any, Callable

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

from ..reps.tensorreps import TensorReps
from .attention import LLoCaAttention


@torch.jit.script
def delta_phi(a, b):
    return (a - b + math.pi) % (2 * math.pi) - math.pi


@torch.jit.script
def delta_r2(eta1, phi1, eta2, phi2):
    return (eta1 - eta2) ** 2 + delta_phi(phi1, phi2) ** 2


def to_pt2(x, eps=1e-8):
    pt2 = x[:, :2].square().sum(dim=1, keepdim=True)
    if eps is not None:
        pt2 = pt2.clamp(min=eps)
    return pt2


def to_m2(x, eps=1e-8):
    m2 = x[:, 3:4].square() - x[:, :3].square().sum(dim=1, keepdim=True)
    if eps is not None:
        m2 = m2.clamp(min=eps)
    return m2


def to_ptrapphim(x, return_mass=True, eps=1e-8):
    # x: (N, 4, ...), dim1 : (px, py, pz, E)
    px, py, pz, energy = x.split((1, 1, 1, 1), dim=1)
    pt = torch.sqrt(to_pt2(x, eps=eps))
    # rapidity = 0.5 * torch.log((energy + pz) / (energy - pz))
    rapidity = 0.5 * torch.log(
        (1 + (2 * pz) / (energy - pz).clamp(min=1e-20)).clamp(min=1e-20)
    )
    phi = torch.atan2(py, px)
    if not return_mass:
        return torch.cat((pt, rapidity, phi), dim=1)
    else:
        m = torch.sqrt(to_m2(x, eps=eps))
        return torch.cat((pt, rapidity, phi, m), dim=1)


def boost(x, boostp4, eps=1e-8):
    # boost x to the rest frame of boostp4
    # x: (N, 4, ...), dim1 : (px, py, pz, E)
    p3 = -boostp4[:, :3] / boostp4[:, 3:].clamp(min=eps)
    b2 = p3.square().sum(dim=1, keepdim=True)
    gamma = (1 - b2).clamp(min=eps) ** (-0.5)
    gamma2 = (gamma - 1) / b2
    gamma2.masked_fill_(b2 == 0, 0)
    bp = (x[:, :3] * p3).sum(dim=1, keepdim=True)
    v = x[:, :3] + gamma2 * bp * p3 + x[:, 3:] * gamma * p3
    return v


def p3_norm(p, eps=1e-8):
    return p[:, :3] / p[:, :3].norm(dim=1, keepdim=True).clamp(min=eps)


def to_energy_momentum(x, return_unit_vector=True):
    energy = x[:, 3:4]
    p2 = x[:, :3].square().sum(dim=1, keepdim=True)
    mom = torch.sqrt(p2)
    if return_unit_vector:
        return energy, mom, x[:, :3] / mom.clamp(min=1e-8)
    else:
        return energy, mom


def to_cos_sin_angles(xi, xj, normed_inputs=False, eps=1e-8):
    if normed_inputs:
        ni, nj = xi, xj
    else:
        ni, nj = p3_norm(xi, eps), p3_norm(xj, eps)
    cos = (ni * nj).sum(dim=1, keepdim=True).clamp(min=-1, max=1)
    sin = (
        torch.linalg.cross(ni, nj, dim=1).norm(dim=1, keepdim=True).clamp(min=0, max=1)
    )
    return cos, sin


def pairwise_lv_fts_pp(xi, xj, num_outputs=4, eps=1e-8):
    pti, rapi, phii = to_ptrapphim(xi, False, eps=None).split((1, 1, 1), dim=1)
    ptj, rapj, phij = to_ptrapphim(xj, False, eps=None).split((1, 1, 1), dim=1)

    # modified this for convenience (only lorentz scalars is most conservative)
    xij = xi + xj
    lnm2 = torch.log(to_m2(xij, eps=eps))
    if num_outputs > 0:
        outputs = [lnm2]

    if num_outputs > 1:
        delta = delta_r2(rapi, phii, rapj, phij).sqrt()
        lndelta = torch.log(delta.clamp(min=eps))
        ptmin = torch.minimum(pti, ptj)
        lnkt = torch.log((ptmin * delta).clamp(min=eps))
        lnz = torch.log((ptmin / (pti + ptj).clamp(min=eps)).clamp(min=eps))
        outputs += [lnkt, lnz, lndelta]

    if num_outputs > 4:
        lnds2 = torch.log(torch.clamp(-to_m2(xi - xj, eps=None), min=eps))
        outputs.append(lnds2)

    # the following features are not symmetric for (i, j)
    if num_outputs > 5:
        xj_boost = boost(xj, xij)
        costheta = (p3_norm(xj_boost, eps=eps) * p3_norm(xij, eps=eps)).sum(
            dim=1, keepdim=True
        )
        outputs.append(costheta)

    if num_outputs > 6:
        deltarap = rapi - rapj
        deltaphi = delta_phi(phii, phij)
        outputs += [deltarap, deltaphi]

    assert len(outputs) == num_outputs
    return torch.cat(outputs, dim=1)


def pairwise_lv_fts_ee(xi, xj, num_outputs=6, eps=1e-8):
    # outputs: [lnm2, cos_angle, sin_angle, lnkt, lnz, lnjade]
    lnm2 = torch.log(to_m2(xi + xj, eps=eps))
    outputs = [lnm2]

    if num_outputs > 1:
        ei, pi, ni = to_energy_momentum(xi)
        ej, pj, nj = to_energy_momentum(xj)
        cos_angle, sin_angle = to_cos_sin_angles(ni, nj, normed_inputs=True)
        outputs += [cos_angle, sin_angle]

    if num_outputs > 3:
        pmin = torch.minimum(pi, pj)
        lnkt = torch.log((pmin * sin_angle).clamp(min=eps))
        lnz = torch.log((pmin / (pi + pj).clamp(min=eps)).clamp(min=eps))
        outputs += [lnkt, lnz]

    if num_outputs > 5:
        lnjade = torch.log((ei * ej * (1 - cos_angle)).clamp(min=eps))
        outputs.append(lnjade)

    assert len(outputs) == num_outputs
    return torch.cat(outputs, dim=1)


def build_sparse_tensor(uu, idx, seq_len):
    # inputs: uu (N, C, num_pairs), idx (N, 2, num_pairs)
    # return: (N, C, seq_len, seq_len)
    batch_size, num_fts, num_pairs = uu.size()
    idx = torch.min(idx, torch.ones_like(idx) * seq_len)
    i = torch.cat(
        (
            torch.arange(0, batch_size, device=uu.device)
            .repeat_interleave(num_fts * num_pairs)
            .unsqueeze(0),
            torch.arange(0, num_fts, device=uu.device)
            .repeat_interleave(num_pairs)
            .repeat(batch_size)
            .unsqueeze(0),
            idx[:, :1, :].expand_as(uu).flatten().unsqueeze(0),
            idx[:, 1:, :].expand_as(uu).flatten().unsqueeze(0),
        ),
        dim=0,
    )
    return torch.sparse_coo_tensor(
        i,
        uu.flatten(),
        size=(batch_size, num_fts, seq_len + 1, seq_len + 1),
        device=uu.device,
    ).to_dense()[:, :, :seq_len, :seq_len]


def tril_indices(row, col, offset=0, *, dtype=torch.long, device="cpu"):
    return torch.ones(row, col, dtype=dtype, device=device).tril(offset).nonzero().T


class SequenceTrimmer(nn.Module):
    def __init__(
        self, enabled=False, target=(0.9, 1.02), warmup_steps=5, **kwargs
    ) -> None:
        super().__init__(**kwargs)
        self.enabled = enabled
        self.target = target
        self.warmup_steps = warmup_steps
        self.register_buffer("_counter", torch.LongTensor([0]), persistent=False)

    def forward(self, x, v=None, mask=None, uu=None):
        # x: (N, C, P)
        # v: (N, 4, P) [px,py,pz,energy]
        # mask: (N, 1, P) -- real particle = 1, padded = 0
        # uu: (N, C', P, P)
        if mask is None:
            mask = torch.ones_like(x[:, :1])
        mask = mask.bool()

        if self.enabled:
            if self._counter < self.warmup_steps:
                self._counter.add_(1)
            else:
                if v is not None:
                    if not isinstance(v, (list, tuple)):
                        v = [v]
                if self.training:
                    q = min(1, random.uniform(*self.target))
                    maxlen = torch.quantile(mask.float().sum(dim=-1), q).long()
                    rand = torch.rand_like(mask.float())
                    rand.masked_fill_(~mask, -1)
                    perm = rand.argsort(dim=-1, descending=True)  # (N, 1, P)
                    mask = torch.gather(mask, -1, perm)
                    x = torch.gather(x, -1, perm.expand_as(x))
                    if v is not None:
                        v = [torch.gather(_v, -1, perm.expand_as(_v)) for _v in v]
                    if uu is not None:
                        uu = torch.gather(uu, -2, perm.unsqueeze(-1).expand_as(uu))
                        uu = torch.gather(uu, -1, perm.unsqueeze(-2).expand_as(uu))
                else:
                    maxlen = mask.sum(dim=-1).max()
                maxlen = max(maxlen, 1)
                if maxlen < mask.size(-1):
                    mask = mask[:, :, :maxlen]
                    x = x[:, :, :maxlen]
                    if v is not None:
                        v = [_v[:, :, :maxlen] for _v in v]
                    if uu is not None:
                        uu = uu[:, :, :maxlen, :maxlen]
                if v is not None:
                    if len(v) == 1:
                        v = v[0]

        return x, v, mask, uu


class SwiGLUFFN(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: Optional[int] = None,
        out_features: Optional[int] = None,
        drop: float = 0.0,
        bias: bool = True,
    ) -> None:
        super().__init__()
        hidden_features = hidden_features or in_features
        out_features = out_features or in_features
        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
        self.drop = nn.Dropout(drop)
        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x12 = self.w12(x)
        x1, x2 = x12.chunk(2, dim=-1)
        hidden = F.silu(x1) * x2
        hidden = self.drop(hidden)
        return self.w3(hidden)


class Embed(nn.Module):
    def __init__(self, input_dim, dims, normalize_input=True, activation="gelu"):
        super().__init__()

        self.input_bn = nn.BatchNorm1d(input_dim) if normalize_input else None
        module_list = []
        for dim in dims:
            module_list.extend(
                [
                    nn.LayerNorm(input_dim),
                    nn.Linear(input_dim, dim),
                    nn.GELU() if activation == "gelu" else nn.ReLU(),
                ]
            )
            input_dim = dim
        self.embed = nn.Sequential(*module_list)

    def forward(self, x):
        if self.input_bn is not None:
            # x: (batch, embed_dim, seq_len)
            x = self.input_bn(x)
            x = x.transpose(1, 2).contiguous()
        # x: (batch, seq_len, embed_dim)
        return self.embed(x)


class PairEmbed(nn.Module):
    def __init__(
        self,
        pairwise_lv_dim,
        pairwise_input_dim,
        dims,
        pairwise_lv_type="pp",
        remove_self_pair=False,
        use_pre_activation_pair=True,
        normalize_input=True,
        activation="gelu",
        eps=1e-8,
        for_onnx=False,
        sparse_eval=None,
    ):
        super().__init__()

        self.pairwise_lv_dim = pairwise_lv_dim
        self.pairwise_input_dim = pairwise_input_dim
        self.remove_self_pair = remove_self_pair
        self.for_onnx = for_onnx
        self.sparse_eval = (not for_onnx) if sparse_eval is None else sparse_eval
        self.out_dim = dims[-1]

        if pairwise_lv_type == "pp":
            self.is_symmetric = (pairwise_lv_dim <= 5) and (pairwise_input_dim == 0)
            self.pairwise_lv_fts = partial(
                pairwise_lv_fts_pp, num_outputs=pairwise_lv_dim, eps=eps
            )
        elif pairwise_lv_type == "ee":
            self.is_symmetric = (pairwise_lv_dim <= 6) and (pairwise_input_dim == 0)
            self.pairwise_lv_fts = partial(
                pairwise_lv_fts_ee, num_outputs=pairwise_lv_dim, eps=eps
            )
        else:
            raise RuntimeError(
                "Invalid value for `pairwise_lv_type`: " + pairwise_lv_type
            )

        if pairwise_lv_dim > 0:
            input_dim = pairwise_lv_dim
            module_list = [nn.BatchNorm1d(input_dim)] if normalize_input else []
            for dim in dims:
                module_list.extend(
                    [
                        nn.Conv1d(input_dim, dim, 1),
                        nn.BatchNorm1d(dim),
                        nn.GELU() if activation == "gelu" else nn.ReLU(),
                    ]
                )
                input_dim = dim
            if use_pre_activation_pair:
                module_list = module_list[:-1]
            self.embed = nn.Sequential(*module_list)

        if pairwise_input_dim > 0:
            input_dim = pairwise_input_dim
            module_list = [nn.BatchNorm1d(input_dim)] if normalize_input else []
            for dim in dims:
                module_list.extend(
                    [
                        nn.Conv1d(input_dim, dim, 1),
                        nn.BatchNorm1d(dim),
                        nn.GELU() if activation == "gelu" else nn.ReLU(),
                    ]
                )
                input_dim = dim
            if use_pre_activation_pair:
                module_list = module_list[:-1]
            self.fts_embed = nn.Sequential(*module_list)

    def _forward_dense(self, x, uu=None, mask=None):
        # x: (batch, v_dim, seq_len)
        # uu: (batch, v_dim, seq_len, seq_len)
        assert x is not None or uu is not None
        with torch.no_grad():
            if x is not None:
                batch_size, _, seq_len = x.size()
            else:
                batch_size, _, seq_len, _ = uu.size()
            if self.is_symmetric:
                tril_indices_fn = tril_indices if self.for_onnx else torch.tril_indices
                i, j = tril_indices_fn(
                    seq_len,
                    seq_len,
                    offset=-1 if self.remove_self_pair else 0,
                    device=(x if x is not None else uu).device,
                )
                if x is not None:
                    x = x.unsqueeze(-1).repeat(1, 1, 1, seq_len)
                    xi = x[:, :, i, j]  # (batch, dim, seq_len*(seq_len+1)/2)
                    xj = x[:, :, j, i]
                    x = self.pairwise_lv_fts(xi, xj)
                if uu is not None:
                    # (batch, dim, seq_len*(seq_len+1)/2)
                    uu = uu[:, :, i, j]
            else:
                if x is not None:
                    x = self.pairwise_lv_fts(x.unsqueeze(-1), x.unsqueeze(-2))
                    if self.remove_self_pair:
                        i = torch.arange(0, seq_len, device=x.device)
                        x[:, :, i, i] = 0
                    x = x.view(-1, self.pairwise_lv_dim, seq_len * seq_len)
                if uu is not None:
                    uu = uu.view(-1, self.pairwise_input_dim, seq_len * seq_len)

        # with grad
        elements = 0
        if x is not None:
            elements = elements + self.embed(x)
        if uu is not None:
            elements = elements + self.fts_embed(uu)

        if self.is_symmetric:
            y = torch.zeros(
                batch_size,
                self.out_dim,
                seq_len,
                seq_len,
                dtype=elements.dtype,
                device=elements.device,
            )
            y[:, :, i, j] = elements
            y[:, :, j, i] = elements
        else:
            y = elements.view(-1, self.out_dim, seq_len, seq_len)
        return y

    def _forward_sparse(self, x, uu=None, mask=None):
        # x: (batch, v_dim, seq_len)
        # uu: (batch, v_dim, seq_len, seq_len)
        assert x is not None or uu is not None
        with torch.no_grad():
            if x is not None:
                batch_size, _, seq_len = x.size()
            else:
                batch_size, _, seq_len, _ = uu.size()

            i0, i1, i2, i3 = (Ellipsis,) * 4
            if mask is not None:
                mask = mask.unsqueeze(-1) * mask.unsqueeze(
                    -2
                )  # (batch_size, 1, seq_len, seq_len)
                if self.is_symmetric:
                    offset = -1 if self.remove_self_pair else 0
                    i0, _, i2, i3 = mask.float().tril(offset).nonzero(as_tuple=True)
                else:
                    i0, _, i2, i3 = mask.nonzero(as_tuple=True)

            if x is not None:
                x = self.pairwise_lv_fts(x.unsqueeze(-1), x.unsqueeze(-2))
                x = x.permute(0, 2, 3, 1)[
                    i0, i2, i3, :
                ]  # (num_elements, pairwise_lv_dim)
                x = x.T.unsqueeze(0).contiguous()  # (1, pairwise_lv_dim, num_elements)
            if uu is not None:
                uu = uu.permute(0, 2, 3, 1)[
                    i0, i2, i3, :
                ]  # (num_elements, pairwise_input_dim)
                uu = uu.T.unsqueeze(
                    0
                ).contiguous()  # (1, pairwise_input_dim, num_elements)

        # with grad
        elements = 0
        if x is not None:
            elements = elements + self.embed(x)
        if uu is not None:
            elements = elements + self.fts_embed(uu)
        elements = elements.squeeze(0).T  # (num_elements, out_dim)

        y = torch.zeros(
            batch_size,
            seq_len,
            seq_len,
            self.out_dim,
            dtype=elements.dtype,
            device=elements.device,
        )
        y[i0, i2, i3, :] = elements
        if self.is_symmetric:
            y[i0, i3, i2, :] = elements
        y = y.permute(0, 3, 1, 2).contiguous()

        return y

    def forward(self, x, uu=None, mask=None):
        if self.sparse_eval:
            return self._forward_sparse(x, uu=uu, mask=mask)
        else:
            return self._forward_dense(x, uu=uu, mask=mask)


def _canonical_mask(
    mask: Optional[torch.Tensor],
    mask_name: str,
    other_type: Optional[Any],
    other_name: str,
    target_type: Any,
    check_other: bool = True,
) -> Optional[torch.Tensor]:

    if mask is not None:
        _mask_dtype = mask.dtype
        _mask_is_float = torch.is_floating_point(mask)
        if _mask_dtype != torch.bool and not _mask_is_float:
            raise AssertionError(
                f"only bool and floating types of {mask_name} are supported"
            )
        if not _mask_is_float:
            mask = torch.zeros_like(mask, dtype=target_type).masked_fill_(
                mask, float("-inf")
            )
    return mask


def _none_or_dtype(input: Optional[torch.Tensor]):
    if input is None:
        return None
    elif isinstance(input, torch.Tensor):
        return input.dtype
    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")


class Attention(torch.nn.Module):
    def __init__(
        self,
        attention,
        embed_dim,
        num_heads,
        dropout=0.0,
        bias=True,
        device=None,
        dtype=None,
    ) -> None:
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert (
            self.head_dim * num_heads == self.embed_dim
        ), "embed_dim must be divisible by num_heads"

        self.in_proj = torch.nn.Linear(
            embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
        )
        self.out_proj = torch.nn.Linear(
            embed_dim, embed_dim, bias=bias, **factory_kwargs
        )
        self.attention = attention

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):

        for k in state_dict.keys():
            if k.endswith("in_proj_weight"):
                state_dict[k.replace("_weight", ".weight")] = state_dict.pop(k)
            elif k.endswith("in_proj_bias"):
                state_dict[k.replace("_bias", ".bias")] = state_dict.pop(k)

        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        key_padding_mask: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

        bsz, tgt_len, _ = query.shape
        _, src_len, _ = key.shape

        # (bsz, src_len)
        key_padding_mask = _canonical_mask(
            mask=key_padding_mask,
            mask_name="key_padding_mask",
            other_type=_none_or_dtype(attn_mask),
            other_name="attn_mask",
            target_type=query.dtype,
        )

        # (bsz, num_heads, tgt_len, src_len)
        attn_mask = _canonical_mask(
            mask=attn_mask,
            mask_name="attn_mask",
            other_type=None,
            other_name="",
            target_type=query.dtype,
            check_other=False,
        )

        # merge key padding and attention masks
        if key_padding_mask is not None:
            assert key_padding_mask.shape == (
                bsz,
                src_len,
            ), f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
            key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).expand(
                -1, self.num_heads, -1, -1
            )
            if attn_mask is None:
                attn_mask = key_padding_mask
            else:
                assert attn_mask.shape == (
                    bsz,
                    self.num_heads,
                    tgt_len,
                    src_len,
                ), f"expecting attn_mask shape of {(bsz, self.num_heads, tgt_len, src_len)}, but got {attn_mask.shape}"
                attn_mask = attn_mask + key_padding_mask

        # (bsz, seq_len, num_heads*head_dim)
        q, k, v = F._in_projection_packed(
            query, key, value, self.in_proj.weight, self.in_proj.bias
        )

        # -> (bsz, num_heads, src/tgt_len, head_dim)
        q = (
            q.view(bsz, tgt_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        k = (
            k.view(bsz, src_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )
        v = (
            v.view(bsz, src_len, self.num_heads, self.head_dim)
            .transpose(1, 2)
            .contiguous()
        )

        dropout_p = self.dropout if self.training else 0.0

        if self.attention is not None:
            # particle attention
            attn_output = self.attention(
                q,
                k,
                v,
                attn_mask=attn_mask,
                dropout_p=dropout_p,
            )
        else:
            # class token attention
            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p)

        attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.embed_dim)
        attn_output = self.out_proj(attn_output)
        return attn_output, None


class LayerScale(nn.Module):
    def __init__(
        self,
        dim: int,
        init_values: float = 1e-5,
        inplace: bool = False,
    ) -> None:
        super().__init__()
        self.inplace = inplace
        self.gamma = nn.Parameter(init_values * torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.mul_(self.gamma) if self.inplace else x * self.gamma


def drop_path(
    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

    def extra_repr(self):
        return f"drop_prob={round(self.drop_prob,3):0.3f}"


class Block(nn.Module):
    def __init__(
        self,
        attention,
        embed_dim=128,
        num_heads=8,
        ffn_ratio=4,
        dropout=0.1,
        attn_dropout=0.1,
        activation_dropout=0.1,
        activation="gelu",
        layer_scale_init_values=None,
        drop_path_rate=0.0,
        scale_attn_mask=False,
        scale_attn=True,
        scale_fc=True,
        scale_heads=True,
        scale_resids=True,
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.ffn_dim = embed_dim * ffn_ratio

        self.pre_attn_norm = nn.LayerNorm(embed_dim)
        self.attn = Attention(attention, embed_dim, num_heads, dropout=attn_dropout)
        self.post_attn_norm = nn.LayerNorm(embed_dim) if scale_attn else nn.Identity()
        self.dropout = nn.Dropout(dropout)
        self.ls1 = (
            LayerScale(embed_dim, init_values=layer_scale_init_values)
            if layer_scale_init_values
            else nn.Identity()
        )
        self.drop_path1 = (
            DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
        )

        self.pre_fc_norm = nn.LayerNorm(embed_dim)
        self.fc1 = nn.Linear(embed_dim, self.ffn_dim)
        if activation == "swiglu":
            self.fc1_g = nn.Linear(embed_dim, self.ffn_dim)
            self.act = nn.SiLU()
        else:
            self.fc1_g = None
            self.act = nn.GELU() if activation == "gelu" else nn.ReLU()
        self.act_dropout = nn.Dropout(activation_dropout)
        self.post_fc_norm = nn.LayerNorm(self.ffn_dim) if scale_fc else nn.Identity()
        self.fc2 = nn.Linear(self.ffn_dim, embed_dim)
        self.ls2 = (
            LayerScale(embed_dim, init_values=layer_scale_init_values)
            if layer_scale_init_values
            else nn.Identity()
        )
        self.drop_path2 = (
            DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
        )

        self.c_mask = (
            nn.Parameter(torch.ones(1), requires_grad=True) if scale_attn_mask else None
        )
        self.c_attn = (
            nn.Parameter(torch.ones(num_heads), requires_grad=True)
            if scale_heads
            else None
        )
        self.w_resid = (
            nn.Parameter(torch.ones(embed_dim), requires_grad=True)
            if scale_resids
            else None
        )

    def forward(self, x, x_cls=None, padding_mask=None, attn_mask=None):
        """
        Args:
            x (Tensor): input to the layer of shape `(batch, seq_len, embed_dim)`
            x_cls (Tensor, optional): class token input to the layer of shape `(batch, 1, embed_dim)`
            padding_mask (ByteTensor, optional): binary
                ByteTensor of shape `(batch, seq_len)` where padding
                elements are indicated by ``True``.

        Returns:
            encoded output of shape `(batch, seq_len, embed_dim)`
        """

        if x_cls is not None:
            with torch.no_grad():
                # prepend one element for x_cls: -> (batch, 1+seq_len)
                padding_mask = torch.cat(
                    (torch.zeros_like(padding_mask[:, :1]), padding_mask), dim=1
                )
            # class attention: https://arxiv.org/pdf/2103.17239.pdf
            residual = x_cls
            u = torch.cat((x_cls, x), dim=1)  # (batch, 1+seq_len, embed_dim)
            u = self.pre_attn_norm(u)

            # default attention for convenience (could be more fancy here)
            x = self.attn(x_cls, u, u, key_padding_mask=padding_mask,)[
                0
            ]  # (1, batch, embed_dim)
        else:
            if self.c_mask is not None and attn_mask is not None:
                attn_mask = torch.mul(self.c_mask, attn_mask)
            residual = x
            x = self.pre_attn_norm(x)
            x = self.attn(x, x, x, key_padding_mask=padding_mask, attn_mask=attn_mask)[
                0
            ]  # (batch, seq_len, embed_dim)

        if self.c_attn is not None:
            bsz, tgt_len, _ = x.size()
            x = x.view(bsz, tgt_len, self.num_heads, self.head_dim)
            x = torch.einsum("bthd,h->btdh", x, self.c_attn)
            x = x.reshape(bsz, tgt_len, self.embed_dim)
        x = self.post_attn_norm(x)
        x = self.dropout(x)
        x = self.drop_path1(self.ls1(x))
        x += residual

        residual = x
        x = self.pre_fc_norm(x)
        if self.fc1_g is None:
            x = self.act(self.fc1(x))
        else:
            x_gate = self.fc1_g(x)
            x = self.fc1(x)
            x = self.act(x_gate) * x
        x = self.act_dropout(x)
        x = self.post_fc_norm(x)
        x = self.fc2(x)
        x = self.dropout(x)
        x = self.drop_path2(self.ls2(x))
        if self.w_resid is not None:
            residual = torch.mul(self.w_resid, residual)
        x += residual

        return x


class ParticleTransformer(nn.Module):
    def __init__(
        self,
        input_dim,
        attn_reps,
        num_classes=None,
        # network configurations
        pair_input_type="pp",
        pair_input_dim=None,
        pair_extra_dim=0,
        remove_self_pair=False,
        use_pre_activation_pair=True,
        embed_dims=(128, 512, 128),
        pair_embed_dims=(64, 64, 64),
        num_heads=8,
        num_layers=8,
        num_cls_layers=2,
        block_params=None,
        cls_block_params=None,
        fc_params=(),
        activation="gelu",
        # misc
        version=1,
        weight_init="moco",
        fix_init=True,
        trim=True,
        for_inference=False,
        for_segmentation=False,
        use_amp=False,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.trimmer = SequenceTrimmer(enabled=trim and not for_inference)
        self.for_inference = for_inference
        self.for_segmentation = for_segmentation
        self.use_amp = use_amp

        self.embed_dim = embed_dims[-1] if len(embed_dims) > 0 else input_dim
        attn_reps = TensorReps(attn_reps)
        assert attn_reps.dim * num_heads == self.embed_dim
        self.attention = LLoCaAttention(attn_reps, num_heads)
        default_cfg = dict(
            embed_dim=self.embed_dim,
            num_heads=num_heads,
            ffn_ratio=4,
            dropout=0.1,
            attn_dropout=0.1,
            activation_dropout=0.1,
            activation=activation,
            layer_scale_init_values=None,
            drop_path_rate=0.0,
            scale_attn_mask=False,
            scale_fc=True,
            scale_attn=True,
            scale_heads=True,
            scale_resids=True,
        )
        if version > 1:
            default_cfg.update(
                activation="swiglu",
                scale_fc=False,
                scale_attn=False,
                scale_heads=False,
                scale_resids=False,
            )

        cfg_block = copy.deepcopy(default_cfg)
        if block_params is not None:
            cfg_block.update(block_params)

        cfg_cls_block = copy.deepcopy(default_cfg)
        cfg_cls_block.update({"dropout": 0, "attn_dropout": 0, "activation_dropout": 0})
        if cls_block_params is not None:
            cfg_cls_block.update(cls_block_params)

        self.embed = (
            Embed(input_dim, embed_dims, activation=activation)
            if len(embed_dims) > 0
            else nn.Identity()
        )

        if pair_input_dim is None:
            pair_input_dim = 4 if pair_input_type == "pp" else 6
        self.pair_extra_dim = pair_extra_dim
        self.pair_embed = (
            PairEmbed(
                pair_input_dim,
                pair_extra_dim,
                (*pair_embed_dims, cfg_block["num_heads"]),
                pairwise_lv_type=pair_input_type,
                remove_self_pair=remove_self_pair,
                use_pre_activation_pair=use_pre_activation_pair,
                for_onnx=for_inference,
            )
            if pair_embed_dims is not None and pair_input_dim + pair_extra_dim > 0
            else None
        )
        self.blocks = nn.ModuleList(
            [Block(attention=self.attention, **cfg_block) for _ in range(num_layers)]
        )
        self.cls_blocks = (
            nn.ModuleList(
                [Block(attention=None, **cfg_cls_block) for _ in range(num_cls_layers)]
            )
            if num_cls_layers > 0
            else None
        )
        self.norm = nn.LayerNorm(self.embed_dim)

        if fc_params is not None:
            fcs = []
            in_dim = self.embed_dim
            for param in fc_params:
                try:
                    out_dim, drop_rate, act = param
                except ValueError:
                    (out_dim, drop_rate), act = param, "relu"
                if act == "swiglu":
                    layer = nn.Sequential(
                        SwiGLUFFN(in_dim, out_dim * 4, out_dim, drop=drop_rate),
                        nn.LayerNorm(out_dim),
                    )
                else:
                    layer = nn.Sequential(
                        nn.Linear(in_dim, out_dim),
                        nn.GELU() if act == "gelu" else nn.ReLU(),
                        nn.Dropout(drop_rate),
                    )
                fcs.append(layer)
                in_dim = out_dim
            fcs.append(nn.Linear(in_dim, num_classes))
            self.fc = nn.Sequential(*fcs)
        else:
            self.fc = None

        # cls tokens
        if not self.for_segmentation and num_cls_layers > 0:
            self.cls_token = nn.Parameter(
                torch.zeros(1, 1, self.embed_dim), requires_grad=True
            )
            nn.init.trunc_normal_(self.cls_token, std=0.02)
        else:
            self.cls_token = None

        # weight initialization
        if weight_init is not None:
            self.init_weights(weight_init)
        if fix_init:
            self.fix_init_weight()

    def fix_init_weight(self):
        def rescale(param, _layer_id):
            param.div_(math.sqrt(2.0 * _layer_id))

        for layer_id, layer in enumerate(self.blocks):
            rescale(layer.attn.out_proj.weight.data, layer_id + 1)
            rescale(layer.fc2.weight.data, layer_id + 1)

    def init_weights(self, mode: str = "") -> None:
        assert mode in ("timm", "moco")
        if mode == "timm":
            named_apply(init_weights_vit_timm, self)
        elif mode == "moco":
            named_apply(init_weights_vit_moco, self)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {
            "cls_token",
        }

    def _forward_encoder(self, x, v=None, mask=None, uu=None, uu_idx=None):
        with torch.no_grad():
            if not self.for_inference:
                if uu_idx is not None:
                    uu = build_sparse_tensor(uu, uu_idx, x.size(-1))
            x, v, mask, uu = self.trimmer(x, v, mask, uu)
            padding_mask = ~mask.squeeze(1)  # (batch_size, seq_len)

        with torch.autocast("cuda", enabled=self.use_amp):
            # input embedding
            x = self.embed(x).masked_fill(
                ~mask.transpose(1, 2), 0
            )  # (batch_size, seq_len, num_fts)
            attn_mask = None
            if (v is not None or uu is not None) and self.pair_embed is not None:
                attn_mask = self.pair_embed(
                    v, uu=uu, mask=mask
                )  # (batch_size, num_heads, seq_len, seq_len)

            # transform
            for block in self.blocks:
                x = block(x, x_cls=None, padding_mask=padding_mask, attn_mask=attn_mask)

        # x: (batch, seq_len, embed_dim)
        # padding_mask: (batch, seq_len)
        return x, padding_mask

    def _forward_aggregator(self, x, padding_mask):
        with torch.autocast("cuda", enabled=self.use_amp):
            if self.cls_blocks is not None:
                # for classification: extract using class token
                cls_tokens = self.cls_token.expand(
                    x.size(0), 1, -1
                )  # (batch, 1, embed_dim)
                for block in self.cls_blocks:
                    cls_tokens = block(
                        x, x_cls=cls_tokens, padding_mask=padding_mask
                    )  # (batch, 1, embed_dim)
                cls_tokens = cls_tokens.squeeze(1)  # (batch, embed_dim)
            else:
                # for classification: simple average pooling
                mask = ~padding_mask.unsqueeze(1)  # (batch, 1, seq_len)
                x = x.transpose(1, 2).contiguous()  # (batch, embed_dim, seq_len)
                counts = mask.float().sum(-1)  # (batch, 1)
                counts = torch.max(counts, torch.ones_like(counts))  # >=1
                cls_tokens = (x * mask).sum(-1) / counts  # (batch, embed_dim)

            x_cls = self.norm(cls_tokens)  # (batch, embed_dim)
        return x_cls

    def forward(self, x, frames, v=None, mask=None, uu=None, uu_idx=None):
        # x: (batch_size, num_fts, seq_len)
        # v: (batch_size, 4, seq_len) [px,py,pz,energy]
        # mask: (batch_size, 1, seq_len) -- real particle = 1, padded = 0
        # for pytorch: uu (batch_size, C', num_pairs), uu_idx (batch_size, 2, num_pairs)
        # for onnx: uu (batch_size, C', seq_len, seq_len), uu_idx=None
        self.attention.prepare_frames(frames)

        x, padding_mask = self._forward_encoder(x, v=v, mask=mask, uu=uu, uu_idx=uu_idx)

        if self.cls_blocks is None and self.fc is None:
            # x: (batch, seq_len, embed_dim)
            # padding_mask: (batch, seq_len)
            return x, padding_mask

        with torch.autocast("cuda", enabled=self.use_amp):
            # === for segmentation ===
            if self.for_segmentation:
                x = self.norm(x)
                if self.fc is not None:
                    x = self.fc(x)
                # x: (P, N, C) -> output: (N, C, P)
                output = x.transpose(1, 2).contiguous()
                if self.for_inference:
                    output = torch.softmax(output, dim=1)
                # print('output:\n', output)
                return output

            x_cls = self._forward_aggregator(x, padding_mask)
            if self.fc is None:
                return x_cls

            # fc
            output = self.fc(x_cls)
            if self.for_inference:
                output = torch.softmax(output, dim=1)
            # print('output:\n', output)
            return output


### weight initialization methods ###
def init_weights_vit_timm(module: nn.Module, name: str = "") -> None:
    """ViT weight initialization, original timm impl (for reproducibility)"""
    if isinstance(module, nn.Linear):
        nn.init.trunc_normal_(module.weight, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif hasattr(module, "init_weights"):
        module.init_weights()


def init_weights_vit_moco(module: nn.Module, name: str = "") -> None:
    """ViT weight initialization, matching moco-v3 impl minus fixed PatchEmbed"""
    if isinstance(module, nn.Linear):
        if "in_proj" in name:
            # treat the weights of Q, K, V separately
            val = math.sqrt(
                6.0 / float(module.weight.shape[0] // 3 + module.weight.shape[1])
            )
            nn.init.uniform_(module.weight, -val, val)
        else:
            nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif hasattr(module, "init_weights"):
        module.init_weights()


def named_apply(
    fn: Callable,
    module: nn.Module,
    name="",
    depth_first: bool = True,
    include_root: bool = False,
) -> nn.Module:
    if not depth_first and include_root:
        fn(module=module, name=name)
    for child_name, child_module in module.named_children():
        child_name = ".".join((name, child_name)) if name else child_name
        named_apply(
            fn=fn,
            module=child_module,
            name=child_name,
            depth_first=depth_first,
            include_root=True,
        )
    if depth_first and include_root:
        fn(module=module, name=name)
    return module
